Add visit_withmodule

Project.toml

@@ -1,7 +1,7 @@
 name = "MethodAnalysis"
 uuid = "85b6ec6f-f7df-4429-9514-a64bcd9ee824"
 authors = ["Tim Holy <[email protected]>"]
-version = "0.4.7"
+version = "0.4.8"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

src/MethodAnalysis.jl

@@ -6,7 +6,8 @@ using Base: Callable, IdSet
 using Core: MethodInstance, CodeInfo, SimpleVector, MethodTable
 using Base.Meta: isexpr
 
-export visit, call_type, methodinstance, methodinstances, worlds  # findcallers is exported from its own file
+export visit, visit_withmodule
+export call_type, methodinstance, methodinstances, worlds  # findcallers is exported from its own file
 export visit_backedges, all_backedges, with_all_backedges, terminal_backedges, direct_backedges
 export child_modules, methodinstances_owned_by
 export hasbox

src/visit.jl

@@ -68,17 +68,52 @@ function visit(@nospecialize(operation); print::Bool=false)
     return nothing
 end
 
+struct ModuleWrapper
+    operation
+    parent::Union{Module,Nothing}
+end
+(w::ModuleWrapper)(@nospecialize(x)) = w.operation(x, w.parent)
+
+rewrap_operation(@nospecialize(operation), ::Module) = operation
+rewrap_operation(w::ModuleWrapper, mod::Module) = ModuleWrapper(w.operation, mod)
+
+"""
+    visit_withmodule(operation; print::Bool=false)
+    visit_withmodule(operation, obj, mod; print::Bool=false)
+
+Similar to [`visit`](@ref), except that `operation` should have signature `operation(x, mod)` where `mod` is either:
+
+- the module in which `x` was found, or
+- `nothing` if `x` is itself a top-level module.
+
+If you're visiting underneath a specific object `obj`, you must supply `mod`, the module (or `nothing`) in which
+`obj` would be found.
+"""
+function visit_withmodule(@nospecialize(operation); print::Bool=false)
+    visited = IdSet{Any}()
+    wrapped = ModuleWrapper(operation, nothing)
+    for mod in Base.loaded_modules_array()
+        _visit(wrapped, mod, visited, print)
+    end
+    return nothing
+end
+
+function visit_withmodule(@nospecialize(operation), @nospecialize(obj), mod::Union{Module,Nothing}; print::Bool=false)
+    return _visit(ModuleWrapper(operation, mod), obj, IdSet{Any}(), print)
+end
+
 # These are non-keyword functions due to https://github.com/JuliaLang/julia/issues/34516
 
 function _visit(@nospecialize(operation), mod::Module, visited::IdSet{Any}, print::Bool)
     mod ∈ visited && return nothing
     push!(visited, mod)
     print && println("Module ", mod)
     if operation(mod)
+        newop = rewrap_operation(operation, mod)
         for nm in names(mod; all=true)
             if isdefined(mod, nm)
                 obj = getfield(mod, nm)
-                _visit(operation, obj, visited, print)
+                _visit(newop, obj, visited, print)
             end
         end
     end

test/runtests.jl

@@ -21,6 +21,14 @@ module Outer
     f2(x, y::Number) = x + y
 end
 
+module One
+    f(x) = 1
+end
+module Two
+    using ..One
+    One.f(x::Int) = x + 2
+end
+
 
 @testset "visit" begin
     @test Outer.Inner.g("hi") == 0
@@ -72,6 +80,37 @@ end
     end
 end
 
+@testset "visit_withmodule" begin
+    found = Dict{Module,Set{Method}}()
+    visit_withmodule(Main, nothing) do item, mod
+        if item === Main
+            @test mod === nothing
+            return true
+        end
+        if mod == Main && item isa Module
+            return item ∈ (Outer, Outer.Inner, Outer.OtherInner, One, Two)
+        end
+        item isa Method || return true
+        push!(get!(Set{Any}, found, mod), item)
+        return false
+    end
+    s = found[One]
+    @test methods(One.f) ⊆ s
+    @test any(m -> m.module == One, s)
+    @test any(m -> m.module == Two, s)
+    @test methods(Outer.f2) ⊆ found[Outer]
+    @test methods(Outer.Inner.h) ⊆ found[Outer.Inner]
+
+    found = Dict{Union{Module,Nothing}, Set{Module}}()
+    visit_withmodule() do item, mod
+        item isa Module || return false
+        push!(get!(Set{Module}, found, mod), item)
+        return true
+    end
+    @test Main ∈ union(found[nothing], found[Core])
+    @test One ∈ found[Main]
+end
+
 @testset "child_modules" begin
     m = Module()
     Base.eval(m, :(