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

Question

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

Answer 1

一种方法是使用 generated function .例如，假设有问题的函数是

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

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

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