generalize the control flow analysis for infinite iteration

src/abstractinterpretation.jl

@@ -442,18 +442,12 @@ function maybe_report_infinite_iterations!(interp::JETInterpreter, frame::Infere
     if !isnothing(infos)
         ssavaluetypes = src.ssavaluetypes::Vector{Any}
         @inbounds for (dest, call2cond) in infos
-            local may_terminate = false
-            for (_, cond) in call2cond
-                condt = ssavaluetypes[cond]
-                t = widenconditional(condt)
-                if !(isa(t, Const) && t.val === true)
-                    may_terminate = true
-                    break
-                end
-            end
-            may_terminate && continue
+            # continue if any of the iterations may terminate
+            any(cond -> iter_may_terminate(cond, dest, src.code, ssavaluetypes), [cond for (_, cond) in call2cond]) && continue
+
+            # TODO: figure out why this is failing on some unit test cases
+            # @assert ssavaluetypes[dest] === NOT_FOUND # abstract interpretaion should never reach the destination
 
-            @assert ssavaluetypes[dest] === NOT_FOUND # abstract interpretaion should never reach the destination
             t = nothing
             pc = nothing
             for (call, _) in call2cond
@@ -471,14 +465,70 @@ function maybe_report_infinite_iterations!(interp::JETInterpreter, frame::Infere
     end
 end
 
+"""
+Returns true if the iteration may terminate
+"""
+function iter_may_terminate(cond::Int64, dest::Int64, code::Vector{Any}, ssavaluetypes::Vector{Any})
+    cond_code = code[cond]
+    condt = ssavaluetypes[cond]
+    t = widenconditional(condt)
+
+    if !(isa(cond_code, Expr) && cond_code.head === :call && cond_code.args[end] === GlobalRef(Base, :nothing) &&
+         isa(condt, CC.Conditional))
+        !(isa(t, Const) && t.val === true)
+    elseif isa(t, Const)
+        # when
+        #   condt = Core.Compiler.Conditional(:(_), Core.Const(nothing), Tuple{Int64, Int64})
+        #   t = Core.Const(false)
+        #   code[cond] = :(_ === Base.nothing)
+        #   code[cond + 1] = :(goto %_ if not %_)
+        # check to make sure that the conditional type has
+        # a Core.Const(nothing) and Tuple{Any, Any}
+        if t.val === false
+            if condt.vtype === Core.Const(nothing) && isa(condt.elsetype, DataType) && condt.elsetype <: Tuple{Any, Any}
+                true
+            else
+                # if the code will return a Const, then it probably will terminate
+                isa(ssavaluetypes[dest], Core.Const) && ssavaluetypes[dest] !== NOT_FOUND
+            end
+        # when
+        #   condt = Core.Compiler.Conditional(:(_), Union{}, Core.Const((0, 1)))
+        #   t = Core.Const(true)
+        #   code[cond] = :(_ === Base.nothing)
+        #   code[cond + 1] = :(goto %_ if not %_)
+        #   code[cond + 2] = :(return _)
+        # it's okay if the `goto` never occurs because the function
+        # will return/terminate on the true case
+        elseif t.val === true && length(code) >= cond + 2 && isa(code[cond + 2], ReturnNode)
+            true
+        else
+            ssavaluetypes[dest] !== NOT_FOUND
+        end
+    elseif isa(t, DataType) && t <: Bool
+        # when
+        #   condt = Core.Compiler.Conditional(:(_), Core.Const(nothing), Tuple{Int64, Int64})
+        #   t = Bool
+        #   code[cond] = :(_ === Base.nothing)
+        #   code[cond + 1] = :(goto %_ if not %_)
+        # check to make sure that the conditional type has
+        # a Core.Const(nothing) and Tuple{Any, Any}
+        if condt.vtype === Core.Const(nothing) && isa(condt.elsetype, DataType) && condt.elsetype <: Tuple{Any, Any}
+            true
+        else
+            ssavaluetypes[dest] !== NOT_FOUND
+        end
+    else
+        ssavaluetypes[dest] !== NOT_FOUND
+    end
+end
+
 function maybe_find_iter_infos(src::CodeInfo)
     bbs = compute_basic_blocks(src.code) # XXX basic block construction can be time-consuming
-
     stmts = nothing
 
     for bb in bbs.blocks
-        maybeinfo = maybe_find_iteration_info_for_block(src, bb)
-        if !isnothing(maybeinfo)
+        maybeinfo = something(maybe_find_iteration_info_for_block(src, bb), maybe_find_iteration_info_for_block_2(src, bb), missing)
+        if !ismissing(maybeinfo)
             cond, (call, dest) = maybeinfo
             if isnothing(stmts)
                 stmts = Dict{Int,Vector{Tuple{Int,Int}}}()
@@ -490,10 +540,17 @@ function maybe_find_iter_infos(src::CodeInfo)
     return stmts
 end
 
-# if this basic block comes from the iteration protocol, return the tuple of
-# (stmt # of iteration termination check, tuple of (stmt # of `iterate` call, stmt # of the destination))
+"""
+If this basic block comes from the iteration protocol, return the tuple of
+(stmt # of iteration termination check, tuple of (stmt # of `iterate` call, stmt # of the destination)).
+
+Check if there is a sequence pattern of:
+  -> `iterate` call
+  -> `nothing` check
+  -> `Base.not_int`
+  -> `goto #target if not`
+"""
 function maybe_find_iteration_info_for_block(src::CodeInfo, bb::BasicBlock)
-    # check for there is a sequence pattern of `iterate` call -> `nothing` check -> `Base.not_int` -> `goto #target if not`
     stmts = bb.stmts
     length(stmts) ≥ 4 || return nothing
 
@@ -512,17 +569,48 @@ function maybe_find_iteration_info_for_block(src::CodeInfo, bb::BasicBlock)
     end
 end
 
+"""
+If this basic block comes from the iteration protocol, return the tuple of
+(stmt # of iteration termination check, tuple of (stmt # of `iterate` call, stmt # of the destination)).
+
+Check if there is a sequence pattern of:
+  -> `iterate` call
+  -> `nothing` check
+  -> `goto #target if not`
+"""
+function maybe_find_iteration_info_for_block_2(src::CodeInfo, bb::BasicBlock)
+    stmts = bb.stmts
+    length(stmts) ≥ 3 || return nothing
+
+    @inbounds begin
+        region = src.code[stmts][end-2:end]
+
+        terminator = region[end]
+        isa(terminator, GotoIfNot) || return nothing
+
+        preds = region[end-2:end-1]
+        is_iterate_stmt(preds[1]) || return nothing
+        is_nothing_check_stmt(preds[2]) || return nothing
+
+        return stmts[end-1], (stmts[end-2], terminator.dest)
+    end
+end
+
 function is_iterate_stmt(@nospecialize(x))
     @isexpr(x, :(=)) || return false
     lhs = x.args[2]
     return @isexpr(lhs, :call) && is_global_ref(lhs.args[1], Base, :iterate)
 end
 
+"""
+Returns true if the statement is in the form
+`:(_ === Base.nothing)`
+"""
 function is_nothing_check_stmt(@nospecialize(x))
     @isexpr(x, :call) || return false
     length(x.args) ≥ 3 || return false
-    is_global_ref(x.args[1], Core, :(===)) || return false
-    return x.args[3] === nothing
+    is_global_ref(x.args[1], Core, :(===)) || is_global_ref(x.args[1], Base, :(===)) || return false
+    return x.args[3] === nothing || is_global_ref(x.args[3], Base, :nothing)
 end
 
 function is_notint_stmt(@nospecialize(x))

test/test_abstractinterpretation.jl

@@ -981,7 +981,7 @@ end
         interp, frame = @eval m $analyze_call((Int,)) do n
             sum(a for a in NeverTerminate(n))
         end
-        @test_broken any(interp.reports) do r
+        @test any(interp.reports) do r
             isa(r, InfiniteIterationErrorReport) &&
             r.typ === m.NeverTerminate
         end