fix code_coverage_effect

Author: Roger-luo

Project.toml

@@ -3,6 +3,7 @@ uuid = "6b7a57c9-7cc1-4fdf-b7f5-e857abae3638"
 version = "0.1.5"
 
 [deps]
+CodeInfoTools = "bc773b8a-8374-437a-b9f2-0e9785855863"
 Expronicon = "6b7a57c9-7cc1-4fdf-b7f5-e857abae3636"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 

src/patches.jl

@@ -101,3 +101,7 @@ Base.iterate(p::Core.Compiler.Pair) = Core.Compiler.iterate(p)
 Base.iterate(p::Core.Compiler.Pair, st) = Core.Compiler.iterate(p, st)
 
 Base.getindex(m::Core.Compiler.MethodLookupResult, idx::Int) = Core.Compiler.getindex(m, idx)
+
+Base.length(x::Core.Compiler.BitSet) = Core.Compiler.length(x)
+Base.iterate(x::Core.Compiler.BitSet) = Core.Compiler.iterate(x)
+Base.iterate(x::Core.Compiler.BitSet, st) = Core.Compiler.iterate(x, st)

src/utils.jl

@@ -85,6 +85,18 @@ function rm_code_coverage_effect(ci::CodeInfo)
     for (v, stmt) in new
         if Meta.isexpr(stmt, :code_coverage_effect)
             delete!(new, v)
+        elseif stmt isa GotoIfNot
+            # idx always exists since code_coverage_effect was inserted
+            idx = findfirst(stmt.dest:length(ci.code)) do k
+                !Meta.isexpr(ci.code[k], :code_coverage_effect)
+            end
+            new[v] = GotoIfNot(stmt.cond, stmt.dest+idx-1)
+        elseif stmt isa GotoNode
+            # idx always exists since code_coverage_effect was inserted
+            idx = findfirst(stmt.label:length(ci.code)) do k
+                !Meta.isexpr(ci.code[k], :code_coverage_effect)
+            end
+            new[v] = GotoNode(stmt.label+idx-1)
         end
     end
     new_ci = finish(new)

test/codeinfo.jl

@@ -156,3 +156,27 @@ end
         #=%8=# ReturnNode(SSAValue(7))::Float64     
     end
 end
+
+@testset "test code_coverage_effect bb start" begin
+    ci = @make_codeinfo begin
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:call, GlobalRef(Main, :foo), 2)::Nothing
+        Expr(:call, GlobalRef(Main, :measure), 2)::Int
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:call, GlobalRef(Main, :measure_cmp), SSAValue(6), 1)::Bool
+        GotoIfNot(SSAValue(9), 14)
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:code_coverage_effect)::Nothing
+        Expr(:call, GlobalRef(Main, :foo), 2)::Nothing
+        Expr(:code_coverage_effect)::Nothing
+        ReturnNode(nothing)::Nothing
+    end
+
+    stmt = CompilerPluginTools.rm_code_coverage_effect(ci).code[4]
+    @test stmt.cond == SSAValue(3)
+    @test stmt.dest == 6
+end