Support Core.Compiler.BackedgeIterator

src/MethodAnalysis.jl

@@ -18,6 +18,19 @@ include("backedges.jl")
 if !hasmethod(==, Tuple{Core.PhiNode,Core.PhiNode})
     Base.:(==)(stmt1::Core.PhiNode, stmt2::Core.PhiNode) = stmt1.edges == stmt2.edges && stmt1.values == stmt2.values
 end
+if isdefined(Core.Compiler, :BackedgeIterator)
+    getmi(mi::MethodInstance) = mi
+    getmi(pr::Pair{Any,MethodInstance}) = pr.second
+    getmi(pr::Core.Compiler.BackedgePair) = pr.caller
+
+    stdbe(::Nothing, caller::MethodInstance) = caller
+    stdbe(@nospecialize(invokesig), caller::MethodInstance) = Pair{Any,MethodInstance}(invokesig, caller)
+    stdbe(pr::Core.Compiler.BackedgePair) = stdbe(pr.sig, pr.caller)
+
+    if !hasmethod(iterate, Tuple{Core.Compiler.BackedgeIterator})
+        Base.iterate(iter::Core.Compiler.BackedgeIterator, state...) = Core.Compiler.iterate(iter, state...)
+    end
+end
 
 """
     call_type(tt)

src/backedges.jl

@@ -72,17 +72,31 @@ function direct_backedges(f::Union{Method,Callable}; skip::Bool=true)
         elseif isa(item, MethodInstance)
             callee = item::MethodInstance
             if isdefined(callee, :backedges)
-                for caller in callee.backedges
-                    skip && caller ∈ _skip && continue
-                    push!(bes, callee=>caller)
-                end
+                push_unskipped_backedges!(bes, callee, skip, _skip)
             end
             return false
         end
         true
     end
     return bes
 end
+if isdefined(Core.Compiler, :BackedgeIterator)
+    function push_unskipped_backedges!(bes, callee, skip, _skip)
+        for caller in Core.Compiler.BackedgeIterator(callee.backedges)
+            skip && getmi(caller) ∈ _skip && continue
+            push!(bes, stdbe(caller.sig, callee)=>caller.caller)
+        end
+        return bes
+    end
+else
+    function push_unskipped_backedges!(bes, callee, skip, _skip)
+        for caller in callee.backedges
+            skip && caller ∈ _skip && continue
+            push!(bes, callee=>caller)
+        end
+        return bes
+    end
+end
 
 """
     direct_backedges(mi::MethodInstance)

src/visit.jl

@@ -229,8 +229,7 @@ _visit(@nospecialize(operation), @nospecialize(x), visited::IdSet{Any}, print::B
 Visit the backedges of `obj` and apply `operation` to each.
 `operation` may need to be able to handle two call forms, `operation(mi)` and
 `operation(sig=>mi)`, where `mi` is a `MethodInstance` and `sig` is a `Tuple`-type.
-The latter arises from `MethodTable` backedges and can be ignored if `obj` does not
-contain `MethodTable`s.
+The latter arises from either `invoke` calls or `MethodTable` backedges.
 
 `operation(edge)` should return `true` if the backedges of `edge` should in turn be visited,
 `false` otherwise.
@@ -267,15 +266,28 @@ function visit_backedges(@nospecialize(operation), obj, visited::IdSet{MethodIns
     visit(opwrapper, obj)
 end
 
-function _visit_backedges(@nospecialize(operation), mi::MethodInstance, visited)
-    mi ∈ visited && return nothing
-    push!(visited, mi)
-    if operation(mi) && isdefined(mi, :backedges)
-        for edge in mi.backedges
-            _visit_backedges(operation, edge, visited)
+if isdefined(Core.Compiler, :BackedgeIterator)
+    function _visit_backedges(@nospecialize(operation), callee::Union{MethodInstance,Pair{Any,MethodInstance}}, visited)
+        getmi(callee) ∈ visited && return nothing
+        push!(visited, callee)
+        if operation(callee) && isdefined(getmi(callee), :backedges)
+            for be in Core.Compiler.BackedgeIterator(getmi(callee).backedges)
+                _visit_backedges(operation, stdbe(be), visited)
+            end
         end
+        return nothing
+    end
+else
+    function _visit_backedges(@nospecialize(operation), mi::MethodInstance, visited)
+        mi ∈ visited && return nothing
+        push!(visited, mi)
+        if operation(mi) && isdefined(mi, :backedges)
+            for edge in mi.backedges
+                _visit_backedges(operation, edge, visited)
+            end
+        end
+        return nothing
     end
-    return nothing
 end
 
 function _visit_backedges(@nospecialize(operation), misig::Pair{Any,MethodInstance}, visited)

test/runtests.jl

@@ -21,6 +21,10 @@ module Outer
     f2(x, y::Number) = x + y
 
     fkw(x; y=0) = 2x + y
+
+    invk(::Int) = 1
+    invk(::Integer) = 2
+    geninvk(x) = invoke(invk, Tuple{Integer}, x)
 end
 
 module One
@@ -266,6 +270,16 @@ end
     callnocallers(x) = nocallers(x)
     callnocallers(3)
     @test !isempty(direct_backedges(mi))
+
+    if isdefined(Core.Compiler, :BackedgeIterator)
+        @test Outer.geninvk(3) == 2
+        m = which(Outer.invk, (Integer,))
+        bes = direct_backedges(m)
+        be = only(bes)
+        @test be.first.second.specTypes.parameters[2] === Int
+        @test be.first.first === Tuple{typeof(Main.Outer.invk), Integer}
+        @test be.second == methodinstance(Outer.geninvk, (Int,))
+    end
 end
 
 @testset "call_type" begin