Extend PartialTuple for structs

This aims to address the following issue:

Say we have:
```
function foo(b)
    a = 1
    f = ()->Val(a)
    f()
end
```

This infers beatifully, because the closure struct gets `Const`'ed
and thus inter-procedural constant prop takes care of it. However,
if we instead do

```
function foo(b)
    a = 1
    f = ()->(println(b); Val(a))
    f()
end
```

the best inference can tell us about this is that we get `Val`, because
the captured arguments are no longer constant. This leads to significant
inference problems for Zygote, because the backwards pass is always
specified as a closure. Thus, even if constant information is present
in the forward pass, it is often lost for the part of the function
that's the backwards pass, because it has to pass through the closure
struct. This fixes that by keeping track of field types individually.

Author: Keno

base/compiler/abstractinterpretation.jl

@@ -149,6 +149,21 @@ function abstract_call_gf_by_type(@nospecialize(f), argtypes::Vector{Any}, @nosp
     return rettype
 end
 
+
+function const_prop_profitable(arg)
+    # have new information from argtypes that wasn't available from the signature
+    if isa(arg, PartialStruct)
+        for b in arg.fields
+            isconstType(b) && return true
+            const_prop_profitable(b) && return true
+        end
+    elseif !isa(arg, Const) || (isa(arg.val, Symbol) || isa(arg.val, Type) || (!isa(arg.val, String) && isimmutable(arg.val)))
+        # don't consider mutable values or Strings useful constants
+        return true
+    end
+    return false
+end
+
 function abstract_call_method_with_const_args(@nospecialize(rettype), @nospecialize(f), argtypes::Vector{Any}, match::SimpleVector, sv::InferenceState)
     method = match[3]::Method
     nargs::Int = method.nargs
@@ -158,12 +173,8 @@ function abstract_call_method_with_const_args(@nospecialize(rettype), @nospecial
     for a in argtypes
         a = widenconditional(a)
         if has_nontrivial_const_info(a)
-            # have new information from argtypes that wasn't available from the signature
-            if !isa(a, Const) || (isa(a.val, Symbol) || isa(a.val, Type) || (!isa(a.val, String) && isimmutable(a.val)))
-                # don't consider mutable values or Strings useful constants
-                haveconst = true
-                break
-            end
+            haveconst = const_prop_profitable(a)
+            haveconst && break
         end
     end
     haveconst || improvable_via_constant_propagation(rettype) || return Any
@@ -189,7 +200,7 @@ function abstract_call_method_with_const_args(@nospecialize(rettype), @nospecial
             # in this case, see if all of the arguments are constants
             for a in argtypes
                 a = widenconditional(a)
-                if !isa(a, Const) && !isconstType(a)
+                if !isa(a, Const) && !isconstType(a) && !isa(a, PartialStruct)
                     return Any
                 end
             end
@@ -384,7 +395,7 @@ end
 # Union of Tuples of the same length is converted to Tuple of Unions.
 # returns an array of types
 function precise_container_type(@nospecialize(typ), vtypes::VarTable, sv::InferenceState)
-    if isa(typ, PartialTuple)
+    if isa(typ, PartialStruct) && typ.typ.name === Tuple.name
         return typ.fields
     end
 
@@ -498,8 +509,9 @@ end
 # do apply(af, fargs...), where af is a function value
 function abstract_apply(@nospecialize(aft), aargtypes::Vector{Any}, vtypes::VarTable, sv::InferenceState,
                         max_methods = sv.params.MAX_METHODS)
-    if !isa(aft, Const) && (!isType(aft) || has_free_typevars(aft))
-        if !isconcretetype(aft) || (aft <: Builtin)
+    aftw = widenconst(aft)
+    if !isa(aft, Const) && (!isType(aftw) || has_free_typevars(aftw))
+        if !isconcretetype(aftw) || (aftw <: Builtin)
             # non-constant function of unknown type: bail now,
             # since it seems unlikely that abstract_call will be able to do any better after splitting
             # this also ensures we don't call abstract_call_gf_by_type below on an IntrinsicFunction or Builtin
@@ -891,26 +903,37 @@ function abstract_eval(@nospecialize(e), vtypes::VarTable, sv::InferenceState)
         t = instanceof_tfunc(abstract_eval(e.args[1], vtypes, sv))[1]
         if isconcretetype(t) && !t.mutable
             args = Vector{Any}(undef, length(e.args)-1)
-            isconst = true
+            ats = Vector{Any}(undef, length(e.args)-1)
+            anyconst = false
+            allconst = true
             for i = 2:length(e.args)
                 at = abstract_eval(e.args[i], vtypes, sv)
+                if !anyconst
+                    anyconst = has_nontrivial_const_info(at)
+                end
+                ats[i-1] = at
                 if at === Bottom
                     t = Bottom
-                    isconst = false
+                    allconst = anyconst = false
                     break
                 elseif at isa Const
                     if !(at.val isa fieldtype(t, i - 1))
                         t = Bottom
-                        isconst = false
+                        allconst = anyconst = false
                         break
                     end
                     args[i-1] = at.val
                 else
-                    isconst = false
+                    allconst = false
                 end
             end
-            if isconst
-                t = Const(ccall(:jl_new_structv, Any, (Any, Ptr{Cvoid}, UInt32), t, args, length(args)))
+            # For now, don't allow partially initialized Const/PartialStruct
+            if t !== Bottom && fieldcount(t) == length(ats)
+                if allconst
+                    t = Const(ccall(:jl_new_structv, Any, (Any, Ptr{Cvoid}, UInt32), t, args, length(args)))
+                elseif anyconst
+                    t = PartialStruct(t, ats)
+                end
             end
         end
     elseif e.head === :splatnew
@@ -1077,7 +1100,7 @@ function typeinf_local(frame::InferenceState)
             elseif hd === :return
                 pc´ = n + 1
                 rt = widenconditional(abstract_eval(stmt.args[1], s[pc], frame))
-                if !isa(rt, Const) && !isa(rt, Type) && !isa(rt, PartialTuple)
+                if !isa(rt, Const) && !isa(rt, Type) && !isa(rt, PartialStruct)
                     # only propagate information we know we can store
                     # and is valid inter-procedurally
                     rt = widenconst(rt)

base/compiler/inferenceresult.jl

@@ -22,7 +22,7 @@ end
 function is_argtype_match(@nospecialize(given_argtype),
                           @nospecialize(cache_argtype),
                           overridden_by_const::Bool)
-    if isa(given_argtype, Const) || isa(given_argtype, PartialTuple)
+    if isa(given_argtype, Const) || isa(given_argtype, PartialStruct)
         return is_lattice_equal(given_argtype, cache_argtype)
     end
     return !overridden_by_const
@@ -66,7 +66,7 @@ function matching_cache_argtypes(linfo::MethodInstance, ::Nothing)
     nargs::Int = toplevel ? 0 : linfo.def.nargs
     cache_argtypes = Vector{Any}(undef, nargs)
     # First, if we're dealing with a varargs method, then we set the last element of `args`
-    # to the appropriate `Tuple` type or `PartialTuple` instance.
+    # to the appropriate `Tuple` type or `PartialStruct` instance.
     if !toplevel && linfo.def.isva
         if linfo.specTypes == Tuple
             if nargs > 1

base/compiler/ssair/inlining.jl

@@ -583,7 +583,8 @@ function rewrite_apply_exprargs!(ir::IRCode, idx::Int, argexprs::Vector{Any}, at
     for i in 3:length(argexprs)
         def = argexprs[i]
         def_type = atypes[i]
-        if def_type isa PartialTuple
+        if def_type isa PartialStruct
+            # def_type.typ <: Tuple is assumed
             def_atypes = def_type.fields
         else
             def_atypes = Any[]

base/compiler/tfuncs.jl

@@ -716,9 +716,12 @@ function getfield_tfunc(@nospecialize(s00), @nospecialize(name))
             end
         end
         s = typeof(sv)
-    elseif isa(s, PartialTuple)
+    elseif isa(s, PartialStruct)
         if isa(name, Const)
             nv = name.val
+            if isa(nv, Symbol)
+                nv = fieldindex(widenconst(s), nv, false)
+            end
             if isa(nv, Int) && 1 <= nv <= length(s.fields)
                 return s.fields[nv]
             end
@@ -1139,7 +1142,7 @@ function tuple_tfunc(atypes::Vector{Any})
     typ = Tuple{params...}
     # replace a singleton type with its equivalent Const object
     isdefined(typ, :instance) && return Const(typ.instance)
-    return anyinfo ? PartialTuple(typ, atypes) : typ
+    return anyinfo ? PartialStruct(typ, atypes) : typ
 end
 
 function array_type_undefable(@nospecialize(a))

base/compiler/typelattice.jl

@@ -70,7 +70,7 @@ struct StateUpdate
     state::VarTable
 end
 
-struct PartialTuple
+struct PartialStruct
     typ
     fields::Vector{Any} # elements are other type lattice members
 end
@@ -125,8 +125,8 @@ function ⊑(@nospecialize(a), @nospecialize(b))
     elseif isa(b, Conditional)
         return false
     end
-    if isa(a, PartialTuple)
-        if isa(b, PartialTuple)
+    if isa(a, PartialStruct)
+        if isa(b, PartialStruct)
             if !(length(a.fields) == length(b.fields) && a.typ <: b.typ)
                 return false
             end
@@ -137,9 +137,15 @@ function ⊑(@nospecialize(a), @nospecialize(b))
             return true
         end
         return isa(b, Type) && a.typ <: b
-    elseif isa(b, PartialTuple)
+    elseif isa(b, PartialStruct)
         if isa(a, Const)
             nfields(a.val) == length(b.fields) || return false
+            widenconst(b).name === widenconst(a).name || return false
+            # We can skip the subtype check if b is a Tuple, since in that
+            # case, the ⊑ of the elements is sufficient.
+            if b.typ.name !== Tuple.name && !(widenconst(a) <: widenconst(b))
+                return false
+            end
             for i in 1:nfields(a.val)
                 # XXX: let's handle varargs later
                 ⊑(Const(getfield(a.val, i)), b.fields[i]) || return false
@@ -173,15 +179,16 @@ end
 # `a ⊑ b && b ⊑ a` but with extra performance optimizations.
 function is_lattice_equal(@nospecialize(a), @nospecialize(b))
     a === b && return true
-    if isa(a, PartialTuple)
-        isa(b, PartialTuple) || return false
+    if isa(a, PartialStruct)
+        isa(b, PartialStruct) || return false
         length(a.fields) == length(b.fields) || return false
+        widenconst(a) == widenconst(b) || return false
         for i in 1:length(a.fields)
             is_lattice_equal(a.fields[i], b.fields[i]) || return false
         end
         return true
     end
-    isa(b, PartialTuple) && return false
+    isa(b, PartialStruct) && return false
     a isa Const && return false
     b isa Const && return false
     return a ⊑ b && b ⊑ a
@@ -200,7 +207,7 @@ function widenconst(c::Const)
 end
 widenconst(m::MaybeUndef) = widenconst(m.typ)
 widenconst(c::PartialTypeVar) = TypeVar
-widenconst(t::PartialTuple) = t.typ
+widenconst(t::PartialStruct) = t.typ
 widenconst(@nospecialize(t)) = t
 
 issubstate(a::VarState, b::VarState) = (a.typ ⊑ b.typ && a.undef <= b.undef)

base/compiler/typelimits.jl

@@ -317,6 +317,27 @@ function tmerge(@nospecialize(typea), @nospecialize(typeb))
         end
         return Bool
     end
+    if (isa(typea, PartialStruct) || isa(typea, Const)) &&
+       (isa(typeb, PartialStruct) || isa(typeb, Const)) &&
+        widenconst(typea) === widenconst(typeb)
+
+       typea_nfields = nfields_tfunc(typea)
+       typeb_nfields = nfields_tfunc(typeb)
+       if !isa(typea_nfields, Const) || !isa(typea_nfields, Const) || typea_nfields.val !== typeb_nfields.val
+            return widenconst(typea)
+       end
+
+       type_nfields = typea_nfields.val::Int
+       fields = Vector{Any}(undef, type_nfields)
+       anyconst = false
+       for i = 1:type_nfields
+            fields[i] = tmerge(getfield_tfunc(typea, Const(i)),
+                               getfield_tfunc(typeb, Const(i)))
+            anyconst |= has_nontrivial_const_info(fields[i])
+       end
+       return anyconst ? PartialStruct(widenconst(typea), fields) :
+            widenconst(typea)
+    end
     # no special type-inference lattice, join the types
     typea, typeb = widenconst(typea), widenconst(typeb)
     typea === typeb && return typea

base/compiler/typeutils.jl

@@ -33,7 +33,7 @@ function issingletontype(@nospecialize t)
 end
 
 function has_nontrivial_const_info(@nospecialize t)
-    isa(t, PartialTuple) && return true
+    isa(t, PartialStruct) && return true
     return isa(t, Const) && !isdefined(typeof(t.val), :instance) && !(isa(t.val, Type) && issingletontype(t.val))
 end
 

test/compiler/inference.jl

@@ -1338,8 +1338,8 @@ let egal_tfunc
     @test egal_tfunc(Union{Int64, Float64}, AbstractArray) === Const(false)
 end
 
-using Core.Compiler: PartialTuple, nfields_tfunc, sizeof_tfunc, sizeof_nothrow
-let PT = PartialTuple(Tuple{Int64,UInt64}, Any[Const(10, false), UInt64])
+using Core.Compiler: PartialStruct, nfields_tfunc, sizeof_tfunc, sizeof_nothrow
+let PT = PartialStruct(Tuple{Int64,UInt64}, Any[Const(10, false), UInt64])
     @test sizeof_tfunc(PT) === Const(16, false)
     @test nfields_tfunc(PT) === Const(2, false)
     @test sizeof_nothrow(PT) === true
@@ -2261,3 +2261,30 @@ f_incr(x::Tuple, y::Tuple, args...) = f_incr((x, y), args...)
 f_incr(x::Tuple) = x
 @test @inferred(f_incr((), (), (), (), (), (), (), ())) ==
     ((((((((), ()), ()), ()), ()), ()), ()), ())
+
+# Test PartialStruct for closures
+@noinline use30783(x) = nothing
+function foo30783(b)
+    a = 1
+    f = ()->(use30783(b); Val(a))
+    f()
+end
+@test @inferred(foo30783(2)) == Val(1)
+
+# PartialStruct tmerge
+using Core.Compiler: PartialStruct, tmerge, Const, ⊑
+struct FooPartial
+    a::Int
+    b::Int
+    c::Int
+end
+let PT1 = PartialStruct(FooPartial, Any[Const(1), Const(2), Int]),
+    PT2 = PartialStruct(FooPartial, Any[Const(1), Int, Int]),
+    PT3 = PartialStruct(FooPartial, Any[Const(1), Int, Const(3)])
+
+    @test PT1 ⊑ PT2
+    @test !(PT1 ⊑ PT3) && !(PT2 ⊑ PT1)
+    let (==) = (a, b)->(a ⊑ b && b ⊑ a)
+        @test tmerge(PT1, PT3) == PT2
+    end
+end