Skip to content

Commit

Permalink
Add visit_withmodule (#32)
Browse files Browse the repository at this point in the history
This allows the operation to know the module in which
the current item was found. Note that this is distinct
from module ownership:

module A f(x) = 1 end
module B
    using ..A
    A.f(x::Int) = 2
end

In this case both methods are found while traversing A,
even though `which(A.f, (Int,)).module == B`.
  • Loading branch information
timholy authored Sep 19, 2022
1 parent 81f05f1 commit f9098fe
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 3 deletions.
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "MethodAnalysis"
uuid = "85b6ec6f-f7df-4429-9514-a64bcd9ee824"
authors = ["Tim Holy <tim.holy@gmail.com>"]
version = "0.4.7"
version = "0.4.8"

[deps]
AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
Expand Down
3 changes: 2 additions & 1 deletion src/MethodAnalysis.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
37 changes: 36 additions & 1 deletion src/visit.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 39 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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, :(
Expand Down

2 comments on commit f9098fe

@timholy
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/68539

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.4.8 -m "<description of version>" f9098fe5b51a8c52ee86e6374e5e112d9c521918
git push origin v0.4.8

Please sign in to comment.