[PRISM] Fix BEGIN{} execution order

prism_compile.c

@@ -6808,27 +6808,38 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         // ^^^^^^^^
         const pm_pre_execution_node_t *cast = (const pm_pre_execution_node_t *) node;
 
-        DECL_ANCHOR(pre_ex);
-        INIT_ANCHOR(pre_ex);
+        LINK_ANCHOR *outer_pre = scope_node->pre_execution_anchor;
+        RUBY_ASSERT(outer_pre != NULL);
+
+        // BEGIN{} nodes can be nested, so here we're going to do the same thing
+        // that we did for the top-level compilation where we create two
+        // anchors and then join them in the correct order into the resulting
+        // anchor.
+        DECL_ANCHOR(inner_pre);
+        INIT_ANCHOR(inner_pre);
+        scope_node->pre_execution_anchor = inner_pre;
+
+        DECL_ANCHOR(inner_body);
+        INIT_ANCHOR(inner_body);
 
         if (cast->statements != NULL) {
             const pm_node_list_t *body = &cast->statements->body;
+
             for (size_t index = 0; index < body->size; index++) {
-                pm_compile_node(iseq, body->nodes[index], pre_ex, true, scope_node);
+                pm_compile_node(iseq, body->nodes[index], inner_body, true, scope_node);
             }
         }
 
         if (!popped) {
-            PUSH_INSN(pre_ex, location, putnil);
+            PUSH_INSN(inner_body, location, putnil);
         }
 
-        pre_ex->last->next = ret->anchor.next;
-        ret->anchor.next = pre_ex->anchor.next;
-        ret->anchor.next->prev = pre_ex->anchor.next;
-
-        if (ret->last == (LINK_ELEMENT *)ret) {
-            ret->last = pre_ex->last;
-        }
+        // Now that everything has been compiled, join both anchors together
+        // into the correct outer pre execution anchor, and reset the value so
+        // that subsequent BEGIN{} nodes can be compiled correctly.
+        ADD_SEQ(outer_pre, inner_pre);
+        ADD_SEQ(outer_pre, inner_body);
+        scope_node->pre_execution_anchor = outer_pre;
 
         return;
       }
@@ -8342,6 +8353,20 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
     }
 }
 
+/** True if the given iseq can have pre execution blocks. */
+static inline bool
+pm_iseq_pre_execution_p(rb_iseq_t *iseq)
+{
+    switch (ISEQ_BODY(iseq)->type) {
+      case ISEQ_TYPE_TOP:
+      case ISEQ_TYPE_EVAL:
+      case ISEQ_TYPE_MAIN:
+        return true;
+      default:
+        return false;
+    }
+}
+
 /**
  * This is the main entry-point into the prism compiler. It accepts the iseq
  * that it should be compiling instruction into and a pointer to the scope node
@@ -8355,7 +8380,32 @@ pm_iseq_compile_node(rb_iseq_t *iseq, pm_scope_node_t *node)
     DECL_ANCHOR(ret);
     INIT_ANCHOR(ret);
 
-    pm_compile_node(iseq, (const pm_node_t *) node, ret, false, node);
+    if (pm_iseq_pre_execution_p(iseq)) {
+        // Because these ISEQs can have BEGIN{}, we're going to create two
+        // anchors to compile them, a "pre" and a "body". We'll mark the "pre"
+        // on the scope node so that when BEGIN{} is found, its contents will be
+        // added to the "pre" anchor.
+        DECL_ANCHOR(pre);
+        INIT_ANCHOR(pre);
+        node->pre_execution_anchor = pre;
+
+        // Now we'll compile the body as normal. We won't compile directly into
+        // the "ret" anchor yet because we want to add the "pre" anchor to the
+        // beginning of the "ret" anchor first.
+        DECL_ANCHOR(body);
+        INIT_ANCHOR(body);
+        pm_compile_node(iseq, (const pm_node_t *) node, body, false, node);
+
+        // Now we'll join both anchors together so that the content is in the
+        // correct order.
+        ADD_SEQ(ret, pre);
+        ADD_SEQ(ret, body);
+    }
+    else {
+        // In other circumstances, we can just compile the node directly into
+        // the "ret" anchor.
+        pm_compile_node(iseq, (const pm_node_t *) node, ret, false, node);
+    }
 
     CHECK(iseq_setup_insn(iseq, ret));
     return iseq_setup(iseq, ret);

prism_compile.h

@@ -13,6 +13,9 @@ typedef struct pm_local_index_struct {
     int index, level;
 } pm_local_index_t;
 
+// A declaration for the struct that lives in compile.c.
+struct iseq_link_anchor;
+
 // ScopeNodes are helper nodes, and will never be part of the AST. We manually
 // declare them here to avoid generating them.
 typedef struct pm_scope_node {
@@ -40,6 +43,12 @@ typedef struct pm_scope_node {
 
     ID *constants;
     st_table *index_lookup_table;
+
+    /**
+     * This will only be set on the top-level scope node. It will contain all of
+     * the instructions pertaining to BEGIN{} nodes.
+     */
+    struct iseq_link_anchor *pre_execution_anchor;
 } pm_scope_node_t;
 
 void pm_scope_node_init(const pm_node_t *node, pm_scope_node_t *scope, pm_scope_node_t *previous);

spec/prism.mspec

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 ## Language
-MSpec.register(:exclude, "The BEGIN keyword runs multiple begins in FIFO order")
 MSpec.register(:exclude, "Executing break from within a block works when passing through a super call")
 MSpec.register(:exclude, "The defined? keyword when called with a method name in a void context warns about the void context when parsing it")
 MSpec.register(:exclude, "Hash literal expands an '**{}' or '**obj' element with the last key/value pair taking precedence")