Add child_modules and module-owned MIs (#30)

`methodinstances_owned_by` captures the MethodInstances that can be serialized in current versions of Julia.

Author: timholy

Project.toml

@@ -11,9 +11,10 @@ AbstractTrees = "0.3"
 julia = "1"
 
 [extras]
+ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Logging", "Pkg", "Test"]
+test = ["ImageCore", "Logging", "Pkg", "Test"]

docs/src/index.md

@@ -28,6 +28,8 @@ julia> Base.FastMath ∈ mods
 true
 ```
 
+You can do this more easily with the convenience utility [`child_modules`](@ref).
+
 ### Collecting all Methods in Core.Compiler
 
 `visit` also descends into functions, methods, and MethodInstances:
@@ -187,8 +189,11 @@ with_all_backedges
 ```@docs
 methodinstance
 methodinstances
+methodinstances_owned_by
+child_modules
 call_type
 findcallers
+hasbox
 worlds
 ```
 

src/MethodAnalysis.jl

@@ -8,6 +8,7 @@ using Base.Meta: isexpr
 
 export visit, 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
 
 include("visit.jl")
@@ -150,7 +151,8 @@ julia> methodinstances(m)
 ```
 
 Note the method `m` was broader than the signature we queried with, and the returned `MethodInstance`s reflect that breadth.
-See [`methodinstances`](@ref) for a more restrictive subset.
+See [`methodinstances`](@ref) for a more restrictive subset, and [`methodinstances_owned_by`](@ref) for collecting
+MethodInstances owned by specific modules.
 """
 function methodinstances(top=())
     if isa(top, Module) || isa(top, Function) || isa(top, Type) || isa(top, Method) || isa(top, Base.MethodList)
@@ -190,6 +192,96 @@ function methodinstances(@nospecialize(types::Type))
     return methodinstances(f, types)
 end
 
+"""
+    mods = child_modules(mod::Module; external::Bool=false)
+
+Return a list that includes `mod` and all sub-modules of `mod`.
+By default, modules loaded from other sources (e.g., packages or those
+defined by Julia itself) are excluded, even if exported (or `@reexport`ed,
+see https://github.com/simonster/Reexport.jl), unless you set `external=true`.
+
+# Examples
+
+```jldoctest
+julia> module Outer
+       module Inner
+       export Base
+       end
+       end
+Main.Outer
+
+julia> child_modules(Outer)
+2-element Vector{Module}:
+ Main.Outer
+ Main.Outer.Inner
+
+julia> child_modules(Outer.Inner)
+1-element Vector{Module}:
+ Main.Outer.Inner
+```
+
+# Extended help
+
+In the example above, because of the `export Base`, the following `visit`-based implementation would
+also collect `Base` and all of its sub-modules:
+
+```jldoctest
+julia> mods = Module[]
+Module[]
+
+julia> visit(Outer) do item
+           if item isa Module
+               push!(mods, item)
+               return true
+           end
+           return false
+       end
+
+julia> Base ∈ mods
+true
+
+julia> length(mods) > 20
+true
+```
+"""
+function child_modules(mod::Module; external::Bool=false)
+    function rootmodule(m::Module)
+        m == mod && return m   # anything under `mod` has a root of `mod`
+        pm = parentmodule(m)
+        m == pm && return m
+        return rootmodule(pm)
+    end
+    mods = Module[]
+    visit(mod) do item
+        if item isa Module && (external || rootmodule(item) == mod)
+            push!(mods, item)
+            return true
+        end
+        return false  # don't recurse into Methods, MethodTables, MethodInstances, etc.
+    end
+    return mods
+end
+
+"""
+    mis = methodinstances_owned_by(mod::Module; include_child_modules::Bool=true, kwargs...)
+
+Return a list of `MethodInstance`s that are owned by `mod`. If `include_child_modules` is `true`,
+this includes sub-modules of `mod`, in which case `kwargs` are passed to [`child_modules`](@ref).
+
+The primary difference between `methodinstances(mod)` and `methodinstances_owned_by(mod)` is that
+the latter excludes `MethodInstances` that belong to re-exported dependent packages.
+"""
+function methodinstances_owned_by(mod::Module; include_child_modules::Bool=true, kwargs...)
+    mods = include_child_modules ? child_modules(mod; kwargs...) : [mod]
+    # get all MethodInstances owned by one of the modules in `mods`
+    # these are the only MethodInstances that can be precompiled in current versions of Julia
+    return filter(methodinstances(mod)) do mi
+        m = mi.def
+        m isa Method && return m.module ∈ mods
+        return m ∈ mods
+    end
+end
+
 if isdefined(Base, :code_typed_by_type)
     function hasbox(mi::MethodInstance)
         try

test/runtests.jl

@@ -1,6 +1,7 @@
 using MethodAnalysis
 using Test
 using Logging
+using ImageCore
 using Pkg
 
 module Outer
@@ -71,6 +72,30 @@ end
     end
 end
 
+@testset "child_modules" begin
+    m = Module()
+    Base.eval(m, :(
+        module Inner
+        export Base
+        end))
+    mmods = child_modules(m)
+    @test m ∈ mmods
+    @test m.Inner ∈ mmods
+    @test length(mmods) == 2
+    imods = child_modules(m.Inner)
+    @test m.Inner ∈ mmods
+    @test length(imods) == 1
+
+    # Base is interesting because it's not its own parent
+    bmods = child_modules(Base)
+    @test Base ∈ bmods
+    @test Base.Sort ∈ bmods
+    @test Main ∉ bmods
+    smods = child_modules(Base.Sort)
+    @test Base ∉ smods
+    @test Base.Sort ∈ smods
+end
+
 @testset "methodinstance(s)" begin
     sum(1:3)
     mi = methodinstance(sum, (UnitRange{Int},))
@@ -102,6 +127,32 @@ end
     end
 end
 
+@testset "AbstractTrees integration" begin
+    mi = methodinstance(findfirst, (BitVector,))
+    io = IOBuffer()
+    MethodAnalysis.AbstractTrees.print_tree(io, mi)
+    str = String(take!(io))
+    @test occursin("├─", str)
+end
+
+@testset "methodinstances_owned_by" begin
+    function owned_by(mi::Core.MethodInstance, mod::Module)
+        m = mi.def
+        m isa Method && return m.module == mod
+        return m == mod
+    end
+    # ImageCore is a package that does a lot of `@reexport`ing
+    mis = methodinstances(ImageCore)
+    @test  any(mi -> owned_by(mi, ImageCore), mis)
+    @test  any(mi -> owned_by(mi, ImageCore.ColorTypes), mis)  # ColorTypes is a dependency of ImageCore
+    mis = methodinstances_owned_by(ImageCore)
+    @test  any(mi -> owned_by(mi, ImageCore), mis)
+    @test !any(mi -> owned_by(mi, ImageCore.ColorTypes), mis)
+    mis = methodinstances_owned_by(ImageCore; external=true)
+    @test  any(mi -> owned_by(mi, ImageCore), mis)
+    @test  any(mi -> owned_by(mi, ImageCore.ColorTypes), mis)
+end
+
 @testset "Backedges" begin
     mi = methodinstance(Outer.Inner.h, (Int,))
     @test length(all_backedges(mi)) == 2