ast.rb: RubyVM::AST.parse and .of accepts `save_script_lines: true`

This option makes the parser keep the original source as an array of
the original code lines. This feature exploits the mechanism of
`SCRIPT_LINES__` but records only the specified code that is passed to
RubyVM::AST.of or .parse, instead of recording all parsed program texts.

1 files changed, 42 insertions, 6 deletions

diff --git a/ast.rb b/ast.rb
index 9d4b05bdf1..ce99f53c45 100644
--- a/ast.rb
+++ b/ast.rb
@@ -29,8 +29,8 @@ module RubyVM::AbstractSyntaxTree
   #
   #    RubyVM::AbstractSyntaxTree.parse("x = 1 + 2")
   #    # => #<RubyVM::AbstractSyntaxTree::Node:SCOPE@1:0-1:9>
-  def self.parse string
-    Primitive.ast_s_parse string
+  def self.parse string, save_script_lines: false
+    Primitive.ast_s_parse string, save_script_lines
   end
 
   #  call-seq:
@@ -44,8 +44,8 @@ module RubyVM::AbstractSyntaxTree
   #
   #     RubyVM::AbstractSyntaxTree.parse_file("my-app/app.rb")
   #     # => #<RubyVM::AbstractSyntaxTree::Node:SCOPE@1:0-31:3>
-  def self.parse_file pathname
-    Primitive.ast_s_parse_file pathname
+  def self.parse_file pathname, save_script_lines: false
+    Primitive.ast_s_parse_file pathname, save_script_lines
   end
 
   #  call-seq:
@@ -63,8 +63,8 @@ module RubyVM::AbstractSyntaxTree
   #
   #     RubyVM::AbstractSyntaxTree.of(method(:hello))
   #     # => #<RubyVM::AbstractSyntaxTree::Node:SCOPE@1:0-3:3>
-  def self.of body
-    Primitive.ast_s_of body
+  def self.of body, save_script_lines: false
+    Primitive.ast_s_of body, save_script_lines
   end
 
   # RubyVM::AbstractSyntaxTree::Node instances are created by parse methods in
@@ -139,5 +139,41 @@ module RubyVM::AbstractSyntaxTree
     def inspect
       Primitive.ast_node_inspect
     end
+
+    #  call-seq:
+    #     node.script_lines -> array
+    #
+    #  Returns the original source code as an array of lines.
+    #
+    #  Note that this is an API for ruby internal use, debugging,
+    #  and research. Do not use this for any other purpose.
+    #  The compatibility is not guaranteed.
+    def script_lines
+      Primitive.ast_node_script_lines
+    end
+
+    #  call-seq:
+    #     node.source -> string
+    #
+    #  Returns the code fragment that corresponds to this AST.
+    #
+    #  Note that this is an API for ruby internal use, debugging,
+    #  and research. Do not use this for any other purpose.
+    #  The compatibility is not guaranteed.
+    #
+    #  Also note that this API may return an incomplete code fragment
+    #  that does not parse; for example, a here document following
+    #  an expression may be dropped.
+    def source
+      lines = script_lines
+      if lines
+        lines = lines[first_lineno - 1 .. last_lineno - 1]
+        lines[-1] = lines[-1][0...last_column]
+        lines[0] = lines[0][first_column..-1]
+        lines.join
+      else
+        nil
+      end
+    end
   end
 end