Add iterator for method specializations (JuliaLang#49116)

There have been longstanding reasons to want an API for extracting
all extant specializations of a method, but this is even more
true after JuliaLang#49071.

Co-authored-by: Shuhei Kadowaki <40514306+aviatesk@users.noreply.github.com>
Co-authored-by: Jameson Nash <vtjnash@gmail.com>

base/reflection.jl

@@ -1110,6 +1110,34 @@ 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, nothing)
+ return iterate(specs, 0)
+end
+iterate(specs::MethodSpecializations, ::Nothing) = nothing
+function iterate(specs::MethodSpecializations, i::Int)
+ s = specs.specializations::Core.SimpleVector
+ 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
+length(specs::MethodSpecializations) = count(Returns(true), specs)
 
 function length(mt::Core.MethodTable)
  n = 0

doc/src/manual/performance-tips.md

@@ -584,7 +584,7 @@ h_vararg(x::Vararg{Any, N}) where {N} = tuple(x...)
 Note that [`@code_typed`](@ref) and friends will always show you specialized code, even if Julia
 would not normally specialize that method call. You need to check the
 [method internals](@ref ast-lowered-method) if you want to see whether specializations are generated
-when argument types are changed, i.e., if `(@which f(...)).specializations` contains specializations
+when argument types are changed, i.e., if `Base.specializations(@which f(...))` contains specializations
 for the argument in question.
 
 ## Break functions into multiple definitions

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)
 

test/reflection.jl

@@ -1039,3 +1039,10 @@ ambig_effects_test(a, b) = 1
 end
 
 @test Base._methods_by_ftype(Tuple{}, -1, Base.get_world_counter()) == Any[]
+
+@testset "specializations" begin
+ f(x) = 1
+ f(1)
+ f("hello")
+ @test length(Base.specializations(only(methods(f)))) == 2
+end

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