Add iterator for method specializations

There have been longstanding reasons to want an API for extracting all extant specializations of a method, but this is even more true after #49071.
JuliaLang · Mar 20, 2023 · 63e194f · 63e194f
1 parent 3e96c76
commit 63e194f
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 29 deletions.
diff --git a/base/reflection.jl b/base/reflection.jl
@@ -1107,6 +1107,33 @@ function visit(f, d::Core.TypeMapEntry)
     end
     nothing
 end
+struct MethodSpecializations
+    specializations::Union{Nothing, Core.MethodInstance, Core.SimpleVector}
+end
+"""
+    specializations(m::Method) → itr
+
+Return an iterator `itr` of all compiler-generated specializations of `m`.
+"""
+specializations(m::Method) = MethodSpecializations(isdefined(m, :specializations) ? m.specializations : nothing)
+function iterate(specs::MethodSpecializations)
+    s = specs.specializations
+    s === nothing && return nothing
+    isa(s, Core.MethodInstance) && return (s, false)
+    return iterate(specs, 0)
+end
+iterate(specs::MethodSpecializations, ::Bool) = nothing
+function iterate(specs::MethodSpecializations, i::Int)
+    s = specs.specializations
+    n = length(s)
+    i >= n && return nothing
+    item = nothing
+    while i < n && item === nothing
+        item = s[i+=1]
+    end
+    item === nothing && return nothing
+    return (item, i)
+end
 
 function length(mt::Core.MethodTable)
     n = 0

diff --git a/test/precompile.jl b/test/precompile.jl
@@ -665,10 +665,8 @@ precompile_test_harness("code caching") do dir
     # size(::Vector) has an inferred specialization for Vector{X}
     msize = which(size, (Vector{<:Any},))
     hasspec = false
-    msizespecs = msize.specializations::Core.SimpleVector
-    for i = 1:length(msizespecs)
-        mi = msizespecs[i]
-        if isa(mi, Core.MethodInstance) && mi.specTypes == Tuple{typeof(size),Vector{Cacheb8321416e8a3e2f1.X}}
+    for mi in Base.specializations(msize)
+        if mi.specTypes == Tuple{typeof(size),Vector{Cacheb8321416e8a3e2f1.X}}
             if isdefined(mi, :cache) && isa(mi.cache, Core.CodeInstance) && mi.cache.max_world == typemax(UInt) && mi.cache.inferred !== nothing
                 hasspec = true
                 break
@@ -786,9 +784,7 @@ precompile_test_harness("code caching") do dir
     MB = getfield(@__MODULE__, RootB)
     M = getfield(MA, RootModule)
     m = which(M.f, (Any,))
-    mspecs = m.specializations
-    mspecs isa Core.SimpleVector || (mspecs = Core.svec(mspecs))
-    for mi in mspecs
+    for mi in Base.specializations(m)
         mi === nothing && continue
         mi = mi::Core.MethodInstance
         if mi.specTypes.parameters[2] === Int8
@@ -925,8 +921,7 @@ precompile_test_harness("code caching") do dir
     # Reporting test (ensure SnoopCompile works)
     @test all(i -> isassigned(invalidations, i), eachindex(invalidations))
     m = only(methods(MB.call_nbits))
-    for mi in m.specializations::Core.SimpleVector
-        mi === nothing && continue
+    for mi in Base.specializations(m)
         hv = hasvalid(mi, world)
         @test mi.specTypes.parameters[end] === Integer ? !hv : hv
     end
@@ -1088,13 +1083,13 @@ precompile_test_harness("invoke") do dir
     end
 
     m = get_method_for_type(M.h, Real)
-    @test m.specializations === Core.svec()
+    @test isempty(Base.specializations(m))
     m = get_method_for_type(M.hnc, Real)
-    @test m.specializations === Core.svec()
+    @test isempty(Base.specializations(m))
     m = only(methods(M.callq))
-    @test m.specializations === Core.svec() || nvalid(m.specializations::Core.MethodInstance) == 0
+    @test isempty(Base.specializations(m)) || nvalid(m.specializations::Core.MethodInstance) == 0
     m = only(methods(M.callqnc))
-    @test m.specializations === Core.svec() || nvalid(m.specializations::Core.MethodInstance) == 0
+    @test isempty(Base.specializations(m)) || nvalid(m.specializations::Core.MethodInstance) == 0
     m = only(methods(M.callqi))
     @test (m.specializations::Core.MethodInstance).specTypes == Tuple{typeof(M.callqi), Int}
     m = only(methods(M.callqnci))
@@ -1112,7 +1107,7 @@ precompile_test_harness("invoke") do dir
     m_any, m_int = sort(collect(methods(invokeme)); by=m->(m.file,m.line))
     @test precompile(invokeme, (Int,), m_any)
     @test (m_any.specializations::Core.MethodInstance).specTypes === Tuple{typeof(invokeme), Int}
-    @test m_int.specializations === Core.svec()
+    @test isempty(Base.specializations(m_int))
 end
 
 # test --compiled-modules=no command line option
@@ -1581,7 +1576,7 @@ precompile_test_harness("issue #46296") do load_path
         """
         module CodeInstancePrecompile
 
-        mi = first(methods(identity)).specializations[1]
+        mi = first(Base.specializations(first(methods(identity))))
         ci = Core.CodeInstance(mi, Any, nothing, nothing, zero(Int32), typemin(UInt),
                                typemax(UInt), zero(UInt32), zero(UInt32), nothing, 0x00)
 

diff --git a/test/worlds.jl b/test/worlds.jl
@@ -233,21 +233,10 @@ function method_instance(f, types=Base.default_tt(f))
     m = which(f, types)
     inst = nothing
     tt = Base.signature_type(f, types)
-    specs = m.specializations
-    if isa(specs, Core.SimpleVector)
-        for i = 1:length(specs)
-            mi = specs[i]
-            if mi isa Core.MethodInstance
-                if mi.specTypes <: tt && tt <: mi.specTypes
-                    inst = mi
-                    break
-                end
-            end
-        end
-    else
-        mi = specs::Core.MethodInstance
-        if mi.specTypes === tt
+    for mi in Base.specializations(m)
+        if mi.specTypes <: tt && tt <: mi.specTypes
             inst = mi
+            break
         end
     end
     return inst