Correctly compile splats in for-loop index in prism

Fixes [Bug #21648]

This is a followup to https://github.com/ruby/ruby/pull/13597.
The added test passed but didn't emit the same instructions.

This also handles bare splats and aligns instructions for all cases

2 files changed, 17 insertions, 8 deletions

diff --git a/prism_compile.c b/prism_compile.c
index 6b4e32f629..811ce47816 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -5338,8 +5338,7 @@ pm_compile_for_node_index(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *c
       case PM_INSTANCE_VARIABLE_TARGET_NODE:
       case PM_CONSTANT_PATH_TARGET_NODE:
       case PM_CALL_TARGET_NODE:
-      case PM_INDEX_TARGET_NODE:
-      case PM_SPLAT_NODE: {
+      case PM_INDEX_TARGET_NODE: {
         // For other targets, we need to potentially compile the parent or
         // owning expression of this target, then retrieve the value, expand it,
         // and then compile the necessary writes.
@@ -5359,6 +5358,7 @@ pm_compile_for_node_index(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *c
         pm_multi_target_state_update(&state);
         break;
       }
+      case PM_SPLAT_NODE:
       case PM_MULTI_TARGET_NODE: {
         DECL_ANCHOR(writes);
         DECL_ANCHOR(cleanup);
@@ -5394,6 +5394,12 @@ pm_compile_for_node_index(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *c
         PUSH_INSN(ret, location, pop);
 
         PUSH_LABEL(ret, not_single);
+
+        if (PM_NODE_TYPE_P(node, PM_SPLAT_NODE)) {
+            const pm_splat_node_t *cast = (const pm_splat_node_t *) node;
+            PUSH_INSN2(ret, location, expandarray, INT2FIX(0), INT2FIX(cast->expression == NULL ? 0 : 1));
+        }
+
         PUSH_SEQ(ret, writes);
         PUSH_SEQ(ret, cleanup);
         break;
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index b95add5bd4..76b961b37e 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -1046,16 +1046,19 @@ module Prism
     end
 
     def test_ForNode
-      assert_prism_eval("for i in [1,2] do; i; end")
-      assert_prism_eval("for @i in [1,2] do; @i; end")
-      assert_prism_eval("for $i in [1,2] do; $i; end")
+      assert_prism_eval("r = []; for i in [1,2] do; r << i; end; r")
+      assert_prism_eval("r = []; for @i in [1,2] do; r << @i; end; r")
+      assert_prism_eval("r = []; for $i in [1,2] do; r << $i; end; r")
 
-      assert_prism_eval("for foo, in  [1,2,3] do end")
+      assert_prism_eval("r = []; for foo, in  [1,2,3] do r << foo end; r")
 
-      assert_prism_eval("for i, j in {a: 'b'} do; i; j; end")
+      assert_prism_eval("r = []; for i, j in {a: 'b'} do; r << [i, j]; end; r")
 
       # Test splat node as index in for loop
-      assert_prism_eval("for *x in [[1,2], [3,4]] do; x; end")
+      assert_prism_eval("r = []; for *x in [[1,2], [3,4]] do; r << x; end; r")
+      assert_prism_eval("r = []; for * in [[1,2], [3,4]] do; r << 'ok'; end; r")
+      assert_prism_eval("r = []; for x, * in [[1,2], [3,4]] do; r << x; end; r")
+      assert_prism_eval("r = []; for x, *y in [[1,2], [3,4]] do; r << [x, y]; end; r")
     end
 
     ############################################################################