Redesign visit (#3)

visit: Bool recursion control, print off by default, etc

This makes many changes to `visit`:
- it requires `operation` to return a Bool to control descent
- it extends the "visit-once" logic to all types
- it extends printing to all objects but turns it off by default
- it greatly improves performance with nospecialize annotations
- it separates user-interface from internals (visit vs _visit)
- it supports more types (MethodTable, CodeInstance)

Closes #2

Also redesign visit_backedges to exploit new visit.
This orthogonalizes the two functions.

Author: timholy

Project.toml

@@ -7,8 +7,8 @@ version = "0.1.0"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 
 [compat]
-julia = "1"
 AbstractTrees = "0.3"
+julia = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

docs/src/index.md

@@ -3,7 +3,7 @@
 This package facilitates introspection of Julia's internals, with a particular focus on its MethodInstances and their backedges.
 
 !!! warning
-    Julia's internals are not subject to the same stability guarantee that the rest of the language enjoys.
+    Julia's internals are not subject to the same interface compatibility guarantee that the rest of the language enjoys.
 
 ## Demonstrations
 
@@ -16,10 +16,12 @@ julia> using MethodAnalysis
 
 julia> mods = Module[];
 
-julia> visit(Base; print=false) do obj
+julia> visit(Base) do obj
            if isa(obj, Module)
                push!(mods, obj)
+               return true     # descend into submodules
            end
+           false   # but don't descend into anything else (MethodTables, etc.)
        end
 
 julia> Base.FastMath ∈ mods
@@ -34,15 +36,10 @@ true
 julia> meths = []
 Any[]
 
-julia> visit(Core.Compiler) do item    # without print=false it will display the modules it visits
+julia> visit(Core.Compiler) do item
            isa(item, Method) && push!(meths, item)
+           true   # walk through everything
        end
-Module Core.Compiler
-Module Core.Compiler.CoreDocs
-Module Core.Compiler.Iterators
-Module Core.Compiler.Order
-Module Core.Compiler.Sort
-Module Core.Compiler.Sort.Float
 
 julia> first(methods(Core.Compiler.typeinf_ext)) ∈ meths
 true
@@ -72,6 +69,7 @@ julia> mis = Core.MethodInstance[];
 
 julia> visit(findfirst) do item
            isa(item, Core.MethodInstance) && length(Base.unwrap_unionall(item.specTypes).parameters) == 2 && push!(mis, item)
+           true
        end
 
 julia> mis

src/MethodAnalysis.jl

@@ -2,6 +2,7 @@ module MethodAnalysis
 
 using AbstractTrees
 
+using Base: Callable, IdSet
 using Core: MethodInstance, SimpleVector, MethodTable
 
 export visit, call_type, instance, worlds
@@ -77,7 +78,9 @@ function instance(f, types)
             if mi.specTypes === tt
                 inst = mi
             end
+            return false
         end
+        true
     end
     return inst
 end

src/backedges.jl

@@ -18,14 +18,11 @@ their backedges. The result includes both MethodTable and MethodInstance backedg
 """
 function with_all_backedges(iter)
     backedges = Set{MethodInstance}()
-    visited = Set{Union{MethodInstance,MethodTable}}()
+    visited = IdSet{MethodInstance}()
     for item in iter
         visit_backedges(item, visited) do edge
             if isa(edge, MethodInstance)
                 push!(backedges, edge)
-            else
-                sig, mi = edge
-                push!(backedges, mi)
             end
             true
         end
@@ -57,26 +54,32 @@ Collect all backedges for a function `f` as pairs `instance=>caller` or `sig=>ca
 The latter occur for MethodTable backedges.
 If `skip` is `true`, any `caller` listed in a MethodTable backedge is omitted from the instance backedges. 
 """
-function direct_backedges(f::Function; skip::Bool=true)
+function direct_backedges(f::Union{Method,Callable}; skip::Bool=true)
     bes = []
     _skip = Set{MethodInstance}()
-    mths = methods(f).ms
-    callee = nothing
-    visit_backedges(f) do item
-        if isa(item, Pair)
-            push!(bes, item)
-            push!(_skip, item.second)
+    visit(f) do item
+        if isa(item, MethodTable)
+            mt = item::MethodTable
+            if isdefined(mt, :backedges)
+                sigmis = mt.backedges::Vector{Any}
+                for i = 1:2:length(sigmis)
+                    sig, mi = sigmis[i], sigmis[i+1]
+                    push!(bes, Pair{Any,MethodInstance}(sig, mi))
+                    push!(_skip, mi)
+                end
+            end
             return false
-        else
-            mi = item::MethodInstance
-            if mi.def ∈ mths
-                callee = mi
-                return true
-            else
-                (!skip || mi ∉ _skip) && push!(bes, callee=>mi)
-                return false
+        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
             end
+            return false
         end
+        true
     end
     return bes
 end