Fix coverage support for RubyVM::ISeq.compile

[Bug #22018]

ISeq returned by `RubyVM::InstructionSequene.load_iseq` weren't
handled by the coverage module.

2 files changed, 75 insertions, 6 deletions

diff --git a/iseq.c b/iseq.c
index be883e7790..099520582c 100644
--- a/iseq.c
+++ b/iseq.c
@@ -862,10 +862,6 @@ set_compile_option_from_hash(rb_compile_option_t *option, VALUE opt)
 static rb_compile_option_t *
 set_compile_option_from_ast(rb_compile_option_t *option, const rb_ast_body_t *ast)
 {
-#define SET_COMPILE_OPTION(o, a, mem) \
-    ((a)->mem < 0 ? 0 : ((o)->mem = (a)->mem > 0))
-    SET_COMPILE_OPTION(option, ast, coverage_enabled);
-#undef SET_COMPILE_OPTION
     if (ast->frozen_string_literal >= 0) {
         option->frozen_string_literal = ast->frozen_string_literal;
     }
@@ -1019,8 +1015,13 @@ rb_iseq_new_eval(const VALUE ast_value, VALUE name, VALUE path, VALUE realpath,
         }
     }
 
+    rb_compile_option_t option = COMPILE_OPTION_DEFAULT;
+    rb_ast_t *ast = rb_ruby_ast_data_get(ast_value);
+    if (ast->body.coverage_enabled >= 0) {
+        option.coverage_enabled = ast->body.coverage_enabled;
+    }
     return rb_iseq_new_with_opt(ast_value, name, path, realpath, first_lineno,
-                                parent, isolated_depth, ISEQ_TYPE_EVAL, &COMPILE_OPTION_DEFAULT,
+                                parent, isolated_depth, ISEQ_TYPE_EVAL, &option,
                                 Qnil);
 }
 
@@ -1188,7 +1189,7 @@ rb_iseq_load_iseq(VALUE fname)
     VALUE iseqv = rb_check_funcall(rb_cISeq, rb_intern("load_iseq"), 1, &fname);
 
     if (!SPECIAL_CONST_P(iseqv) && RBASIC_CLASS(iseqv) == rb_cISeq) {
-        return  iseqw_check(iseqv);
+        return iseqw_check(iseqv);
     }
 
     return NULL;
@@ -1376,6 +1377,7 @@ rb_iseq_compile_with_option(VALUE src, VALUE file, VALUE realpath, VALUE line, V
         rb_exc_raise(GET_EC()->errinfo);
     }
     else {
+        iseq_new_setup_coverage(file, ast_line_count(ast_value));
         iseq = rb_iseq_new_with_opt(ast_value, name, file, realpath, ln,
                                     NULL, 0, ISEQ_TYPE_TOP, &option,
                                     Qnil);
@@ -1440,6 +1442,7 @@ pm_iseq_compile_with_option(VALUE src, VALUE file, VALUE realpath, VALUE line, V
 
     if (error == Qnil) {
         int error_state;
+        iseq_new_setup_coverage(file, (int) (result.node.parser->newline_list.size - 1));
         iseq = pm_iseq_new_with_opt(&result.node, name, file, realpath, ln, NULL, 0, ISEQ_TYPE_TOP, &option, &error_state);
 
         pm_parse_result_free(&result);
@@ -1825,6 +1828,7 @@ iseqw_s_compile_file(int argc, VALUE *argv, VALUE self)
     parser = rb_parser_new();
     rb_parser_set_context(parser, NULL, FALSE);
     ast_value = rb_parser_load_file(parser, file);
+    iseq_new_setup_coverage(file, ast_line_count(ast_value));
     ast = rb_ruby_ast_data_get(ast_value);
     if (!ast->body.root) exc = GET_EC()->errinfo;
 
@@ -1899,6 +1903,7 @@ iseqw_s_compile_file_prism(int argc, VALUE *argv, VALUE self)
         make_compile_option(&option, opt);
 
         int error_state;
+        iseq_new_setup_coverage(file, (int) (result.node.parser->newline_list.size - 1));
         rb_iseq_t *iseq = pm_iseq_new_with_opt(&result.node, rb_fstring_lit("<main>"),
                                                file,
                                                rb_realpath_internal(Qnil, file, 1),
diff --git a/test/coverage/test_coverage.rb b/test/coverage/test_coverage.rb
index 80f8930472..adcd4a946c 100644
--- a/test/coverage/test_coverage.rb
+++ b/test/coverage/test_coverage.rb
@@ -82,6 +82,70 @@ class TestCoverage < Test::Unit::TestCase
     }
   end
 
+  def test_coverage_snapshot_iseq_compile
+    Dir.mktmpdir {|tmp|
+      Dir.chdir(tmp) {
+        File.open("test.rb", "w") do |f|
+          f.puts <<-EOS
+            def coverage_test_snapshot
+              :ok
+            end
+          EOS
+        end
+
+        assert_in_out_err(ARGV, <<-"end;", ["[1, 0, nil]", "[1, 1, nil]", "[1, 1, nil]"], [])
+          class RubyVM::InstructionSequence
+            def self.load_iseq(path)
+              compile(File.read(path), path, path)
+            end
+          end
+
+          Coverage.start
+          tmp = Dir.pwd
+          require tmp + "/test.rb"
+          cov = Coverage.peek_result[tmp + "/test.rb"]
+          coverage_test_snapshot
+          cov2 = Coverage.peek_result[tmp + "/test.rb"]
+          p cov
+          p cov2
+          p Coverage.result[tmp + "/test.rb"]
+        end;
+      }
+    }
+  end
+
+  def test_coverage_snapshot_iseq_compile_file
+    Dir.mktmpdir {|tmp|
+      Dir.chdir(tmp) {
+        File.open("test.rb", "w") do |f|
+          f.puts <<-EOS
+            def coverage_test_snapshot
+              :ok
+            end
+          EOS
+        end
+
+        assert_in_out_err(ARGV, <<-"end;", ["[1, 0, nil]", "[1, 1, nil]", "[1, 1, nil]"], [])
+          class RubyVM::InstructionSequence
+            def self.load_iseq(path)
+              compile_file(path)
+            end
+          end
+
+          Coverage.start
+          tmp = Dir.pwd
+          require tmp + "/test.rb"
+          cov = Coverage.peek_result[tmp + "/test.rb"]
+          coverage_test_snapshot
+          cov2 = Coverage.peek_result[tmp + "/test.rb"]
+          p cov
+          p cov2
+          p Coverage.result[tmp + "/test.rb"]
+        end;
+      }
+    }
+  end
+
   def test_restarting_coverage
     Dir.mktmpdir {|tmp|
       Dir.chdir(tmp) {