support test_codeinfo

Author: Roger-luo

src/CompilerPluginTools.jl

@@ -57,6 +57,7 @@ export
     method_instance,
     @make_codeinfo,
     @make_ircode,
+    @test_codeinfo,
     @intrinsic_stub
 
 using Base:

src/utils.jl

@@ -74,3 +74,39 @@ function method_instance(@nospecialize(f), @nospecialize(tt), world=Base.get_wor
     return ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
         (Any, Any, Any, UInt), meth, ti, env, world)
 end
+
+macro test_codeinfo(ci, ex)
+    esc(test_codeinfo_m(ci, ex))
+end
+
+function test_codeinfo_m(ci, ex::Expr)
+    Meta.isexpr(ex, :block) || error("expect a begin ... end")
+    ret = Expr(:block)
+    stmt_count = 1
+    for each in ex.args
+        @switch each begin
+            @case :($stmt::$type)
+                push!(ret.args, quote
+                    @test $MLStyle.@match $ci.code[$stmt_count] begin
+                        $stmt => true
+                        _ => false
+                    end
+                end)
+                push!(ret.args, quote
+                    @test $ci.ssavaluetypes[$stmt_count] == $type
+                end)
+                stmt_count += 1
+            @case ::LineNumberNode
+                continue
+            @case _
+                push!(ret.args, quote
+                    @test $MLStyle.@match $ci.code[$stmt_count] begin
+                        $each => true
+                        _ => false
+                    end
+                end)
+                stmt_count += 1
+        end        
+    end
+    return ret
+end

test/codeinfo.jl

@@ -20,12 +20,14 @@ end
         end
     end
 
-    @test test_ci.code[1] == :(overdub($(GlobalRef(Main, :>)), $(SlotNumber(2)), 0))
-    @test test_ci.code[2] isa GotoIfNot
-    @test test_ci.code[3] == :(overdub($(GlobalRef(Main, :*)), 2, $(SlotNumber(2))))
-    @test test_ci.code[4] == ReturnNode(SSAValue(3))
-    @test test_ci.code[5] == :(overdub($(GlobalRef(Main, :*)), 3, $(SlotNumber(2))))
-    @test test_ci.code[6] == ReturnNode(SSAValue(5))
+    @test_codeinfo test_ci begin
+        :(overdub($(GlobalRef(Main, :>)), $(SlotNumber(2)), 0))
+        ::GotoIfNot
+        :(overdub($(GlobalRef(Main, :*)), 2, $(SlotNumber(2))))
+        ReturnNode(SSAValue(3))
+        :(overdub($(GlobalRef(Main, :*)), 3, $(SlotNumber(2))))
+        ReturnNode(SSAValue(5))
+    end
 end
 
 @testset "NewCodeInfo" begin
@@ -45,12 +47,17 @@ end
     end
     test_ci = finish(new)
 
-    @test test_ci.code[1] == :(1 + 1)
-    @test test_ci.code[2] isa GotoIfNot
-    @test test_ci.code[3] == "variable: %3"
-    @test test_ci.code[4] == :(overdub($(SSAValue(3)), $(GlobalRef(Main, :*)), 2, $(SlotNumber(2))))
-    @test test_ci.code[5] == ReturnNode(SSAValue(4))
-    @test test_ci.code[6] == "variable: %5"
+    @test_codeinfo test_ci begin
+        :(1 + 1)
+        ::GotoIfNot
+
+        "variable: %3"
+    
+        :(overdub($(SSAValue(3)), $(GlobalRef(Main, :*)), 2, $(SlotNumber(2))))
+        ReturnNode(SSAValue(4))
+    
+        "variable: %5"
+    end
 
     @testset "slot insertion" begin
         ci = code_lowered(foo, (Float64, ))[1]