Improve the API and add specializations demo

Author: timholy

demos/specialization.jl

@@ -0,0 +1,71 @@
+using MethodAnalysis
+using OrderedCollections
+using CodeTracking
+
+const fully_specialized = OrderedSet{Core.MethodInstance}()
+const nonspecialized = OrderedSet{Core.MethodInstance}()
+
+function categorizemi(mi::Core.MethodInstance)
+    bin = all(isconcretetype, Base.unwrap_unionall(mi.specTypes).parameters) ?
+        fully_specialized : nonspecialized
+    push!(bin, mi)
+    return nothing
+end
+
+categorizemi(x) = nothing
+
+visit(categorizemi)
+# We're especially interested in the nonspecialized ones.
+# Start with the ones in Base
+ubase = collect(filter(x->x.def.module === Base, nonspecialized))
+# Find specializations that have a TypeVar parameter
+function hastv(typ)
+    isa(typ, UnionAll) && return true
+    if isa(typ, DataType)
+        for p in typ.parameters
+            hastv(p) && return true
+        end
+    end
+    return false
+end
+ubasetv = filter(mi->hastv(mi.specTypes), ubase)
+println("There are ", length(ubasetv), " MethodInstances that have a TypeVar in specTypes")
+
+# Let's analyze a specific case
+m = which(similar, (Vector,))
+mitv = Ref{Any}(nothing)
+visit(m.specializations) do mi
+    mitv[] isa Core.MethodInstance && return nothing
+    hastv(mi.specTypes) && (mitv[] = mi)
+    return nothing
+end
+mi = mitv[]
+@assert(mi isa Core.MethodInstance)
+println("\n## Performing an analysis of ", mi, '\n')
+# Find the last callers with TypeVar specializations
+callers = Core.MethodInstance[]
+visit_backedges(mi) do caller
+    hastv(caller.specTypes) || return false
+    if !isdefined(caller, :backedges)
+        push!(callers, caller)
+        return false
+    end
+    foundone = false
+    for edge in caller.backedges
+        if !hastv(edge.specTypes)
+            push!(callers, caller)
+            foundone = true
+            break
+        end
+    end
+    return !foundone
+end
+# Let's look at the code of these callers
+for caller in callers
+    for be in caller.backedges
+        f, t = call_type(be.specTypes)
+        println("caller: ", be.def, " with argtypes ", t)
+        code_warntype(f, t)
+        println()
+    end
+end

src/MethodAnalysis.jl

@@ -2,9 +2,21 @@ module MethodAnalysis
 
 using Core: MethodInstance, SimpleVector
 
-export visit, visit_backedges, all_backedges, terminal_backedges
+export visit, visit_backedges, all_backedges, terminal_backedges, call_type
 
 include("visit.jl")
 include("backedges.jl")
 
+"""
+    call_type(tt)
+
+Split a signature type like `Tuple{typeof(f),ArgTypes...}` back out to `(f, Tuple{ArgTypes...})`
+"""
+function call_type(tt)
+    ft = tt.parameters[1]
+    argt = Tuple{tt.parameters[2:end]...}
+    name = ft.name
+    return (getfield(name.module, Symbol(String(name.name)[2:end])), argt)
+end
+
 end # module

src/backedges.jl

@@ -1,6 +1,6 @@
 function all_backedges(mi)
     backedges = Set{MethodInstance}()
-    visit_backedges(x->push!(backedges, x), mi)
+    visit_backedges(x->(push!(backedges, x); true), mi)
     delete!(backedges, mi)
     return collect(backedges)
 end
@@ -11,6 +11,7 @@ function terminal_backedges(mi)
         if !isdefined(x, :backedges) || isempty(x.backedges)
             push!(backedges, x)
         end
+        true
     end
     delete!(backedges, mi)
     return collect(backedges)

src/visit.jl

@@ -1,3 +1,9 @@
+"""
+    visit(operation)
+
+Scan all loaded modules with `operation`. `operation(x)` should handle `x::Module`, `x::Function`,
+`x::Method`, `x::MethodInstance`. Any return value from `operation` will be discarded.
+"""
 function visit(operation)
     visiting=Set{Module}()
     for mod in Base.loaded_modules_array()
@@ -73,16 +79,23 @@ end
 
 visit(operation, x) = nothing
 
+"""
+    visit_backedges(operation, mi::MethodInstance)
+
+Visit the backedges of `mi` and apply `operation`.
+`operation(edge::MethodInstance)` should return `true` if the backedges of `edge` should in turn be visited,
+`false` otherwise.
+"""
 visit_backedges(operation, mi::MethodInstance) =
     visit_backedges(operation, mi, Set{MethodInstance}())
 
 function visit_backedges(operation, mi, visited)
     mi ∈ visited && return nothing
     push!(visited, mi)
-    operation(mi)
-    if isdefined(mi, :backedges)
-        for be in mi.backedges
-            visit_backedges(operation, be, visited)
+    status = operation(mi)
+    if status && isdefined(mi, :backedges)
+        for edge in mi.backedges
+            visit_backedges(operation, edge, visited)
         end
     end
     return nothing

test/runtests.jl

@@ -23,34 +23,25 @@ end
     @test !Outer.callh(1.0)
     @test Outer.callcallh(1)
 
-    function mmatches(mi, mod, name, spectypes)
-        mi.def.module === mod || return false
-        mi.def.name === name || return false
-        for (t1, t2) in zip(Iterators.drop(mi.specTypes.parameters, 1), spectypes)
-            t1 === t2 || return false
-        end
-        return true
-    end
-
     mis = Core.MethodInstance[]
     visit(Outer) do x
         isa(x, Core.MethodInstance) && push!(mis, x)
     end
-    @test any(mi->mmatches(mi, Outer, :f, (Nothing,)), mis)
-    @test count(mi->mmatches(mi, Outer.Inner, :g, (String,)), mis) == 1
-    @test !any(mi->mmatches(mi, Outer.Inner, :g, (SubString,)), mis)
-    @test any(mi->mmatches(mi, Outer.Inner, :h, (Int,)), mis)
-    @test any(mi->mmatches(mi, Outer.Inner, :h, (Float64,)), mis)
-    @test any(mi->mmatches(mi, Outer, :callh, (Int,)), mis)
-    @test any(mi->mmatches(mi, Outer, :callh, (Float64,)), mis)
-    @test any(mi->mmatches(mi, Outer, :callcallh, (Int,)), mis)
-    @test !any(mi->mmatches(mi, Outer, :callcallh, (Float64,)), mis)
+    @test any(mi->mi.specTypes === Tuple{typeof(Outer.f), Nothing}, mis)
+    @test count(mi->mi.specTypes === Tuple{typeof(Outer.Inner.g), String}, mis) == 1
+    @test !any(mi->mi.specTypes=== Tuple{typeof(Outer.Inner.g), SubString{String}}, mis)
+    @test any(mi->mi.specTypes === Tuple{typeof(Outer.Inner.h), Int}, mis)
+    @test any(mi->mi.specTypes === Tuple{typeof(Outer.Inner.h), Float64}, mis)
+    @test any(mi->mi.specTypes === Tuple{typeof(Outer.callh), Int}, mis)
+    @test any(mi->mi.specTypes === Tuple{typeof(Outer.callh), Float64}, mis)
+    @test any(mi->mi.specTypes === Tuple{typeof(Outer.callcallh), Int}, mis)
+    @test !any(mi->mi.specTypes === Tuple{typeof(Outer.callcallh), Float64}, mis)
 
-    mi = mis[findfirst(mi->mmatches(mi, Outer.Inner, :h, (Int,)), mis)]
+    mi = mis[findfirst(mi->mi.specTypes === Tuple{typeof(Outer.Inner.h), Int}, mis)]
     @test length(all_backedges(mi)) == 2
-    @test terminal_backedges(mi) == [mis[findfirst(mi->mmatches(mi, Outer, :callcallh, (Int,)), mis)]]
-    mi = mis[findfirst(mi->mmatches(mi, Outer.Inner, :h, (Float64,)), mis)]
+    @test terminal_backedges(mi) == [mis[findfirst(mi->mi.specTypes === Tuple{typeof(Outer.callcallh), Int}, mis)]]
+    mi = mis[findfirst(mi->mi.specTypes === Tuple{typeof(Outer.Inner.h), Float64}, mis)]
     @test length(all_backedges(mi)) == 1
-    @test terminal_backedges(mi) == [mis[findfirst(mi->mmatches(mi, Outer, :callh, (Float64,)), mis)]]
+    @test terminal_backedges(mi) == [mis[findfirst(mi->mi.specTypes === Tuple{typeof(Outer.callh), Float64}, mis)]]
     @test all_backedges(mi) == terminal_backedges(mi)
 end