Type stability in tree constructors

[kirtsar]

Given some data X, i want to create tree (KDTree, or Ball)

X = [1. 2; 3 4]
@code_warntype KDTree(X)

Body::KDTree{_1,Euclidean,_2} where _2 where _1
75 1 ─ %1 = %new(Distances.Euclidean, 0.0)::Euclidean           │╻╷ Type
   │   %2 = $(Expr(:foreigncall, :(:jl_alloc_array_2d), Array{Float64,2}, svec(Any, Int64, Int64), :(:ccall), 3, Array{Float64,2}, 0, 0, 0, 0))::Array{Float64,2}
   │   %3 = invoke NearestNeighbors.:(#KDTree#17)(10::Int64, true::Bool, true::Bool, %2::Array{Float64,2}, _1::Type, _2::Array{Float64,2}, %1::Euclidean)::KDTree{_1,Euclidean,_2} where _2 where _1
   └──      return %3  

@code_warntype BallTree(X)

Body::BallTree{_1,_2,_3,Euclidean} where _3 where _2 where _1
82 1 ─ %1 = %new(Distances.Euclidean, 0.0)::Euclidean           │╻╷ Type
   │   %2 = $(Expr(:foreigncall, :(:jl_alloc_array_2d), Array{Float64,2}, svec(Any, Int64, Int64), :(:ccall), 3, Array{Float64,2}, 0, 0, 0, 0))::Array{Float64,2}
   │   %3 = invoke NearestNeighbors.:(#BallTree#19)(10::Int64, true::Bool, true::Bool, %2::Array{Float64,2}, _1::Type, _2::Array{Float64,2}, %1::Euclidean)::BallTree{_1,_2,_3,Euclidean} where _3 where _2 where _1
   └──      return %3  

[KristofferC]

We can't really fix this completely when the input is a matrix because we specialize on the dimension if the points. If the input was a vector of static vectors then it should be theoretically possible to have it be type stable. In practice it is unlikely to be too much of an issue since in most cases there is either a function barrier or the call has enough points that the dynamic dispatch is negligible.

[KristofferC]

Note

julia> using StaticArrays

julia> x = [rand(SVector{3, Float64}) for i in 1:10];

julia> @code_warntype KDTree(x)
MethodInstance for KDTree(::Vector{SVector{3, Float64}})
  from KDTree(data::AbstractVector{V}; ...) where V<:AbstractArray @ NearestNeighbors ~/JuliaPkgs/NearestNeighbors.jl/src/kd_tree.jl:19
Static Parameters
  V = SVector{3, Float64}
Arguments
  #self#::Type{KDTree}
  data::Vector{SVector{3, Float64}}
Body::KDTree{SVector{3, Float64}, Euclidean, Float64, SVector{3, Float64}}