Type-unstable application of `map`

Consider the following (minimal?) example:

julia> struct A{N}
           A(::NTuple{N,Int}) where {N} = new{N}()
       end

julia> t = (1,2)
(1, 2)

julia> T = typeof(t)
Tuple{Int64, Int64}

julia> using Test

julia> @inferred A(t) # A itself is type-stable and correctly inferred
A{2}()

julia> @inferred map(A, [t]) # ...but `map`ping A to a vector of T is not!
ERROR: return type Vector{A{2}} does not match inferred return type Union{Vector{Any}, Vector{A{2}}}

julia> map(A, T[]) # indeed, on an empty vector it returns Any[]
Any[]

I’m surprized by this behavior, because I think this is the first time I observe a type instability related to whether map is applied to an empty vector (in contrast to reduce, for example, where such instabilities frequently need to be fixed using a correctly typed initial value). Consider for example the following example, very similar to the above but which infers correctly:

julia> struct B
           B(::NTuple{N,Int}) where {N} = new()
       end

julia> @inferred map(B, [t])
1-element Vector{B}:
 B()

julia> map(B, T[])
B[]

Do you know why type instabilities appear with the parametric type A but not the simpler type B? And perhaps how to avoid them (short of type-asserting the output of map(A, ...)?


PS: if that matters, in my real use-case I can guarantee that map(A,...) is applied to a non-empty vector.

1 Like

I’d classify this as a bug. The function _collect that’s ultimately being called here tries to infer the default eltype by doing @default_eltype(iter). iter here is a Base.Generator, mapping A. The macro does the following check:

julia> @macroexpand Base.@default_eltype(iter)
quote
    #= array.jl:804 =#
    if iter isa Base.Generator && iter.f isa Base.Type
        #= array.jl:805 =#
        var"#3#T" = iter.f
    else
        #= array.jl:807 =#
        var"#3#T" = (Base.Core).Compiler.return_type(Base._iterator_upper_bound, Base.Tuple{Base.typeof(iter)})
    end
    #= array.jl:809 =#
    Base.promote_typejoin_union(var"#3#T")
end

So it checks whether iter isa Base.Generator (it is, because that’s how our map works), and then it checks whether iter.f (the function being mapped) is a type - which is then just used as the direct type for this, resulting in the later mapping only ever assuming an A, and not a A{2}.

If you use an outer constructor (thereby not passing a Type), it works as expected:

julia> a(t) = A(t)
a (generic function with 1 method)

julia> @inferred map(a, [t])
1-element Vector{A{2}}:
 A{2}()

because the iter.f isa Base.Type check is false.

This is a bit hard to find from @code_warntype or Cthulhu.jl alone, because of Color incompletely specified type parameters in `@code_warntype` differently · Issue #41251 · JuliaLang/julia · GitHub… I was only lucky to notice that the inferred types everywhere was Type{A}, instead of Type{A{2}}.

3 Likes

If it would hit the other path, it would infer just fine:

julia> (Base.Core).Compiler.return_type(Base._iterator_upper_bound, Base.Tuple{Base.Generator{Vector{Tuple{Int,Int}}, Type{A}}})
A{2}

So it’d probably be a quick & easy fix to modify @default_eltype to do

quote
    #= array.jl:804 =#
    if iter isa Base.Generator && iter.f isa Base.Type && !(iter.f isa Base.UnionAll)
        #= array.jl:805 =#
        var"#3#T" = iter.f
    else
        #= array.jl:807 =#
        var"#3#T" = (Base.Core).Compiler.return_type(Base._iterator_upper_bound, Base.Tuple{Base.typeof(iter)})
    end
    #= array.jl:809 =#
    Base.promote_typejoin_union(var"#3#T")
end

instead. Whether that is desirable though, I can’t say - there may be some other inference or compile time concerns at play here too.

2 Likes

Incredible! Thanks a lot!

I have to admit I stared a long time at the Cthulhu output, but did not manage to figure out where instabilities were introduced.

Although there is an easy workaround, I submitted an issue in case it would be more beneficial to the project to fix this once and for all

4 Likes