From 8897ee0b97bd5feadfd11cabc9184e673db38a1a Mon Sep 17 00:00:00 2001
From: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Date: Thu, 22 Feb 2024 20:41:07 +0900
Subject: [PATCH] make `code_lowered` type stable (#53416)

---
 base/expr.jl       |  4 ++--
 base/reflection.jl | 20 +++++++++++---------
 test/reflection.jl |  3 +++
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/base/expr.jl b/base/expr.jl
index b90f75e75a76f4..94abe5b2eab67b 100644
--- a/base/expr.jl
+++ b/base/expr.jl
@@ -79,7 +79,7 @@ function copy(c::CodeInfo)
     cnew.slotnames = copy(cnew.slotnames)
     cnew.slotflags = copy(cnew.slotflags)
     if cnew.slottypes !== nothing
-        cnew.slottypes = copy(cnew.slottypes)
+        cnew.slottypes = copy(cnew.slottypes::Vector{Any})
     end
     cnew.codelocs  = copy(cnew.codelocs)
     cnew.linetable = copy(cnew.linetable::Union{Vector{Any},Vector{Core.LineInfoNode}})
@@ -1019,7 +1019,7 @@ function remove_linenums!(@nospecialize ex)
         return ex
     elseif ex isa CodeInfo
         ex.codelocs .= 0
-        length(ex.linetable) > 1 && resize!(ex.linetable, 1)
+        length(ex.linetable::Vector) > 1 && resize!(ex.linetable::Vector, 1)
         return ex
     else
         return ex
diff --git a/base/reflection.jl b/base/reflection.jl
index 7a7d60ba25f1ea..202573b02f1bfd 100644
--- a/base/reflection.jl
+++ b/base/reflection.jl
@@ -1127,20 +1127,23 @@ function code_lowered(@nospecialize(f), @nospecialize(t=Tuple); generated::Bool=
     end
     world = get_world_counter()
     world == typemax(UInt) && error("code reflection cannot be used from generated functions")
-    return map(method_instances(f, t, world)) do m
+    ret = CodeInfo[]
+    for m in method_instances(f, t, world)
         if generated && hasgenerator(m)
             if may_invoke_generator(m)
-                return ccall(:jl_code_for_staged, Any, (Any, UInt), m, world)::CodeInfo
+                code = ccall(:jl_code_for_staged, Any, (Any, UInt), m, world)::CodeInfo
             else
                 error("Could not expand generator for `@generated` method ", m, ". ",
                       "This can happen if the provided argument types (", t, ") are ",
                       "not leaf types, but the `generated` argument is `true`.")
             end
+        else
+            code = uncompressed_ir(m.def::Method)
+            debuginfo === :none && remove_linenums!(code)
         end
-        code = uncompressed_ir(m.def::Method)
-        debuginfo === :none && remove_linenums!(code)
-        return code
+        push!(ret, code)
     end
+    return ret
 end
 
 hasgenerator(m::Method) = isdefined(m, :generator)
@@ -1331,7 +1334,7 @@ function method_instances(@nospecialize(f), @nospecialize(t), world::UInt)
     # this make a better error message than the typeassert that follows
     world == typemax(UInt) && error("code reflection cannot be used from generated functions")
     for match in _methods_by_ftype(tt, -1, world)::Vector
-        instance = Core.Compiler.specialize_method(match)
+        instance = Core.Compiler.specialize_method(match::Core.MethodMatch)
         push!(results, instance)
     end
     return results
@@ -1497,7 +1500,7 @@ function may_invoke_generator(method::Method, @nospecialize(atype), sparams::Sim
     gen_mthds isa Vector || return false
     length(gen_mthds) == 1 || return false
 
-    generator_method = first(gen_mthds).method
+    generator_method = (first(gen_mthds)::Core.MethodMatch).method
     nsparams = length(sparams)
     isdefined(generator_method, :source) || return false
     code = generator_method.source
@@ -2459,7 +2462,7 @@ function isambiguous(m1::Method, m2::Method; ambiguous_bottom::Bool=false)
         min = Ref{UInt}(typemin(UInt))
         max = Ref{UInt}(typemax(UInt))
         has_ambig = Ref{Int32}(0)
-        ms = _methods_by_ftype(ti, nothing, -1, world, true, min, max, has_ambig)::Vector
+        ms = collect(Core.MethodMatch, _methods_by_ftype(ti, nothing, -1, world, true, min, max, has_ambig)::Vector)
         has_ambig[] == 0 && return false
         if !ambiguous_bottom
             filter!(ms) do m::Core.MethodMatch
@@ -2472,7 +2475,6 @@ function isambiguous(m1::Method, m2::Method; ambiguous_bottom::Bool=false)
         # report the other ambiguous pair)
         have_m1 = have_m2 = false
         for match in ms
-            match = match::Core.MethodMatch
             m = match.method
             m === m1 && (have_m1 = true)
             m === m2 && (have_m2 = true)
diff --git a/test/reflection.jl b/test/reflection.jl
index c4e941fa7bb1a2..f2444a1cb8658e 100644
--- a/test/reflection.jl
+++ b/test/reflection.jl
@@ -1193,3 +1193,6 @@ end
 @test Base.isexported(Mod52812, :b)
 @test Base.ispublic(Mod52812, :a)
 @test Base.ispublic(Mod52812, :b)
+
+@test Base.infer_return_type(code_lowered, (Any,)) == Vector{Core.CodeInfo}
+@test Base.infer_return_type(code_lowered, (Any,Any)) == Vector{Core.CodeInfo}