Are "inline assignments" inside loop predicates bad for type inference?

Internals & Design
, , , ,
2

A very small modification makes the error go away for me (on a somewhat recent master - which version did you test?):

function inline_assign_while_loop(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

If we take a look at the lowered representation, the difference is extremely subtle:

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

The reason for the warning by JET seems to be that the SSA slot %4 is checked against nothing, not i. So the issue does not lie in inference per se, but in propagating the information we have about %4 to i as well. This seems very interesting, so I’ll open an issue.

EDIT: Issue opened here.

4 Likes