Propagate type information to assigned variables from slots in lowered form

[Seelengrab]

This was found on discourse:

-function inline_assign_while_loop(permutation::Vector{Int})
+function inline_assign_while_loop_extra(permutation::Vector{Int})
    indices = Int[]
    v = 1
-    while (i=findfirst(==(v),permutation)) !== nothing
+    while ((i=findfirst(==(v),permutation)); i !== nothing)
        push!(indices, i)
        v += 1
    end
    indices
end

Copyable functions

function inline_assign_while_loop(permutation::Vector{Int})
    indices = Int[]
    v = 1
    while (i=findfirst(==(v),permutation)) !== nothing
        push!(indices, i)
        v += 1
    end
    indices
end

function inline_assign_while_loop_extra(permutation::Vector{Int})
    indices = Int[]
    v = 1
    while ((i=findfirst(==(v),permutation)); i !== nothing)
        push!(indices, i)
        v += 1
    end
    indices
end

The version checking the result of the assignment directly has JET complaining about a possible push!(::Vector{Int}, ::Nothing), while the version with i !== nothing correctly notices that i can't be nothing in the loop:

julia> using JET, Random

julia> test_perm = randperm(10);

julia> @report_call inline_assign_while_loop(test_perm)
═════ 1 possible error found ═════
┌ @ REPL[20]:5 Main.push!(indices, i)
│┌ @ array.jl:1050 itemT = Base.convert(_, item)
││ no matching method found for call signature (Tuple{typeof(convert), Type{Int64}, Nothing}): itemT = Base.convert(_::Type{Int64}, item::Nothing)
│└─────────────────


julia> @report_call inline_assign_while_loop_extra(test_perm)
No errors detected

The difference seems to be explainable due to a difference in type information propagation, if I interpret @code_lowered correctly:

julia> @code_lowered inline_assign_while_loop(test_perm)
CodeInfo(
1 ─      indices = Base.getindex(Main.Int)
└──      v = 1
2 ┄ %3 = (==)(v)
│   %4 = Main.findfirst(%3, permutation)
│        i = %4
│   %6 = %4 !== Main.nothing
└──      goto #4 if not %6
3 ─      Main.push!(indices, i)
│        v = v + 1
└──      goto #2
4 ─      return indices
)

julia> @code_lowered inline_assign_while_loop_extra(test_perm)
CodeInfo(
1 ─      indices = Base.getindex(Main.Int)
└──      v = 1
2 ┄ %3 = (==)(v)
│        i = Main.findfirst(%3, permutation)
│   %5 = i !== Main.nothing
└──      goto #4 if not %5
3 ─      Main.push!(indices, i)
│        v = v + 1
└──      goto #2
4 ─      return indices
)

In the output above, it seems like the information about %4 is not propagated to i, even though through the assignment they're identical. I think possible solutions are propagating information we learn about %4 to i as well, or changing lowering to use i instead of %4 here (though that would only fix this specific instance of the problem).