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

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

请注意：如果您交互式使用此类内容，请小心重新定义所依赖的功能。如果出于某种原因而决定在程序包中使用此宏，则请注意，实施类型盗版的人将使您的类型检查无效。

如果有人试图将这种技术应用于重要的代码，我建议您重新考虑，或者认真考虑如何使其更安全。如果您有任何提高安全性的想法，我很乐意听到！每次更改从属方法时，也许可以采取一些技巧来强制重新编译功能。