Allow PartialStruct on inference edges (JuliaLang#36521)

* Allow PartialStruct on inference edges At the moment, we only allow types, or completely Const inference results to propagate along inference edges. In particular PartialStruct is disallowed. I noticed some inference sub-optimalities as a result (e.g. in iterating `Pair`), but more generally, I would prefer if adding an additional return value somewhere (i.e. returning a tuple of two values, rather than a single value) would not degrade the inference precision for the previously returned value. This fixes the issue by allowing `PartialStruct` to propagate along inference edges. * Make sure to widenconst in return_type_tfunc
caleb-allen · Jul 9, 2020 · 0e74b3e · 0e74b3e
1 parent 37cadb4
commit 0e74b3e
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 13 deletions.
diff --git a/base/compiler/tfuncs.jl b/base/compiler/tfuncs.jl
@@ -1492,18 +1492,21 @@ function return_type_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, v
                     if isa(rt, Const)
                         # output was computed to be constant
                         return Const(typeof(rt.val))
-                    elseif hasuniquerep(rt) || rt === Bottom
-                        # output type was known for certain
-                        return Const(rt)
-                    elseif (isa(tt, Const) || isconstType(tt)) &&
-                           (isa(aft, Const) || isconstType(aft))
-                        # input arguments were known for certain
-                        # XXX: this doesn't imply we know anything about rt
-                        return Const(rt)
-                    elseif isType(rt)
-                        return Type{rt}
                     else
-                        return Type{<:widenconst(rt)}
+                        rt = widenconst(rt)
+                        if hasuniquerep(rt) || rt === Bottom
+                            # output type was known for certain
+                            return Const(rt)
+                        elseif (isa(tt, Const) || isconstType(tt)) &&
+                            (isa(aft, Const) || isconstType(aft))
+                            # input arguments were known for certain
+                            # XXX: this doesn't imply we know anything about rt
+                            return Const(rt)
+                        elseif isType(rt)
+                            return Type{rt}
+                        else
+                            return Type{<:rt}
+                        end
                     end
                 end
             end

diff --git a/base/compiler/typeinfer.jl b/base/compiler/typeinfer.jl
@@ -94,6 +94,9 @@ function CodeInstance(result::InferenceResult, min_valid::UInt, max_valid::UInt,
         elseif isconstType(result.result)
             rettype_const = result.result.parameters[1]
             const_flags = 0x2
+        elseif isa(result.result, PartialStruct)
+            rettype_const = (result.result::PartialStruct).fields
+            const_flags = 0x2
         else
             rettype_const = nothing
             const_flags = 0x00
@@ -493,7 +496,11 @@ function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize
     if code isa CodeInstance # return existing rettype if the code is already inferred
         update_valid_age!(min_world(code), max_world(code), caller)
         if isdefined(code, :rettype_const)
-            return Const(code.rettype_const), mi
+            if isa(code.rettype_const, Vector{Any}) && !(Vector{Any} <: code.rettype)
+                return PartialStruct(code.rettype, code.rettype_const), mi
+            else
+                return Const(code.rettype_const), mi
+            end
         else
             return code.rettype, mi
         end
@@ -532,7 +539,7 @@ function typeinf_edge(interp::AbstractInterpreter, method::Method, @nospecialize
 end
 
 function widenconst_bestguess(bestguess)
-    !isa(bestguess, Const) && !isa(bestguess, Type) && return widenconst(bestguess)
+    !isa(bestguess, Const) && !isa(bestguess, PartialStruct) && !isa(bestguess, Type) && return widenconst(bestguess)
     return bestguess
 end
 

diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl
@@ -2668,3 +2668,9 @@ end
 f36531(args...) = tuple((args...)...)
 @test @inferred(f36531(1,2,3)) == (1,2,3)
 @test code_typed(f36531, Tuple{Vararg{Int}}) isa Vector
+
+# PartialStruct results on typeinf edges
+partial_return_1(x) = (x, 1)
+partial_return_2(x) = Val{partial_return_1(x)[2]}
+
+@test Base.return_types(partial_return_2, (Int,)) == Any[Type{Val{1}}]