inference: enable constant propagation for union-split signatures

The inference precision of certain functions really relies on constant
propagation, but currently constant prop' won't happen when a call
signature is union split and so sometimes inference ends up looser
return type: e.g.
```julia
julia>
Base.return_types((Union{Tuple{Int,Nothing},Tuple{Int,Missing}},)) do t
           a, b = t
           a # I expected a::Int, but a::Union{Missing,Nothing,Int}
       end |> first
Union{Missing, Nothing, Int64}
```

This PR:
- enables constant prop' for union signatures (with "non-Bottom"
results)
- adds special cases for some functions (see `force_constant_prop'`) so
  that we keep inferring on them even after the return type from
  non-constant analysis has grown to `Any`, since constant analysis can
  improve the result later on

The added test cases will should showcase the cases where the
inference result could be improved by that.

---

Here is a sample benchmark of the impact of this PR on latency,
from which I guess this PR is acceptable ?

> build time: master (26a721b)
```bash
Sysimage built. Summary:
Total ───────  57.320450 seconds
Base: ───────  24.371516 seconds 42.518%
Stdlibs: ────  32.947400 seconds 57.4793%
    JULIA usr/lib/julia/sys-o.a
Generating REPL precompile statements... 30/30
Executing precompile statements... 1379/1379
Precompilation complete. Summary:
Total ─────── 101.325866 seconds
Generation ──  74.901320 seconds 73.9212%
Execution ───  26.424546 seconds 26.0788%
    LINK usr/lib/julia/sys.dylib
```

> build time: this PR
```bash
Sysimage built. Summary:
Total ───────  58.387562 seconds
Base: ───────  24.403206 seconds 41.7952%
Stdlibs: ────  33.982657 seconds 58.2019%
    JULIA usr/lib/julia/sys-o.a
Generating REPL precompile statements... 30/30
Executing precompile statements... 1362/1362
Precompilation complete. Summary:
Total ───────  99.659102 seconds
Generation ──  73.983763 seconds 74.2368%
Execution ───  25.675339 seconds 25.7632%
    LINK usr/lib/julia/sys.dylib
```

> first time to plot: master (26a721b)
```julia
julia> using Plots; @time plot(rand(10,3))
  3.614168 seconds (5.47 M allocations: 324.564 MiB, 5.73% gc time,
53.02% compilation time)
```

> first time to plot: this PR
```julia
julia> using Plots; @time plot(rand(10,3))
  3.557919 seconds (5.53 M allocations: 328.812 MiB, 2.89% gc time,
51.94% compilation time)
```

---

NOTE: this PR is an alternative for JuliaLang#39296

Author: aviatesk

base/compiler/abstractinterpretation.jl

@@ -97,8 +97,8 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
     napplicable = length(applicable)
     rettype = Bottom
     edgecycle = false
-    edges = Any[]
-    nonbot = 0  # the index of the only non-Bottom inference result if > 0
+    edges = MethodInstance[]
+    nonbots = Int[]  # the indexes of non-Bottom inference results (will be re-analyzed with constant arguments)
     seen = 0    # number of signatures actually inferred
     istoplevel = sv.linfo.def isa Module
     multiple_matches = napplicable > 1
@@ -111,6 +111,8 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
         end
     end
 
+    keep_inference = force_constant_prop(interp, f, argtypes)
+
     for i in 1:napplicable
         match = applicable[i]::MethodMatch
         method = match.method
@@ -122,8 +124,8 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
             break
         end
         sigtuple = unwrap_unionall(sig)::DataType
-        splitunions = false
         this_rt = Bottom
+        splitunions = false
         # TODO: splitunions = 1 < unionsplitcost(sigtuple.parameters) * napplicable <= InferenceParams(interp).MAX_UNION_SPLITTING
         # currently this triggers a bug in inference recursion detection
         if splitunions
@@ -135,7 +137,9 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
                 end
                 edgecycle |= edgecycle1::Bool
                 this_rt = tmerge(this_rt, rt)
-                this_rt === Any && break
+                if !(this_rt !== Any || keep_inference)
+                    break
+                end
             end
         else
             this_rt, edgecycle1, edge = abstract_call_method(interp, method, sig, match.sparams, multiple_matches, sv)
@@ -145,29 +149,32 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
             end
         end
         if this_rt !== Bottom
-            if nonbot === 0
-                nonbot = i
-            else
-                nonbot = -1
-            end
+            push!(nonbots, i)
         end
         seen += 1
         rettype = tmerge(rettype, this_rt)
-        rettype === Any && break
+        if !(rettype !== Any || keep_inference)
+            break
+        end
     end
-    # try constant propagation if only 1 method is inferred to non-Bottom
+    # try constant propagation for non-Bottom inference results
     # this is in preparation for inlining, or improving the return result
     is_unused = call_result_unused(sv)
-    if nonbot > 0 && seen == napplicable && (!edgecycle || !is_unused) &&
+    if !isempty(nonbots) && seen == napplicable && (!edgecycle || !is_unused) &&
             is_improvable(rettype) && InferenceParams(interp).ipo_constant_propagation
         # if there's a possibility we could constant-propagate a better result
         # (hopefully without doing too much work), try to do that now
         # TODO: it feels like this could be better integrated into abstract_call_method / typeinf_edge
-        const_rettype = abstract_call_method_with_const_args(interp, rettype, f, argtypes, applicable[nonbot]::MethodMatch, sv, edgecycle)
-        if const_rettype ⊑ rettype
-            # use the better result, if it's a refinement of rettype
-            rettype = const_rettype
+        const_rettype = Bottom
+        for nonbot in nonbots
+            this_rt = abstract_call_method_with_const_args(interp, rettype, f, argtypes, applicable[nonbot]::MethodMatch, sv, edgecycle)
+            const_rettype = tmerge(const_rettype, this_rt)
+            if !(const_rettype !== rettype && const_rettype ⊑ rettype)
+                const_rettype = rettype
+                break
+            end
         end
+        rettype = const_rettype
     end
     if is_unused && !(rettype === Bottom)
         add_remark!(interp, sv, "Call result type was widened because the return value is unused")
@@ -195,6 +202,21 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
     return CallMeta(rettype, info)
 end
 
+# some functions heavily rely on constant propagation to get precise return type,
+# so we better to keep inference on them even after the return type has grown to `Any`
+function force_constant_prop(interp::AbstractInterpreter, @nospecialize(f), argtypes::Vector{Any})
+    istopfunction(f, :getproperty) && return true
+    istopfunction(f, :setproperty!) && return true
+    la = length(argtypes)
+    # tuple indexing
+    if la ≥ 3 &&
+       (a3 = argtypes[3]; isa(a3, Const) && isa(a3.val, Int)) &&
+       (a2 = argtypes[2]; a2 ⊑ Tuple && isa(a2, Union)) &&
+       istopfunction(f, :getindex)
+        return true
+    end
+    return false
+end
 
 function const_prop_profitable(@nospecialize(arg))
     # have new information from argtypes that wasn't available from the signature

test/compiler/inference.jl

@@ -7,6 +7,11 @@ isdispatchelem(@nospecialize x) = !isa(x, Type) || Core.Compiler.isdispatchelem(
 using Random, Core.IR
 using InteractiveUtils: code_llvm
 
+macro evaltoplevel(ex)
+    m = Core.eval(__module__, :(module $(gensym()) end))::Module
+    return QuoteNode(Core.eval(m, ex))
+end
+
 f39082(x::Vararg{T}) where {T <: Number} = x[1]
 let ast = only(code_typed(f39082, Tuple{Vararg{Rational}}))[1]
     @test ast.slottypes == Any[Const(f39082), Tuple{Vararg{Rational}}]
@@ -3008,3 +3013,66 @@ f38888() = S38888(Base.inferencebarrier(3))
 @test f38888() isa S38888
 g38888() = S38888(Base.inferencebarrier(3), nothing)
 @test g38888() isa S38888
+
+@testset "constant prop' for union split signature" begin
+    # indexing into tuples really relies on constant prop', and we will get looser result
+    # (`Union{Int,String,Char}`) if constant prop' doesn't happen for splitunion signatures
+    tt = (Union{Tuple{Int,String},Tuple{Int,Char}},)
+
+    # `getindex`
+    @test Base.return_types(tt) do t
+        getindex(t, 1)
+    end == Any[Int]
+    @test Base.return_types(tt) do t
+        getindex(t, 2)
+    end == Any[Union{String,Char}]
+
+    # `indexed_iterate`
+    @test Base.return_types(tt) do t
+        a, b = t
+        a
+    end == Any[Int]
+    @test Base.return_types(tt) do t
+        a, b = t
+        b
+    end == Any[Union{String,Char}]
+
+    @test @evaltoplevel begin
+        struct F32
+            val::Float32
+            _v::Int
+        end
+        struct F64
+            val::Float64
+            _v::Int
+        end
+        Base.return_types((Union{F32,F64},)) do f
+            f.val
+        end == Any[Union{Float32,Float64}]
+    end
+
+    # for some of functions, we should force constant propagation by keeping inference
+    # even after the return type of non-constant analysis has grown to `Any` to get precise
+    # return type of some functions
+    @testset "force constant prop'" begin
+        # `getproperty`
+        @test @evaltoplevel begin
+            struct F32
+                val::Float32
+                _v
+            end
+            struct F64
+                val::Float64
+                _v
+            end
+            Base.return_types((Union{F32,F64},)) do f
+                f.val
+            end == Any[Union{Float32,Float64}]
+        end
+
+        # `getindex(::Tuple, ::Const(::Int))`
+        @test Base.return_types((Union{Tuple{Nothing,Any,Any},Tuple{Nothing,Any}},)) do t
+            getindex(t, 1)
+        end == Any[Nothing]
+    end
+end