From 6a51c3e80c0901851c252ed4d1387e43939a452f Mon Sep 17 00:00:00 2001
From: Yusuke Endoh <mame@ruby-lang.org>
Date: Sun, 19 Dec 2021 04:00:51 +0900
Subject: Make AST.of possible even under eval when keep_script_lines is
 enabled

Now the following code works without an exception.

```
RubyVM.keep_script_lines = true

eval(<<END)
def foo
end
END

p RubyVM::AbstractSyntaxTree.of(method(:foo))
```
---
 ast.c                 |  4 +--
 test/ruby/test_ast.rb | 67 ++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/ast.c b/ast.c
index e180213ef7..466e7a6a2e 100644
--- a/ast.c
+++ b/ast.c
@@ -222,11 +222,11 @@ ast_s_of(rb_execution_context_t *ec, VALUE module, VALUE body, VALUE keep_script
     if (!iseq) {
         return Qnil;
     }
-    if (rb_iseq_from_eval_p(iseq)) {
+    lines = iseq->body->variable.script_lines;
+    if (NIL_P(lines) && rb_iseq_from_eval_p(iseq)) {
         rb_raise(rb_eArgError, "cannot get AST for method defined in eval");
     }
     path = rb_iseq_path(iseq);
-    lines = iseq->body->variable.script_lines;
 
     if (!NIL_P(lines) || !NIL_P(lines = script_lines(path))) {
         node = rb_ast_parse_array(lines, keep_script_lines);
diff --git a/test/ruby/test_ast.rb b/test/ruby/test_ast.rb
index 2ab9f4e294..63362096f9 100644
--- a/test/ruby/test_ast.rb
+++ b/test/ruby/test_ast.rb
@@ -226,6 +226,9 @@ class TestAst < Test::Unit::TestCase
   end
 
   def test_of_proc_and_method_under_eval
+    keep_script_lines_back = RubyVM.keep_script_lines
+    RubyVM.keep_script_lines = false
+
     method = self.method(eval("def example_method_#{$$}; end"))
     assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }
 
@@ -246,18 +249,76 @@ class TestAst < Test::Unit::TestCase
 
     method = eval("Class.new{def example_method; end}.instance_method(:example_method)")
     assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(method) }
+
+  ensure
+    RubyVM.keep_script_lines = keep_script_lines_back
+  end
+
+  def test_of_proc_and_method_under_eval_with_keep_script_lines
+    keep_script_lines_back = RubyVM.keep_script_lines
+    RubyVM.keep_script_lines = true
+
+    method = self.method(eval("def example_method_#{$$}; end"))
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))
+
+    method = self.method(eval("def self.example_singleton_method_#{$$}; end"))
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))
+
+    method = eval("proc{}")
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))
+
+    method = self.method(eval("singleton_class.define_method(:example_define_method_#{$$}){}"))
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))
+
+    method = self.method(eval("define_singleton_method(:example_dsm_#{$$}){}"))
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))
+
+    method = eval("Class.new{def example_method; end}.instance_method(:example_method)")
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))
+
+    method = eval("Class.new{def example_method; end}.instance_method(:example_method)")
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, RubyVM::AbstractSyntaxTree.of(method))
+
+  ensure
+    RubyVM.keep_script_lines = keep_script_lines_back
   end
 
   def test_of_backtrace_location_under_eval
+    keep_script_lines_back = RubyVM.keep_script_lines
+    RubyVM.keep_script_lines = false
+
+    m = Module.new do
+      eval(<<-END, nil, __FILE__, __LINE__)
+        def self.sample_backtrace_location
+          caller_locations(0).first
+        end
+      END
+    end
+    backtrace_location = m.sample_backtrace_location
+    assert_raise(ArgumentError) { RubyVM::AbstractSyntaxTree.of(backtrace_location) }
+
+  ensure
+    RubyVM.keep_script_lines = keep_script_lines_back
+  end
+
+  def test_of_backtrace_location_under_eval_with_keep_script_lines
+    keep_script_lines_back = RubyVM.keep_script_lines
+    RubyVM.keep_script_lines = true
+
     m = Module.new do
       eval(<<-END, nil, __FILE__, __LINE__)
         def self.sample_backtrace_location
-          [caller_locations(0).first, __LINE__]
+          caller_locations(0).first
         end
       END
     end
-    backtrace_location, lineno = m.sample_backtrace_location
-    assert_raise(ArgumentError) {  RubyVM::AbstractSyntaxTree.of(backtrace_location) }
+    backtrace_location = m.sample_backtrace_location
+    node = RubyVM::AbstractSyntaxTree.of(backtrace_location)
+    assert_instance_of(RubyVM::AbstractSyntaxTree::Node, node)
+    assert_equal(2, node.first_lineno)
+
+  ensure
+    RubyVM.keep_script_lines = keep_script_lines_back
   end
 
   def test_of_c_method
-- 
cgit v1.2.3