Stateful: fix major bug in implementation

base/compiler/abstractinterpretation.jl

@@ -126,6 +126,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
     length(argtypes) >= nargs || return Any
     haveconst = false
     for a in argtypes
+        a = maybe_widen_conditional(a)
         if isa(a, Const) && !isdefined(typeof(a.val), :instance) && !(isa(a.val, Type) && issingletontype(a.val))
             # have new information from argtypes that wasn't available from the signature
             haveconst = true
@@ -154,6 +155,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
         if !istopfunction(f, :getproperty) && !istopfunction(f, :setproperty!)
             # in this case, see if all of the arguments are constants
             for a in argtypes
+                a = maybe_widen_conditional(a)
                 if !isa(a, Const) && !isconstType(a)
                     return Any
                 end
@@ -167,7 +169,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
         if method.isva
             vargs = argtypes[(nargs + 1):end]
             for i in 1:length(vargs)
-                a = vargs[i]
+                a = maybe_widen_conditional(vargs[i])
                 if i > length(inf_result.vargs)
                     push!(inf_result.vargs, a)
                 elseif a isa Const
@@ -176,7 +178,7 @@ function abstract_call_method_with_const_args(@nospecialize(f), argtypes::Vector
             end
         end
         for i in 1:nargs
-            a = argtypes[i]
+            a = maybe_widen_conditional(argtypes[i])
             if a isa Const
                 atypes[i] = a # inject Const argtypes into inference
             end
@@ -486,8 +488,8 @@ end
 
 function pure_eval_call(@nospecialize(f), argtypes::Vector{Any}, @nospecialize(atype), sv::InferenceState)
     for i = 2:length(argtypes)
-        a = argtypes[i]
-        if !(isa(a,Const) || isconstType(a))
+        a = maybe_widen_conditional(argtypes[i])
+        if !(isa(a, Const) || isconstType(a))
             return false
         end
     end
@@ -505,7 +507,7 @@ function pure_eval_call(@nospecialize(f), argtypes::Vector{Any}, @nospecialize(a
         return false
     end
 
-    args = Any[ (a=argtypes[i]; isa(a,Const) ? a.val : a.parameters[1]) for i in 2:length(argtypes) ]
+    args = Any[ (a = maybe_widen_conditional(argtypes[i]); isa(a, Const) ? a.val : a.parameters[1]) for i in 2:length(argtypes) ]
     try
         value = Core._apply_pure(f, args)
         # TODO: add some sort of edge(s)
@@ -530,13 +532,16 @@ function abstract_call(@nospecialize(f), fargs::Union{Tuple{},Vector{Any}}, argt
 
     if isa(f, Builtin) || isa(f, IntrinsicFunction)
         if f === ifelse && fargs isa Vector{Any} && length(argtypes) == 4 && argtypes[2] isa Conditional
-            cnd = argtypes[2]
+            # try to simulate this as a real conditional (`cnd ? x : y`), so that the penalty for using `ifelse` instead isn't too high
+            cnd = argtypes[2]::Conditional
             tx = argtypes[3]
             ty = argtypes[4]
-            if isa(fargs[3], Slot) && slot_id(cnd.var) == slot_id(fargs[3])
+            a = ssa_def_expr(fargs[3], sv)
+            b = ssa_def_expr(fargs[4], sv)
+            if isa(a, Slot) && slot_id(cnd.var) == slot_id(a)
                 tx = typeintersect(tx, cnd.vtype)
             end
-            if isa(fargs[4], Slot) && slot_id(cnd.var) == slot_id(fargs[4])
+            if isa(b, Slot) && slot_id(cnd.var) == slot_id(b)
                 ty = typeintersect(ty, cnd.elsetype)
             end
             return tmerge(tx, ty)
@@ -554,18 +559,20 @@ function abstract_call(@nospecialize(f), fargs::Union{Tuple{},Vector{Any}}, argt
                 a = ssa_def_expr(fargs[2], sv)
                 if isa(a, Slot)
                     aty = widenconst(argtypes[2])
+                    if rt === Const(false)
+                        return Conditional(a, Union{}, aty)
+                    elseif rt === Const(true)
+                        return Conditional(a, aty, Union{})
+                    end
                     tty_ub, isexact_tty = instanceof_tfunc(argtypes[3])
                     if isexact_tty && !isa(tty_ub, TypeVar)
                         tty_lb = tty_ub # TODO: this would be wrong if !isexact_tty, but instanceof_tfunc doesn't preserve this info
                         if !has_free_typevars(tty_lb) && !has_free_typevars(tty_ub)
                             ifty = typeintersect(aty, tty_ub)
-                            elsety = typesubtract(aty, tty_lb)
-                            if ifty != elsety
-                                return Conditional(a, ifty, elsety)
-                            end
+                            elty = typesubtract(aty, tty_lb)
+                            return Conditional(a, ifty, elty)
                         end
                     end
-                    return Bool
                 end
             elseif f === (===)
                 a = ssa_def_expr(fargs[2], sv)
@@ -574,21 +581,42 @@ function abstract_call(@nospecialize(f), fargs::Union{Tuple{},Vector{Any}}, argt
                 bty = argtypes[3]
                 # if doing a comparison to a singleton, consider returning a `Conditional` instead
                 if isa(aty, Const) && isa(b, Slot)
-                    if isdefined(typeof(aty.val), :instance) # can only widen a if it is a singleton
-                        return Conditional(b, aty, typesubtract(widenconst(bty), typeof(aty.val)))
+                    if rt === Const(false)
+                        aty = Union{}
+                    elseif rt === Const(true)
+                        bty = Union{}
+                    elseif bty isa Type && isdefined(typeof(aty.val), :instance) # can only widen a if it is a singleton
+                        bty = typesubtract(bty, typeof(aty.val))
                     end
-                    return isa(rt, Const) ? rt : Conditional(b, aty, bty)
+                    return Conditional(b, aty, bty)
                 end
                 if isa(bty, Const) && isa(a, Slot)
-                    if isdefined(typeof(bty.val), :instance) # same for b
-                        return Conditional(a, bty, typesubtract(widenconst(aty), typeof(bty.val)))
+                    if rt === Const(false)
+                        bty = Union{}
+                    elseif rt === Const(true)
+                        aty = Union{}
+                    elseif aty isa Type && isdefined(typeof(bty.val), :instance) # same for b
+                        aty = typesubtract(aty, typeof(bty.val))
                     end
-                    return isa(rt, Const) ? rt : Conditional(a, bty, aty)
+                    return Conditional(a, bty, aty)
+                end
+                if isa(b, Slot)
+                    return Conditional(b, bty, bty)
+                end
+                if isa(a, Slot)
+                    return Conditional(a, aty, aty)
                 end
             elseif f === Core.Compiler.not_int
                 aty = argtypes[2]
                 if isa(aty, Conditional)
-                    return Conditional(aty.var, aty.elsetype, aty.vtype)
+                    ifty = aty.elsetype
+                    elty = aty.vtype
+                    if rt === Const(false)
+                        ifty = Union{}
+                    elseif rt === Const(true)
+                        elty = Union{}
+                    end
+                    return Conditional(aty.var, ifty, elty)
                 end
             end
         end
@@ -947,7 +975,7 @@ end
 # determine whether `ex` abstractly evals to constant `c`
 function abstract_evals_to_constant(@nospecialize(ex), @nospecialize(c), vtypes::VarTable, sv::InferenceState)
     av = abstract_eval(ex, vtypes, sv)
-    return isa(av,Const) && av.val === c
+    return isa(av, Const) && av.val === c
 end
 
 # make as much progress on `frame` as possible (without handling cycles)
@@ -1018,7 +1046,7 @@ function typeinf_local(frame::InferenceState)
                 end
             elseif hd === :return
                 pc´ = n + 1
-                rt = abstract_eval(stmt.args[1], s[pc], frame)
+                rt = maybe_widen_conditional(abstract_eval(stmt.args[1], s[pc], frame))
                 if !isa(rt, Const) && !isa(rt, Type)
                     # only propagate information we know we can store
                     # and is valid inter-procedurally

base/compiler/inferenceresult.jl

@@ -123,7 +123,7 @@ function cache_lookup(code::MethodInstance, argtypes::Vector{Any}, cache::Vector
         if cache_code.linfo === code && length(argtypes) === (length(cache_vargs) + nargs)
             cache_match = true
             for i in 1:length(argtypes)
-                a = argtypes[i]
+                a = maybe_widen_conditional(argtypes[i])
                 ca = i <= nargs ? cache_args[i] : cache_vargs[i - nargs]
                 # verify that all Const argument types match between the call and cache
                 if (isa(a, Const) || isa(ca, Const)) && !(a === ca)

base/compiler/ssair/inlining.jl

@@ -745,7 +745,6 @@ function assemble_inline_todo!(ir::IRCode, linetable::Vector{LineInfoNode}, sv::
 
         ft = argextype(arg1, ir, sv.sp)
         has_free_typevars(ft) && continue
-        isa(ft, Conditional) && (ft = Bool)
         f = singleton_type(ft)
         # TODO: llvmcall can contain other calls as arguments, so this code doesn't work on it
         f === Core.Intrinsics.llvmcall && continue
@@ -786,7 +785,6 @@ function assemble_inline_todo!(ir::IRCode, linetable::Vector{LineInfoNode}, sv::
             new_atypes = Any[]
             ft = argextype(stmt.args[2], ir, sv.sp)
             has_free_typevars(ft) && continue
-            isa(ft, Conditional) && (ft = Bool)
             f = singleton_type(ft)
             # Push function type
             push!(new_atypes, ft)

base/compiler/tfuncs.jl

@@ -72,6 +72,9 @@ function instanceof_tfunc(@nospecialize(t))
     elseif isa(t, Union)
         ta, isexact_a = instanceof_tfunc(t.a)
         tb, isexact_b = instanceof_tfunc(t.b)
+        ta === Union{} && return tb, isexact_b
+        tb === Union{} && return ta, isexact_a
+        ta == tb && return ta, isexact_a && isexact_b
         return Union{ta, tb}, false # at runtime, will be exactly one of these
     end
     return Any, false
@@ -206,27 +209,29 @@ add_tfunc(ifelse, 3, 3,
                 return Bottom
             end
         elseif isa(cnd, Conditional)
-            # handled in abstract_call
+            # optimized (if applicable) in abstract_call
         elseif !(Bool ⊑ cnd)
             return Bottom
         end
         return tmerge(x, y)
     end, 1)
 function egal_tfunc(@nospecialize(x), @nospecialize(y))
-    if isa(x, Const) && isa(y, Const)
-        return Const(x.val === y.val)
-    elseif typeintersect(widenconst(x), widenconst(y)) === Bottom
+    xx = maybe_widen_conditional(x)
+    yy = maybe_widen_conditional(y)
+    if isa(x, Conditional) && isa(yy, Const)
+        yy.val === false && return Conditional(x.var, x.elsetype, x.vtype)
+        yy.val === true && return x
+        return x
+    elseif isa(y, Conditional) && isa(xx, Const)
+        xx.val === false && return Conditional(y.var, y.elsetype, y.vtype)
+        xx.val === true && return y
+    elseif isa(xx, Const) && isa(yy, Const)
+        return Const(xx.val === yy.val)
+    elseif typeintersect(widenconst(xx), widenconst(yy)) === Bottom
         return Const(false)
-    elseif (isa(x, Const) && y === typeof(x.val) && isdefined(y, :instance)) ||
-           (isa(y, Const) && x === typeof(y.val) && isdefined(x, :instance))
+    elseif (isa(xx, Const) && y === typeof(xx.val) && isdefined(y, :instance)) ||
+           (isa(yy, Const) && x === typeof(yy.val) && isdefined(x, :instance))
         return Const(true)
-    elseif isa(x, Conditional) && isa(y, Const)
-        y.val === false && return Conditional(x.var, x.elsetype, x.vtype)
-        y.val === true && return x
-        return x
-    elseif isa(y, Conditional) && isa(x, Const)
-        x.val === false && return Conditional(y.var, y.elsetype, y.vtype)
-        x.val === true && return y
     end
     return Bool
 end
@@ -374,7 +379,7 @@ end
 add_tfunc(typeof, 1, 1, typeof_tfunc, 0)
 add_tfunc(typeassert, 2, 2,
           function (@nospecialize(v), @nospecialize(t))
-              t, isexact = instanceof_tfunc(t)
+              t = instanceof_tfunc(t)[1]
               t === Any && return v
               if isa(v, Const)
                   if !has_free_typevars(t) && !isa(v.val, t)
@@ -787,7 +792,7 @@ function apply_type_tfunc(@nospecialize(headtypetype), @nospecialize args...)
     tparams = Any[]
     outervars = Any[]
     for i = 1:largs
-        ai = args[i]
+        ai = maybe_widen_conditional(args[i])
         if isType(ai)
             aip1 = ai.parameters[1]
             canconst &= !has_free_typevars(aip1)
@@ -939,7 +944,7 @@ function builtin_tfunction(@nospecialize(f), argtypes::Array{Any,1},
                            sv::Union{InferenceState,Nothing}, params::Params = sv.params)
     isva = !isempty(argtypes) && isvarargtype(argtypes[end])
     if f === tuple
-        for a in argtypes
+        for a in argtypes # TODO: permit Conditional here too
             if !isa(a, Const)
                 return tuple_tfunc(argtypes_to_type(argtypes))
             end
@@ -1080,7 +1085,8 @@ end
 
 # N.B.: typename maps type equivalence classes to a single value
 function typename_static(@nospecialize(t))
-    t = unwrap_unionall(t)
+    t isa Const && return _typename(t.val)
+    t isa Conditional && return Bool.name
+    t = unwrap_unionall(widenconst(t))
     return isType(t) ? _typename(t.parameters[1]) : Core.TypeName
 end
-typename_static(t::Const) = _typename(t.val)

base/compiler/typelimits.jl

@@ -321,6 +321,10 @@ function tmerge(@nospecialize(typea), @nospecialize(typeb))
                 return Conditional(typea.var, vtype, elsetype)
             end
         end
+        val = maybe_extract_const_bool(typea)
+        if val isa Bool && val === maybe_extract_const_bool(typeb)
+            return Const(val)
+        end
         return Bool
     end
     # no special type-inference lattice, join the types

base/iterators.jl

@@ -1035,13 +1035,13 @@ mutable struct Stateful{T, VS}
     end
     @inline function Stateful(itr::T) where {T}
         VS = approx_iter_type(T)
-        new{T, VS}(itr, iterate(itr)::VS, 0)
+        return new{T, VS}(itr, iterate(itr)::VS, 0)
     end
 end
 
 function reset!(s::Stateful{T,VS}, itr::T) where {T,VS}
     s.itr = itr
-    s.nextvalstate = iterate(itr)
+    setfield!(s, :nextvalstate, iterate(itr))
     s.taken = 0
     s
 end
@@ -1057,7 +1057,8 @@ else
     # having to typesubtract
     function doiterate(itr, valstate::Union{Nothing, Tuple{Any, Any}})
         valstate === nothing && return nothing
-        iterate(itr, tail(valstate))
+        val, st = valstate
+        return iterate(itr, st)
     end
     function _approx_iter_type(itrT::Type, vstate::Type)
         vstate <: Union{Nothing, Tuple{Any, Any}} || return Any

test/compiler/compiler.jl

@@ -1189,7 +1189,7 @@ let isa_tfunc = Core.Compiler.T_FFUNC_VAL[
     @test isa_tfunc(typeof(Union{}), Const(Int)) === Const(false) # any result is ok
     @test isa_tfunc(typeof(Union{}), Const(Union{})) === Const(false)
     @test isa_tfunc(typeof(Union{}), typeof(Union{})) === Const(false)
-    @test isa_tfunc(typeof(Union{}), Union{}) === Union{} # any result is ok
+    @test isa_tfunc(typeof(Union{}), Union{}) === Union{}
     @test isa_tfunc(typeof(Union{}), Type{typeof(Union{})}) === Const(true)
     @test isa_tfunc(typeof(Union{}), Const(typeof(Union{}))) === Const(true)
     let c = Conditional(Core.SlotNumber(0), Const(Union{}), Const(Union{}))
@@ -1204,7 +1204,7 @@ let isa_tfunc = Core.Compiler.T_FFUNC_VAL[
     @test isa_tfunc(Val{1}, Type{Val{T}} where T) === Bool
     @test isa_tfunc(Val{1}, DataType) === Bool
     @test isa_tfunc(Any, Const(Any)) === Const(true)
-    @test isa_tfunc(Any, Union{}) === Union{} # any result is ok
+    @test isa_tfunc(Any, Union{}) === Union{}
     @test isa_tfunc(Any, Type{Union{}}) === Const(false)
     @test isa_tfunc(Union{Int64, Float64}, Type{Real}) === Const(true)
     @test isa_tfunc(Union{Int64, Float64}, Type{Integer}) === Bool
@@ -1468,6 +1468,19 @@ result = f24852_kernel(x, y)
 # TODO: test that `expand_early = true` + inflated `method_for_inference_limit_heuristics`
 # can be used to tighten up some inference result.
 
+f26339(T) = T === Union{} ? 1 : ""
+g26339(T) = T === Int ? 1 : ""
+@test Base.return_types(f26339, (Int,)) == Any[String]
+@test Base.return_types(g26339, (Int,)) == Any[String]
+@test Base.return_types(f26339, (Type{Int},)) == Any[String]
+@test Base.return_types(g26339, (Type{Int},)) == Any[Int]
+@test Base.return_types(f26339, (Type{Union{}},)) == Any[Int]
+@test Base.return_types(g26339, (Type{Union{}},)) == Any[String]
+@test Base.return_types(f26339, (typeof(Union{}),)) == Any[Int]
+@test Base.return_types(g26339, (typeof(Union{}),)) == Any[String]
+@test Base.return_types(f26339, (Type,)) == Any[Union{Int, String}]
+@test Base.return_types(g26339, (Type,)) == Any[Union{Int, String}]
+
 # Test that Conditional doesn't get widened to Bool too quickly
 f25261() = (1, 1)
 f25261(s) = i == 1 ? (1, 2) : nothing