如何编写一个函数来检查每个调用方法的返回类型是否静态可推断?

Question

我想编写一个函数，如果 julia 无法推断出该函数的具体返回类型，它将抛出错误。如何在没有任何运行时开销的情况下执行此操作？

Answer 1

一种方法是使用生成的函数。例如，假设所讨论的函数是

f(x) = x + (rand(Bool) ? 1.0 : 1)

我们可以改为写

_f(x) = x + (rand(Bool) ? 1.0 : 1)
@generated function f(x)
out_type = Core.Compiler.return_type(_f, Tuple{x})
if !isconcretetype(out_type)
error("$f($x) does not infer to a concrete type")
end
:(_f(x))
end

现在我们可以在 repl 上对此进行测试。 浮点输入很好，但整数错误：

julia> f(1.0)
2.0
julia> f(1)
ERROR: f(Int64) does not infer to a concrete type
Stacktrace:
[1] error(::String) at ./error.jl:33
[2] #s28#4(::Any, ::Any) at ./REPL[5]:4
[3] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any,N} where N) at ./boot.jl:524
[4] top-level scope at REPL[8]:1

由于我们使用生成的函数的方式，类型检查和错误抛出仅在编译时发生，因此我们无需为此支付运行时成本。

如果以上看起来样板代码太多了，我们可以写一个宏自动生成内部函数，生成任意函数签名的函数：

using MacroTools: splitdef, combinedef
strip_type_asserts(ex::Expr) = ex.head == :(::) ? ex.args[1] : ex
strip_type_asserts(s) = s
macro checked(fdef)
d = splitdef(fdef)
f = d[:name]
args = d[:args]
whereparams = d[:whereparams]
d[:name] = gensym()
shadow_fdef = combinedef(d)
args_stripped = strip_type_asserts.(args)
quote
$shadow_fdef
@generated function $f($(args...)) where {$(whereparams...)}
d = $d
T = Tuple{$(args_stripped...)}
shadowf = $(d[:name])
out_type = Core.Compiler.return_type(shadowf, T)
sig = collect(T.parameters)
if !isconcretetype(out_type)
f = $f
sig = reduce(*, (", $U" for U in T.parameters[2:end]), init="$(T.parameters[1])")
error("$f($(sig...)) does not infer to a concrete type")
end
args = $args
#Core.println("statically inferred return type was $out_type")
:($(shadowf)($(args...)))
end
end |> esc
end

现在在 repl 中，我们只需要用@checked注释函数定义：

julia> @checked g(x, y) = x + (rand(Bool) ? 1.0 : 1)*y
f (generic function with 2 methods)
julia> g(1, 2.0)
3.0
julia> g(1, 2)
ERROR: g(Int64, Int64) does not infer to a concrete type
Stacktrace:
[1] error(::String) at ./error.jl:33
[2] #s28#5(::Any, ::Any, ::Any) at ./REPL[11]:22
[3] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any,N} where N) at ./boot.jl:524
[4] top-level scope at REPL[14]:1

---

编辑：评论中指出，我违反了在这里使用生成函数的"规则"之一，因为如果有人重新定义@checked函数所依赖的函数，则编译时生成的函数中发生的情况可能会被默默地失效。例如：

julia> g(x) = x + 1;
julia> @checked f(x) = g(x) + 1;
julia> f(1) # shouldn't error
3
julia> g(x) = rand(Bool) ? 1.0 : 1
g (generic function with 1 method)
julia> f(1) # Should error but doesn't!!!
2.0
julia> f(1)
2

所以请注意：如果你以交互方式使用这样的东西，要小心重新定义你所依赖的函数。如果出于某种原因，您决定在包中使用此宏，请注意，提交类型盗版的人将使你的类型检查无效。

如果有人试图将这种技术应用于重要的代码，我建议要么重新考虑，要么认真考虑如何使其更安全。如果您有任何使其更安全的想法，我很想听听！也许您可以做一些技巧来强制在每次更改依赖方法时重新编译函数。