Relax constraints on inlining for some single calls (JuliaLang#43113)

Co-authored-by: Shuhei Kadowaki <[email protected]>
LilithHafner · Mar 8, 2022 · ec79e87 · ec79e87
1 parent 8e85c56
commit ec79e87
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 27 deletions.
diff --git a/base/compiler/ssair/inlining.jl b/base/compiler/ssair/inlining.jl
@@ -1176,26 +1176,24 @@ function analyze_single_call!(
         end
     end
 
-    # if the signature is fully covered and there is only one applicable method,
+    # if the signature is fully or mostly covered and there is only one applicable method,
     # we can try to inline it even if the signature is not a dispatch tuple
-    if atype <: signature_union
-        if length(cases) == 0 && only_method isa Method
-            if length(infos) > 1
-                (metharg, methsp) = ccall(:jl_type_intersection_with_env, Any, (Any, Any),
-                    atype, only_method.sig)::SimpleVector
-                match = MethodMatch(metharg, methsp::SimpleVector, only_method, true)
-            else
-                meth = meth::MethodLookupResult
-                @assert length(meth) == 1
-                match = meth[1]
-            end
-            item = analyze_method!(match, argtypes, flag, state)
-            item === nothing && return
-            push!(cases, InliningCase(match.spec_types, item))
-            fully_covered = true
+    if length(cases) == 0 && only_method isa Method
+        if length(infos) > 1
+            (metharg, methsp) = ccall(:jl_type_intersection_with_env, Any, (Any, Any),
+                atype, only_method.sig)::SimpleVector
+            match = MethodMatch(metharg, methsp::SimpleVector, only_method, true)
+        else
+            meth = meth::MethodLookupResult
+            @assert length(meth) == 1
+            match = meth[1]
         end
+        item = analyze_method!(match, argtypes, flag, state)
+        item === nothing && return
+        push!(cases, InliningCase(match.spec_types, item))
+        fully_covered = match.fully_covers
     else
-        fully_covered = false
+        fully_covered &= atype <: signature_union
     end
 
     # If we only have one case and that case is fully covered, we may either
@@ -1244,17 +1242,15 @@ function maybe_handle_const_call!(
 
     # if the signature is fully covered and there is only one applicable method,
     # we can try to inline it even if the signature is not a dispatch tuple
-    if atype <: signature_union
-        if length(cases) == 0 && length(results) == 1
-            (; mi) = item = InliningTodo(results[1]::InferenceResult, argtypes)
-            state.mi_cache !== nothing && (item = resolve_todo(item, state, flag))
-            validate_sparams(mi.sparam_vals) || return true
-            item === nothing && return true
-            push!(cases, InliningCase(mi.specTypes, item))
-            fully_covered = true
-        end
+    if length(cases) == 0 && length(results) == 1
+        (; mi) = item = InliningTodo(results[1]::InferenceResult, argtypes)
+        state.mi_cache !== nothing && (item = resolve_todo(item, state, flag))
+        validate_sparams(mi.sparam_vals) || return true
+        item === nothing && return true
+        push!(cases, InliningCase(mi.specTypes, item))
+        fully_covered = atype <: mi.specTypes
     else
-        fully_covered = false
+        fully_covered &= atype <: signature_union
     end
 
     # If we only have one case and that case is fully covered, we may either

diff --git a/test/compiler/inline.jl b/test/compiler/inline.jl
@@ -817,3 +817,46 @@ let
     invoke(xs) = validate_unionsplit_inlining(true, xs[1])
     @test invoke(Any[10]) === false
 end
+
+# issue 43104
+
+@inline isGoodType(@nospecialize x::Type) =
+    x !== Any && !(@noinline Base.has_free_typevars(x))
+let # aggressive inlining of single, abstract method match
+    src = code_typed((Type, Any,)) do x, y
+        isGoodType(x), isGoodType(y)
+    end |> only |> first
+    # both callsites should be inlined
+    @test count(isinvoke(:has_free_typevars), src.code) == 2
+    # `isGoodType(y::Any)` isn't fully covered, thus a runtime type check and fallback dynamic dispatch should be inserted
+    @test count(iscall((src,isGoodType)), src.code) == 1
+end
+
+@inline isGoodType2(cnd, @nospecialize x::Type) =
+    x !== Any && !(@noinline (cnd ? Core.Compiler.isType : Base.has_free_typevars)(x))
+let # aggressive inlining of single, abstract method match (with constant-prop'ed)
+    src = code_typed((Type, Any,)) do x, y
+        isGoodType2(true, x), isGoodType2(true, y)
+    end |> only |> first
+    # both callsite should be inlined with constant-prop'ed result
+    @test count(isinvoke(:isType), src.code) == 2
+    @test count(isinvoke(:has_free_typevars), src.code) == 0
+    # `isGoodType(y::Any)` isn't fully convered, thus a runtime type check and fallback dynamic dispatch should be inserted
+    @test count(iscall((src,isGoodType2)), src.code) == 1
+end
+
+@noinline function checkBadType!(@nospecialize x::Type)
+    if x === Any || Base.has_free_typevars(x)
+        println(x)
+    end
+    return nothing
+end
+let # aggressive static dispatch of single, abstract method match
+    src = code_typed((Type, Any,)) do x, y
+        checkBadType!(x), checkBadType!(y)
+    end |> only |> first
+    # both callsites should be resolved statically
+    @test count(isinvoke(:checkBadType!), src.code) == 2
+    # `checkBadType!(y::Any)` isn't fully covered, thus a runtime type check and fallback dynamic dispatch should be inserted
+    @test count(iscall((src,checkBadType!)), src.code) == 1
+end