wip: report infinite iteration

WIP because I found this is better to be generalized to a control flow
analysis that checks for "never reaching exit points".
For As an example, `sum(a for a in NeverTeminate(::Int))` will come down
to `Base._fold_impl`, which directly uses `iterate(::NeverTeminate, [state])`,
against which the current implementation based on the iteration protocol
can't report an error.

> Adapated from https://github.com/JuliaLang/julia/blob/24d9eab45632bdb3120c9e664503745eb58aa2d6/base/reduce.jl#L53-L65

```julia
function _foldl_impl(op::OP, init, itr) where {OP}
    # Unroll the while loop once; if init is known, the call to op may
    # be evaluated at compile time
    y = iterate(itr)
    y === nothing && return init
    v = op(init, y[1])
    while true
        y = iterate(itr, y[2])
        y === nothing && break
        v = op(v, y[1])
    end
    return v
end
```

> In a lowered representation

```julia
CodeInfo(
1 ─       Core.NewvarNode(:(v))
│         y = Base.iterate(itr)
│   %3  = y === Base.nothing
└──       goto #3 if not %3
2 ─       return init
3 ─ %6  = Base.getindex(y, 1)
└──       v = (op)(init, %6)
4 ┄       goto #8 if not true
5 ─ %9  = Base.getindex(y, 2)
│         y = Base.iterate(itr, %9)
│   %11 = y === Base.nothing
└──       goto #7 if not %11
6 ─       goto #8
7 ─ %14 = v
│   %15 = Base.getindex(y, 1)
│         v = (op)(%14, %15)
└──       goto #4
8 ┄       return v
)
```

We can report infinite iteration by checking both `goto #3 if not %3` and
`goto #7 if not %11` never happens and thus this function never returns.

---

- closes #135
- closes #179

Author: aviatesk

src/JET.jl

@@ -47,6 +47,8 @@ import .CC:
     abstract_eval_special_value,
     abstract_eval_value,
     abstract_eval_statement,
+    typeinf_local,
+    abstract_iteration,
     # typeinfer.jl
     typeinf,
     _typeinf,
@@ -108,6 +110,7 @@ import .CC:
     BasicBlock,
     slot_id,
     widenconst,
+    widenconditional,
     ⊑,
     is_throw_call,
     tmerge,
@@ -402,6 +405,13 @@ macro jetconfigurable(funcdef)
 end
 const _JET_CONFIGURATIONS = Dict{Symbol,Symbol}()
 
+macro get!(x, key, default)
+    return quote
+        x, key = $(esc(x)), $(esc(key))
+        haskey(x, key) ? x[key] : (x[key] = $(esc(default)))
+    end
+end
+
 # utils
 # -----
 

src/abstractinterpretation.jl

@@ -428,6 +428,120 @@ function CC.abstract_eval_statement(interp::JETInterpreter, @nospecialize(e), vt
     return @invoke abstract_eval_statement(interp::AbstractInterpreter, e, vtypes::VarTable, sv::InferenceState)
 end
 
+# TODO generalize the following check to "this function will never return"
+
+# in this overload we will work on pre-optimization state
+function CC.typeinf_local(interp::JETInterpreter, frame::InferenceState)
+    @invoke typeinf_local(interp::AbstractInterpreter, frame::InferenceState)
+    maybe_report_infinite_iterations!(interp, frame)
+end
+
+function maybe_report_infinite_iterations!(interp::JETInterpreter, frame::InferenceState)
+    src = frame.src
+    infos = maybe_find_iter_infos(src)
+    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
+
+            @assert ssavaluetypes[dest] === NOT_FOUND # abstract interpretaion should never reach the destination
+            t = nothing
+            pc = nothing
+            for (call, _) in call2cond
+                ea = ((src.code[call]::Expr).args[2]::Expr).args # x_ = iterate(t′, [state])
+                t′ = @invoke abstract_eval_value(interp::AbstractInterpreter, ea[2]::Any, frame.stmt_types[call]::VarTable, frame::InferenceState)
+                @assert isnothing(t) || t === t′
+                t = t′
+
+                if isnothing(pc)
+                    pc = call
+                end
+            end
+            report!(interp, InfiniteIterationErrorReport(interp, frame, t, pc))
+        end
+    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)
+            cond, (call, dest) = maybeinfo
+            if isnothing(stmts)
+                stmts = Dict{Int,Vector{Tuple{Int,Int}}}()
+            end
+            push!(@get!(stmts, dest, Tuple{Int,Int}[]), (call, cond))
+        end
+    end
+
+    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))
+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
+
+    @inbounds begin
+    region = src.code[stmts][end-3:end]
+
+    terminator = region[end]
+    isa(terminator, GotoIfNot) || return nothing
+
+    preds = region[end-3:end-1]
+    is_iterate_stmt(preds[1]) || return nothing
+    is_nothing_check_stmt(preds[2]) || return nothing
+    is_notint_stmt(preds[3]) || return nothing
+
+    return stmts[end-1], (stmts[end-3], 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
+
+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
+end
+
+function is_notint_stmt(@nospecialize(x))
+    @isexpr(x, :call) || return false
+    return is_global_ref(x.args[1], Base, :not_int)
+end
+
+is_global_ref(@nospecialize(x), mod::Module, name::Symbol) = isa(x, GlobalRef) && x.mod === mod && x.name === name
+
+# TODO more detailed error report using the context of abstract iteration
+function CC.abstract_iteration(interp::JETInterpreter, @nospecialize(itft), @nospecialize(itertype), sv::InferenceState)
+    ret = @invoke abstract_iteration(interp::AbstractInterpreter, @nospecialize(itft), @nospecialize(itertype), sv::InferenceState)
+    rts, info = ret
+    if length(rts) == 1 && first(rts) === Bottom && isnothing(info)
+        report!(interp, InfiniteIterationErrorReport(interp, sv, itertype))
+    end
+    return ret
+end
+
 function CC.finish(me::InferenceState, interp::JETInterpreter)
     @invoke finish(me::InferenceState, interp::AbstractInterpreter)
 

src/reports.jl

@@ -508,6 +508,17 @@ let s = sprint(showerror, DivideError())
     global get_msg(::Type{DivideErrorReport}, interp, sv::InferenceState) = s
 end
 
+@reportdef struct InfiniteIterationErrorReport <: InferenceErrorReport
+    @nospecialize(typ)
+end
+# if provided, use program counter where infinite iteration(s) is found
+function InfiniteIterationErrorReport(interp, sv::InferenceState, @nospecialize(typ), pc = nothing)
+    vst = VirtualFrame[get_virtual_frame(interp, sv, isnothing(pc) ? get_currpc(sv) : pc)]
+    msg = "iterate(::$typ) won't terminate"
+    sig = get_sig(interp, sv, isnothing(pc) ? get_stmt(sv) : get_stmt(sv, pc))
+    return InfiniteIterationErrorReport(vst, msg, sig, typ)
+end
+
 # TODO we may want to hoist `InvalidConstXXX` errors into top-level errors
 
 @reportdef struct InvalidConstantRedefinition <: InferenceErrorReport

test/test_abstractinterpretation.jl

@@ -908,3 +908,82 @@ end
         @test isa(r, GeneratorErrorReport) && r.err == "invalid argument"
     end
 end
+
+@testset "infinite iteration" begin
+    m = @fixturedef begin
+        struct NeverTerminate
+            val::Int
+        end
+        Base.iterate(nv::NeverTerminate, state = 0) =
+            state>nv.val ? (state, state+1) : (state, state+1)
+    end
+
+    let # iteration protocol
+        interp, frame = @eval m $analyze_call((Int,)) do n
+            for a in NeverTerminate(n)
+                println(a)
+            end
+        end
+        @test any(interp.reports) do r
+            isa(r, InfiniteIterationErrorReport) &&
+            r.typ === m.NeverTerminate &&
+            any(r.vst) do vf
+                vf.line == (@__LINE__)-8
+            end
+        end
+    end
+
+    let # iteration protocol, nested
+        interp, frame = @eval m $analyze_call((Int,)) do n
+            for a in 1:n
+                for b in NeverTerminate(a)
+                    println(b)
+                end
+            end
+        end
+        @test any(interp.reports) do r
+            isa(r, InfiniteIterationErrorReport) &&
+            r.typ === m.NeverTerminate &&
+            any(r.vst) do vf
+                vf.line == (@__LINE__)-9
+            end
+        end
+    end
+
+    let # iteration protocol, on container type
+        interp, frame = @eval m $analyze_call((Int,)) do n
+            sum((a for a in NeverTerminate(n))...)
+        end
+        @test any(interp.reports) do r
+            isa(r, InfiniteIterationErrorReport)
+        end
+    end
+
+    let # complicated control flow, no false positive
+        interp, frame = analyze_call((Char,Tuple{Char,Char})) do x, itr
+            # adapated from https://github.com/JuliaLang/julia/blob/24d9eab45632bdb3120c9e664503745eb58aa2d6/base/operators.jl#L1278-L1297
+            anymissing = false
+            for y in itr
+                v = (y == x)
+                if ismissing(v)
+                    anymissing = true
+                elseif v
+                    return true
+                end
+            end
+            return anymissing ? missing : false
+        end
+        @test isempty(interp.reports)
+    end
+
+    let # general case of "this function never return"
+        # NOTE comes down to `Base._foldl_impl`
+        interp, frame = @eval m $analyze_call((Int,)) do n
+            sum(a for a in NeverTerminate(n))
+        end
+        @test_broken any(interp.reports) do r
+            isa(r, InfiniteIterationErrorReport) &&
+            r.typ === m.NeverTerminate
+        end
+    end
+end