more tree work

Author: timholy

SnoopCompileCore/src/snoop_inference.jl

@@ -88,7 +88,6 @@ struct InferenceTimingNode
 end
 
 function timingtree(cis, backtraces)
-    remaining = Set(cis)
     root = InferenceTimingNode(Core.Compiler.Timings.ROOTmi.cache, nothing)
     # the cis are added in the order children-before-parents, we need to be able to reverse that
     # We index on MethodInstance rather than CodeInstance, because constprop can result in a distinct
@@ -97,52 +96,56 @@ function timingtree(cis, backtraces)
     backedges = [Int[] for _ in eachindex(cis)]
     for (i, ci) in pairs(cis)
         for e in ci.edges
+            e isa CodeInstance || continue
             eidx = get(miidx, methodinstance(e), nothing)
             if eidx !== nothing
                 push!(backedges[eidx], i)
             end
         end
     end
     # backtraces = Dict{MethodInstance,Any}(backtrace_log[i] => backtrace_log[i+1] for i in 1:2:length(backtrace_log))
-    addchildren!(root, remaining, cis, backedges, miidx, backtraces)
+    addchildren!(root, cis, backedges, miidx, backtraces)
     return root
 end
 
-function addchildren!(parent::InferenceTimingNode, remaining::Set{CodeInstance}, miidx)
+function addchildren!(parent::InferenceTimingNode, handled::Set{CodeInstance}, miidx)
     for ci in parent.ci.edges
+        ci isa CodeInstance || continue
         haskey(miidx, methodinstance(ci)) || continue
+        ci ∈ handled && continue
         child = InferenceTimingNode(ci, nothing, parent)
-        delete!(remaining, ci)
-        addchildren!(child, remaining, miidx)
+        push!(handled, ci)
+        addchildren!(child, handled, miidx)
     end
     return parent
 end
 
-function addchildren!(parent::InferenceTimingNode, remaining::Set{CodeInstance}, cis, backedges, miidx, backtraces)
+function addchildren!(parent::InferenceTimingNode, cis, backedges, miidx, backtraces)
+    handled = Set{CodeInstance}()
     for (i, ci) in pairs(cis)
-        ci ∈ remaining || continue
+        ci ∈ handled && continue
         # Follow the backedges to the root
         j = i
         be = ci
         while true
             found = false
             for k in backedges[j]
                 be = cis[k]
-                if be ∈ remaining
+                if be ∉ handled
                     j = k
                     found = true
                     break
                 end
             end
             found || break
         end
+        be ∈ handled && continue
         # bt1, bt2 = get(backtraces, Core.Compiler.get_ci_mi(be), (nothing, nothing))
         # child = InferenceTimingNode(be, make_stacktrace(bt1, bt2), parent)
         child = InferenceTimingNode(be, get(backtraces, be, nothing), parent)
-        delete!(remaining, be)
-        addchildren!(child, remaining, miidx)
+        push!(handled, be)
+        addchildren!(child, handled, miidx)
     end
-    # `remaining` may not be empty if there is constprop
     return parent
 end