Store BEGIN nodes separately in RootParseNode

Author: bjfish

src/main/java/org/truffleruby/parser/TranslatorDriver.java

@@ -240,6 +240,10 @@ public RubyRootNode parse(RubySource rubySource, ParserContext parserContext, St
                 topLevel);
 
         printParseTranslateExecuteMetric("before-translate", context, source);
+        RubyNode beginNode = null;
+        if (node.getBeginNode() != null) {
+            beginNode = translator.translateNodeOrNil(sourceIndexLength, node.getBeginNode());
+        }
         RubyNode truffleNode = translator.translateNodeOrNil(sourceIndexLength, node.getBodyNode());
         printParseTranslateExecuteMetric("after-translate", context, source);
 
@@ -271,6 +275,12 @@ public RubyRootNode parse(RubySource rubySource, ParserContext parserContext, St
                     Arrays.asList(translator.initFlipFlopStates(sourceIndexLength), truffleNode));
         }
 
+        if (beginNode != null) {
+            truffleNode = Translator.sequence(
+                    sourceIndexLength,
+                    Arrays.asList(beginNode, truffleNode));
+        }
+
         // Catch next
 
         truffleNode = new CatchNextNode(truffleNode);

src/main/java/org/truffleruby/parser/ast/RootParseNode.java

@@ -42,18 +42,24 @@
  * of the code.
  *
  */
-// TODO: Store BEGIN and END information into this node
 public class RootParseNode extends ParseNode {
 
+    private final ParseNode beginNode;
     private final Source source;
     private final ParseNode bodyNode;
     private final String file;
     private final int endPosition;
 
-    public RootParseNode(Source source, SourceIndexLength position, ParseNode bodyNode, String file, int endPosition) {
+    public RootParseNode(
+            Source source,
+            SourceIndexLength position,
+            ParseNode beginNode,
+            ParseNode bodyNode,
+            String file,
+            int endPosition) {
         super(position);
-
         this.source = source;
+        this.beginNode = beginNode;
         this.bodyNode = bodyNode;
         this.file = file;
         this.endPosition = endPosition;
@@ -72,6 +78,10 @@ public String getFile() {
         return file;
     }
 
+    public ParseNode getBeginNode() {
+        return beginNode;
+    }
+
     /**
      * First real AST node to be interpreted
      *

src/main/java/org/truffleruby/parser/parser/ParserSupport.java

@@ -288,28 +288,24 @@ public ParseNode addRootNode(ParseNode topOfAST) {
         final int endPosition = lexer.getEndPosition();
 
         SourceIndexLength position;
-        if (result.getBeginNodes().isEmpty()) {
-            if (topOfAST == null) {
-                topOfAST = NilImplicitParseNode.NIL;
-                position = lexer.getPosition();
-            } else {
-                position = topOfAST.getPosition();
-            }
+
+        if (topOfAST == null) {
+            topOfAST = NilImplicitParseNode.NIL;
+            position = lexer.getPosition();
         } else {
+            position = topOfAST.getPosition();
+        }
+
+        BlockParseNode beginAST = null;
+        if (!result.getBeginNodes().isEmpty()) {
             position = topOfAST != null ? topOfAST.getPosition() : result.getBeginNodes().get(0).getPosition();
-            BlockParseNode newTopOfAST = new BlockParseNode(position);
+            beginAST = new BlockParseNode(position);
             for (ParseNode beginNode : result.getBeginNodes()) {
-                appendToBlock(newTopOfAST, beginNode);
-            }
-
-            // Add real top to new top (unless this top is empty [only begin/end nodes or truly empty])
-            if (topOfAST != null) {
-                newTopOfAST.add(topOfAST);
+                appendToBlock(beginAST, beginNode);
             }
-            topOfAST = newTopOfAST;
         }
 
-        return new RootParseNode(lexer.getSource(), position, topOfAST, lexer.getFile(), endPosition);
+        return new RootParseNode(lexer.getSource(), position, beginAST, topOfAST, lexer.getFile(), endPosition);
     }
 
     /* MRI: block_append */