From 06a552642ef4f927553f8e3bb01ffe0a7bfdf1a4 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Wed, 25 Nov 2020 12:52:33 -0600
Subject: [PATCH] Add `module_roots`

`module_roots` makes it easier to discover precompile-worthy
MethodInstances. The promise of `@snoopi_deep` is that it is not
restricted to just the entrance calls to inference; consequently,
we can find the most-expensive-to-infer calls in a specific module and
precompile those calls.
---
 src/SnoopCompile.jl       |  2 +-
 src/parcel_snoopi_deep.jl | 56 +++++++++++++++++++++++++++++++++++++++
 test/snoopi_deep.jl       | 22 +++++++++++++++
 3 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/src/SnoopCompile.jl b/src/SnoopCompile.jl
index 1b61fa448..5fd37cba9 100644
--- a/src/SnoopCompile.jl
+++ b/src/SnoopCompile.jl
@@ -4,7 +4,7 @@ using SnoopCompileCore
 export @snoopc
 isdefined(SnoopCompileCore, Symbol("@snoopi")) && export @snoopi
 if isdefined(SnoopCompileCore, Symbol("@snoopi_deep"))
-    export @snoopi_deep, flamegraph, flatten_times, accumulate_by_source
+    export @snoopi_deep, flamegraph, flatten_times, accumulate_by_source, module_roots
 end
 if isdefined(SnoopCompileCore, Symbol("@snoopr"))
     export @snoopr, uinvalidated, invalidation_trees, filtermod, findcaller, ascend
diff --git a/src/parcel_snoopi_deep.jl b/src/parcel_snoopi_deep.jl
index 39c2b2aea..03bff13c3 100644
--- a/src/parcel_snoopi_deep.jl
+++ b/src/parcel_snoopi_deep.jl
@@ -78,8 +78,16 @@ struct InclusiveTiming
     children::Vector{InclusiveTiming}
 end
 
+Base.show(io::IO, t::InclusiveTiming) = print(io, "InclusiveTiming: ", t.inclusive_time/10^9, " for ", t.mi_info.mi, " with ", length(t.children), " direct children")
+
 inclusive_time(t::InclusiveTiming) = t.inclusive_time
 
+"""
+    tinc = SnoopCompile.build_inclusive_times(t::Core.Compiler.Timings.Timing)
+
+Calculate times for inference for a node and all its children. `tinc.inclusive_time` records this time for
+node `tinc`; `tinc.children` gives you access to the children of this node.
+"""
 function build_inclusive_times(t::Timing)
     child_times = InclusiveTiming[
         build_inclusive_times(child)
@@ -89,6 +97,54 @@ function build_inclusive_times(t::Timing)
     return InclusiveTiming(t.mi_info, incl_time, t.start_time, child_times)
 end
 
+"""
+    roots = module_roots(mod::Module, t)
+
+Identify the "root" `MethodInstance`s in module `mod` on each branch of the children of `t`, where `t`
+comes from [`@snoopi_deep`](@ref). Each may or may not be the entrance point for a fresh inference tree,
+but on each branch it is the closest-to-the-root of any method in `mod`.
+
+# Example
+
+```julia
+julia> module M
+       f(x) = x > 0.5
+       end
+Main.M
+
+julia> t = @snoopi_deep filter(M.f, [0, 1]);
+
+julia> module_roots(Base, t)
+1-element Vector{SnoopCompile.InclusiveTiming}:
+ InclusiveTiming: 0.001296559 for MethodInstance for filter(::typeof(Main.M.f), ::Vector{Int64}) with 2 direct children
+
+ julia> module_roots(M, t)
+1-element Vector{SnoopCompile.InclusiveTiming}:
+ InclusiveTiming: 0.000274648 for MethodInstance for f(::Int64) with 1 direct children
+```
+
+Be aware that `InclusiveTiming` objects for the "same" node built from two separate calls to `module_roots` will not
+compare as equal. If you need that property, call [`SnoopCompile.build_inclusive_times`](@ref) to construct the tree
+`tinc` of `InclusiveTiming` nodes, and then select the roots from `tinc`.
+"""
+module_roots(mod::Module, t::Timing) = module_roots(mod, build_inclusive_times(t))
+module_roots(mod::Module, t::InclusiveTiming) = sort(module_roots!(InclusiveTiming[], mod, t); by=inclusive_time)
+
+function module_roots!(tmi, mod::Module, t::InclusiveTiming)
+    mi = t.mi_info.mi
+    m = mi.def
+    if isa(m, Method)
+        if m.module === mod
+            push!(tmi, t)
+            return tmi
+        end
+    end
+    foreach(t.children) do c
+        module_roots!(tmi, mod, c)
+    end
+    return tmi
+end
+
 """
     flamegraph(t::Core.Compiler.Timings.Timing; tmin_secs=0.0)
     flamegraph(t::SnoopCompile.InclusiveTiming; tmin_secs=0.0)
diff --git a/test/snoopi_deep.jl b/test/snoopi_deep.jl
index 52373e03e..908e4f900 100644
--- a/test/snoopi_deep.jl
+++ b/test/snoopi_deep.jl
@@ -59,6 +59,28 @@ using AbstractTrees  # For FlameGraphs tests
     @test length(timesmod) == 1
 end
 
+@testset "module_roots" begin
+    @eval module M
+        f(x) = x < 0.5
+    end
+    filter(M.f, [0, 1])  # warmup
+    @eval module M
+        f(x) = x < 0.5
+    end
+    timing = SnoopCompileCore.@snoopi_deep begin
+        filter(M.f, [0, 1])
+    end
+    # It's important to build the inclusive times only once, because otherwise `==` will fail in `mroot ∈ broot.children`
+    timingi = SnoopCompile.build_inclusive_times(timing)
+    mroot = only(module_roots(M, timingi))
+    broot = only(module_roots(Base, timingi))
+    @test mroot != broot
+    @test mroot ∈ broot.children
+    io = IOBuffer()
+    show(io, mroot)
+    @test endswith(String(take!(io)), "MethodInstance for f(::Int64) with 0 direct children")
+end
+
 @testset "flamegraph_export" begin
     @eval module M  # Take another timing
         i(x) = x+5