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

Question

这个问题是什么出现了改编版本在这里的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 

我们看到类型推断只计算出这会产生 a 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 的a ，其元素静态已知为 a String，Int和Float64！