julia - 切片元组时避免类型不稳定

Question

这个问题是出现的内容的改编版本here在 JuliaLang Zulip 服务台。

假设我有一个函数接收异类类型 Tuple并且将返回该元组的切片，其中切片索引可以仅从类型信息静态推断。如何以正确推断输出类型的方式编写我的函数？

例如，假设我的功能是

function f(t::Tuple, A::Array{T, N}) where {T, N}
    if T <: AbstractFloat
        imin = 1
    elseif T <: Integer
        imin = 2
    else
        imin = 3
    end
    imax = N+2    
    t[imin:imax]
end 

我们看到，类型推断仅表明这会产生 Tuple ，即使在编译时所有需要的信息都可用，它不是长度或元素类型吗？

julia> let t = (:a, "b", 2, 3.0, Val(1), 2+im), A = rand(Int, 3,3)
           Base.return_types(f, Tuple{typeof(t), typeof(A)})
       end
1-element Array{Any,1}:
 Tuple

怎么写 f这样有效吗？

Answer 1

我最喜欢的策略(但也许有更简单的方法？)是写一个 @generated手动确保 julia 在编译时执行我想要的类型级别操作的函数:

@generated function f2(t::Tuple, A::Array{T, N}) where {T, N}
    if T <: AbstractFloat
        imin = 1
    elseif T <: Integer
        imin = 2
    else
        imin = 3
    end
    imax = N+2
    out_expr = Expr(:tuple, (:(t[$i]) for i ∈ imin:imax)...)
end 

这里的想法是在生成的函数体中，在编译时，我们确定什么 imin和 imax是，然后我们手动为我们的函数体构建一个表达式，内容为 (t[imin], t[imin+1], ..., t[imax-1], t[imax]) .

无论出于何种原因，Julia 都能够更好地推理 getindex(::Tuple, ::Int) 的序列。不是关于切片元组，即使是静态已知切片，因此通过手动构建此表达式，编译器能够执行我们想要的操作:

julia> let t = (:a, "b", 2, 3.0, Val(1), 2+im), A = rand(Int, 3,3)
           Base.return_types(f2, Tuple{typeof(t), typeof(A)})
       end
1-element Array{Any,1}:
 Tuple{String,Int64,Float64}

瞧，推断的输出类型是 Tuple长度为 3 的元素静态已知为 String , 和 Int和 Float64 !