inference: forward Conditional inter-procedurally

base/compiler/inferenceresult.jl

@@ -3,20 +3,57 @@
 function is_argtype_match(@nospecialize(given_argtype),
  @nospecialize(cache_argtype),
  overridden_by_const::Bool)
- if isa(given_argtype, Const) || isa(given_argtype, PartialStruct) || isa(given_argtype, PartialOpaque)
+ if is_forwardable_argtype(given_argtype)
  return is_lattice_equal(given_argtype, cache_argtype)
  end
  return !overridden_by_const
 end
 
+function is_forwardable_argtype(@nospecialize x)
+ return isa(x, Const) ||
+ isa(x, Conditional) ||
+ isa(x, PartialStruct) ||
+ isa(x, PartialOpaque)
+end
+
 # In theory, there could be a `cache` containing a matching `InferenceResult`
 # for the provided `linfo` and `given_argtypes`. The purpose of this function is
 # to return a valid value for `cache_lookup(linfo, argtypes, cache).argtypes`,
 # so that we can construct cache-correct `InferenceResult`s in the first place.
-function matching_cache_argtypes(linfo::MethodInstance, given_argtypes::Vector, va_override::Bool)
+function matching_cache_argtypes(
+ linfo::MethodInstance, (; fargs, argtypes)::ArgInfo, va_override::Bool)
  @assert isa(linfo.def, Method) # ensure the next line works
  nargs::Int = linfo.def.nargs
- given_argtypes = anymap(widenconditional, given_argtypes)
+ cache_argtypes, overridden_by_const = matching_cache_argtypes(linfo, nothing, va_override)
+ given_argtypes = Vector{Any}(undef, length(argtypes))
+ local condargs = nothing
+ for i in 1:length(argtypes)
+ argtype = argtypes[i]
+ # forward `Conditional` if it conveys a constraint on any other argument
+ if isa(argtype, Conditional) && fargs !== nothing
+ cnd = argtype
+ slotid = find_constrained_arg(cnd, fargs)
+ if slotid !== nothing
+ # using union-split signature, we may be able to narrow down `Conditional`
+ sigt = widenconst(slotid > nargs ? argtypes[slotid] : cache_argtypes[slotid])
+ vtype = tmeet(cnd.vtype, sigt)
+ elsetype = tmeet(cnd.elsetype, sigt)
+ if vtype === Bottom && elsetype === Bottom
+ # we accidentally proved this method match is impossible
+ # TODO bail out here immediately rather than just propagating Bottom ?
+ given_argtypes[i] = Bottom
+ else
+ if condargs === nothing
+ condargs = Tuple{Int,Int}[]
+ end
+ push!(condargs, (slotid, i))
+ given_argtypes[i] = Conditional(SlotNumber(slotid), vtype, elsetype)
+ end
+ continue
+ end
+ end
+ given_argtypes[i] = widenconditional(argtype)
+ end
  isva = va_override || linfo.def.isva
  if isva || isvarargtype(given_argtypes[end])
  isva_given_argtypes = Vector{Any}(undef, nargs)
@@ -30,15 +67,22 @@ function matching_cache_argtypes(linfo::MethodInstance, given_argtypes::Vector,
  last = nargs
  end
  isva_given_argtypes[nargs] = tuple_tfunc(given_argtypes[last:end])
+ # invalidate `Conditional` imposed on varargs
+ if condargs !== nothing
+ for (slotid, i) in condargs
+ if slotid ≥ last
+ isva_given_argtypes[i] = widenconditional(isva_given_argtypes[i])
+ end
+ end
+ end
  end
  given_argtypes = isva_given_argtypes
  end
  @assert length(given_argtypes) == nargs
- cache_argtypes, overridden_by_const = matching_cache_argtypes(linfo, nothing, va_override)
  for i in 1:nargs
  given_argtype = given_argtypes[i]
  cache_argtype = cache_argtypes[i]
- if !is_argtype_match(given_argtype, cache_argtype, overridden_by_const[i])
+ if !is_argtype_match(given_argtype, cache_argtype, false)
  # prefer the argtype we were given over the one computed from `linfo`
  cache_argtypes[i] = given_argtype
  overridden_by_const[i] = true

base/compiler/tfuncs.jl

@@ -964,7 +964,7 @@ function abstract_modifyfield!(interp::AbstractInterpreter, argtypes::Vector{Any
  v = unwrapva(argtypes[5])
  TF = getfield_tfunc(o, f)
  push!(sv.ssavalue_uses[sv.currpc], sv.currpc) # temporarily disable `call_result_unused` check for this call
- callinfo = abstract_call(interp, nothing, Any[op, TF, v], sv, #=max_methods=# 1)
+ callinfo = abstract_call(interp, ArgInfo(nothing, Any[op, TF, v]), sv, #=max_methods=# 1)
  pop!(sv.ssavalue_uses[sv.currpc], sv.currpc)
  TF2 = tmeet(callinfo.rt, widenconst(TF))
  if TF2 === Bottom
@@ -1747,7 +1747,7 @@ function return_type_tfunc(interp::AbstractInterpreter, argtypes::Vector{Any}, s
  if contains_is(argtypes_vec, Union{})
  return CallMeta(Const(Union{}), false)
  end
- call = abstract_call(interp, nothing, argtypes_vec, sv, -1)
+ call = abstract_call(interp, ArgInfo(nothing, argtypes_vec), sv, -1)
  info = verbose_stmt_info(interp) ? ReturnTypeCallInfo(call.info) : false
  rt = widenconditional(call.rt)
  if isa(rt, Const)

base/compiler/types.jl

@@ -17,6 +17,11 @@ If `interp` is an `AbstractInterpreter`, it is expected to provide at least the
 """
 abstract type AbstractInterpreter end
 
+struct ArgInfo
+ fargs::Union{Nothing,Vector{Any}}
+ argtypes::Vector{Any}
+end
+
 """
  InferenceResult
 
@@ -29,8 +34,10 @@ mutable struct InferenceResult
  result # ::Type, or InferenceState if WIP
  src #::Union{CodeInfo, OptimizationState, Nothing} # if inferred copy is available
  valid_worlds::WorldRange # if inference and optimization is finished
- function InferenceResult(linfo::MethodInstance, given_argtypes = nothing, va_override=false)
- argtypes, overridden_by_const = matching_cache_argtypes(linfo, given_argtypes, va_override)
+ function InferenceResult(linfo::MethodInstance,
+ arginfo::Union{Nothing,ArgInfo} = nothing,
+ va_override::Bool = false)
+ argtypes, overridden_by_const = matching_cache_argtypes(linfo, arginfo, va_override)
  return new(linfo, argtypes, overridden_by_const, Any, nothing, WorldRange())
  end
 end

base/compiler/typeutils.jl

@@ -259,15 +259,6 @@ unioncomplexity(u::UnionAll) = max(unioncomplexity(u.body)::Int, unioncomplexity
 unioncomplexity(t::TypeofVararg) = isdefined(t, :T) ? unioncomplexity(t.T)::Int : 0
 unioncomplexity(@nospecialize(x)) = 0
 
-function improvable_via_constant_propagation(@nospecialize(t))
- if isconcretetype(t) && t <: Tuple
- for p in t.parameters
- p === DataType && return true
- end
- end
- return false
-end
-
 # convert a Union of Tuple types to a Tuple of Unions
 function unswitchtupleunion(u::Union)
  ts = uniontypes(u)

test/compiler/inference.jl

@@ -2000,6 +2000,61 @@ function _g_ifelse_isa_()
 end
 @test Base.return_types(_g_ifelse_isa_, ()) == [Int]
 
+@testset "Conditional forwarding" begin
+ # forward `Conditional` if it conveys a constraint on any other argument
+ ifelselike(cnd, x, y) = cnd ? x : y
+
+ @test Base.return_types((Any,Int,)) do x, y
+ ifelselike(isa(x, Int), x, y)
+ end |> only == Int
+
+ # should work nicely with union-split
+ @test Base.return_types((Union{Int,Nothing},)) do x
+ ifelselike(isa(x, Int), x, 0)
+ end |> only == Int
+
+ @test Base.return_types((Any,Int)) do x, y
+ ifelselike(!isa(x, Int), y, x)
+ end |> only == Int
+
+ @test Base.return_types((Any,Int)) do x, y
+ a = ifelselike(x === 0, x, 0) # ::Const(0)
+ if a == 0
+ return y
+ else
+ return nothing # dead branch
+ end
+ end |> only == Int
+
+ # pick up the first if there are multiple constrained arguments
+ @test Base.return_types((Any,)) do x
+ ifelselike(isa(x, Int), x, x)
+ end |> only == Any
+
+ # just propagate multiple constraints
+ ifelselike2(cnd1, cnd2, x, y, z) = cnd1 ? x : cnd2 ? y : z
+ @test Base.return_types((Any,Any)) do x, y
+ ifelselike2(isa(x, Int), isa(y, Int), x, y, 0)
+ end |> only == Int
+
+ # work with `invoke`
+ @test Base.return_types((Any,Any)) do x, y
+ Base.@invoke ifelselike(isa(x, Int), x, y::Int)
+ end |> only == Int
+
+ # don't be confused with vararg method
+ vacond(cnd, va...) = cnd ? va : 0
+ @test Base.return_types((Any,)) do x
+ # at runtime we will see `va::Tuple{Tuple{Int,Int}, Tuple{Int,Int}}`
+ vacond(isa(x, Tuple{Int,Int}), x, x)
+ end |> only == Union{Int,Tuple{Any,Any}}
+
+ # demonstrate extra constraint propagation for Base.ifelse
+ @test Base.return_types((Any,Int,)) do x, y
+ ifelse(isa(x, Int), x, y)
+ end |> only == Int
+end
+
 # Equivalence of Const(T.instance) and T for singleton types
 @test Const(nothing) ⊑ Nothing && Nothing ⊑ Const(nothing)