Test infrastructure for limiting output to one method

When testing the IR generated by closures or complex function desugaring
there can be a lot of methods generated and it can be hard to review changes.
Add a `method_filter` to test cases to limit the output to a method of
interest.

Author: c42f

src/utils.jl

@@ -75,7 +75,46 @@ function subscript_str(i)
              "5"=>"₅", "6"=>"₆", "7"=>"₇", "8"=>"₈", "9"=>"₉")
 end
 
-function print_ir(io::IO, ex, indent="")
+function _deref_ssa(stmts, ex)
+    while kind(ex) == K"SSAValue"
+        ex = stmts[ex.var_id]
+    end
+    ex
+end
+
+function _find_method_lambda(ex, name)
+    @assert kind(ex) == K"code_info"
+    # Heuristic search through outer thunk for the method in question.
+    method_found = false
+    stmts = children(ex[1])
+    for e in stmts
+        if kind(e) == K"method" && numchildren(e) >= 2
+            sig = _deref_ssa(stmts, e[2])
+            @assert kind(sig) == K"call"
+            arg_types = _deref_ssa(stmts, sig[2])
+            @assert kind(arg_types) == K"call"
+            self_type = _deref_ssa(stmts, arg_types[2])
+            if kind(self_type) == K"globalref" && occursin(name, self_type.name_val)
+                return e[3]
+            end
+        end
+    end
+end
+
+function print_ir(io::IO, ex, method_filter=nothing)
+    @assert kind(ex) == K"code_info"
+    if !isnothing(method_filter)
+        filtered = _find_method_lambda(ex, method_filter)
+        if isnothing(filtered)
+            @warn "Method not found with method filter $method_filter"
+        else
+            ex = filtered
+        end
+    end
+    _print_ir(io, ex, "")
+end
+
+function _print_ir(io::IO, ex, indent)
     added_indent = "    "
     @assert (kind(ex) == K"lambda" || kind(ex) == K"code_info") && kind(ex[1]) == K"block"
     if !ex.is_toplevel_thunk && kind(ex) == K"code_info"
@@ -105,7 +144,7 @@ function print_ir(io::IO, ex, indent="")
             print(io, indent, lno, " --- method ", string(e[1]), " ", string(e[2]))
             if kind(e[3]) == K"lambda" || kind(e[3]) == K"code_info"
                 println(io)
-                print_ir(io, e[3], indent*added_indent)
+                _print_ir(io, e[3], indent*added_indent)
             else
                 println(io, " ", string(e[3]))
             end
@@ -116,10 +155,10 @@ function print_ir(io::IO, ex, indent="")
                 print(io, " ", e[i])
             end
             println(io)
-            print_ir(io, e[5], indent*added_indent)
+            _print_ir(io, e[5], indent*added_indent)
         elseif kind(e) == K"code_info"
             println(io, indent, lno, " --- ", e.is_toplevel_thunk ? "thunk" : "code_info")
-            print_ir(io, e, indent*added_indent)
+            _print_ir(io, e, indent*added_indent)
         else
             code = string(e)
             println(io, indent, lno, " ", code)

test/closures_ir.jl

@@ -287,66 +287,26 @@ end
 ########################################
 # Nested captures - here `g` captures `x` because it is needed to initialize
 # the closure `h` which captures both `x` and `y`.
-function f(x)
-    function g(y)
-        function h(z)
+# [method_filter: #f_nest#g_nest##0]
+function f_nest(x)
+    function g_nest(y)
+        function h_nest(z)
             (x,y,z)
         end
     end
 end
 #---------------------
-1   (method TestMod.f)
-2   (call core.svec :x)
-3   (call core.svec false)
-4   (call JuliaLowering.eval_closure_type TestMod :#f#g##4 %₂ %₃)
-5   (call core.svec :x :y)
-6   (call core.svec false false)
-7   (call JuliaLowering.eval_closure_type TestMod :#f#g#h##0 %₅ %₆)
-8   TestMod.#f#g#h##0
-9   (call core.svec %₈ core.Any)
-10  (call core.svec)
-11  SourceLocation::3:18
-12  (call core.svec %₉ %₁₀ %₁₁)
-13  --- method core.nothing %₁₂
-    slots: [slot₁/#self#(!read) slot₂/z]
-    1   (call core.getfield slot₁/#self# :x)
-    2   (call core.getfield slot₁/#self# :y)
-    3   (call core.tuple %₁ %₂ slot₂/z)
-    4   (return %₃)
-14  TestMod.#f#g##4
-15  (call core.svec %₁₄ core.Any)
-16  (call core.svec)
-17  SourceLocation::2:14
-18  (call core.svec %₁₅ %₁₆ %₁₇)
-19  --- method core.nothing %₁₈
-    slots: [slot₁/#self#(!read) slot₂/y(!read) slot₃/h]
-    1   TestMod.#f#g#h##0
-    2   (call core.getfield slot₁/#self# :x)
-    3   (call core.typeof %₂)
-    4   (call core.typeof slot₂/y)
-    5   (call core.apply_type %₁ %₃ %₄)
-    6   (call core.getfield slot₁/#self# :x)
-    7   (new %₅ %₆ slot₂/y)
-    8   (= slot₃/h %₇)
-    9   slot₃/h
-    10  (return %₉)
-20  TestMod.f
-21  (call core.Typeof %₂₀)
-22  (call core.svec %₂₁ core.Any)
-23  (call core.svec)
-24  SourceLocation::1:10
-25  (call core.svec %₂₂ %₂₃ %₂₄)
-26  --- method core.nothing %₂₅
-    slots: [slot₁/#self#(!read) slot₂/x(!read) slot₃/g]
-    1   TestMod.#f#g##4
-    2   (call core.typeof slot₂/x)
-    3   (call core.apply_type %₁ %₂)
-    4   (new %₃ slot₂/x)
-    5   (= slot₃/g %₄)
-    6   slot₃/g
-    7   (return %₆)
-27  TestMod.f
-28  (return %₂₇)
+slots: [slot₁/#self#(!read) slot₂/y(!read) slot₃/h_nest]
+1   TestMod.#f_nest#g_nest#h_nest##0
+2   (call core.getfield slot₁/#self# :x)
+3   (call core.typeof %₂)
+4   (call core.typeof slot₂/y)
+5   (call core.apply_type %₁ %₃ %₄)
+6   (call core.getfield slot₁/#self# :x)
+7   (new %₅ %₆ slot₂/y)
+8   (= slot₃/h_nest %₇)
+9   slot₃/h_nest
+10  (return %₉)
 
 ########################################
 # Global method capturing local variables

test/utils.jl

@@ -126,8 +126,13 @@ function match_ir_test_case(case_str)
                     error("Too many sections in IR test case")
     expect_error = startswith(description, "Error")
     is_broken = startswith(description, "FIXME")
+    method_filter = begin
+        mf = match(r"\[method_filter: *(.*)\]", description)
+        isnothing(mf) ? nothing : strip(mf[1])
+    end
     (; expect_error=expect_error, is_broken=is_broken,
      description=strip(description),
+     method_filter=method_filter,
      input=strip(input), output=strip(output))
 end
 
@@ -152,49 +157,49 @@ function setup_ir_test_module(preamble)
     test_mod
 end
 
-function format_ir_for_test(mod, description, input, expect_error=false, is_broken=false)
-    ex = parsestmt(SyntaxTree, input)
+function format_ir_for_test(mod, case)
+    ex = parsestmt(SyntaxTree, case.input)
     try
         if kind(ex) == K"macrocall" && kind(ex[1]) == K"MacroName" && ex[1].name_val == "@ast_"
             # Total hack, until @ast_ can be implemented in terms of new-style
             # macros.
             ex = JuliaLowering.eval(mod, Expr(ex))
         end
         x = JuliaLowering.lower(mod, ex)
-        if expect_error
-            error("Expected a lowering error in test case \"$description\"")
+        if case.expect_error
+            error("Expected a lowering error in test case \"$(case.description)\"")
         end
-        ir = strip(sprint(JuliaLowering.print_ir, x))
+        ir = strip(sprint(JuliaLowering.print_ir, x, case.method_filter))
         return replace(ir, string(mod)=>"TestMod")
     catch exc
         if exc isa InterruptException
             rethrow()
-        elseif expect_error && (exc isa LoweringError)
+        elseif case.expect_error && (exc isa LoweringError)
             return sprint(io->Base.showerror(io, exc, show_detail=false))
-        elseif expect_error && (exc isa MacroExpansionError)
+        elseif case.expect_error && (exc isa MacroExpansionError)
             return sprint(io->Base.showerror(io, exc))
-        elseif is_broken
+        elseif case.is_broken
             return sprint(io->Base.showerror(io, exc))
         else
-            throw("Error in test case \"$description\"")
+            throw("Error in test case \"$(case.description)\"")
         end
     end
 end
 
 function test_ir_cases(filename::AbstractString)
     preamble, cases = read_ir_test_cases(filename)
     test_mod = setup_ir_test_module(preamble)
-    for (expect_error, is_broken, description, input, ref) in cases
-        if is_broken
+    for case in cases
+        if case.is_broken
             continue
         end
-        output = format_ir_for_test(test_mod, description, input, expect_error)
-        @testset "$description" begin
-            if output != ref
+        output = format_ir_for_test(test_mod, case)
+        @testset "$(case.description)" begin
+            if output != case.output
                 # Do additional error dumping, as @test will not format errors in a nice way
-                @error "Test \"$description\" failed" output=Text(output) ref=Text(ref)
+                @error "Test \"$(case.description)\" failed" output=Text(output) ref=Text(case.output)
             end
-            @test output == ref
+            @test output == case.output
         end
     end
 end
@@ -213,20 +218,20 @@ function refresh_ir_test_cases(filename, pattern=nothing)
         println(io, preamble, "\n")
         println(io, "#*******************************************************************************")
     end
-    for (expect_error, is_broken, description, input, ref) in cases
-        if isnothing(pattern) || occursin(pattern, description)
-            ir = format_ir_for_test(test_mod, description, input, expect_error, is_broken)
-            if rstrip(ir) != ref
-                @info "Refreshing test case $(repr(description)) in $filename"
+    for case in cases
+        if isnothing(pattern) || occursin(pattern, case.description)
+            ir = format_ir_for_test(test_mod, case)
+            if rstrip(ir) != case.output
+                @info "Refreshing test case $(repr(case.description)) in $filename"
             end
         else
-            ir = ref
+            ir = case.output
         end
         println(io,
             """
             ########################################
-            $(comment_description(description))
-            $(strip(input))
+            $(comment_description(case.description))
+            $(strip(case.input))
             #---------------------
             $ir
             """