[PRISM] Fix compilation for SplatNodes within ArrayNodes

SplatNodes within ArrayNodes (e.g. [*1..2, 3]) need to be special
cased in the compiler because they use a combination of concatarray
and newarray instructions to treat each sequence of splat or non-splat
elements as independent arrays which get concatenated. This commit
implements those cases.

2 files changed, 75 insertions, 5 deletions

diff --git a/prism_compile.c b/prism_compile.c
index b66c2e0ed0..fa5a7aafca 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -1575,15 +1575,78 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
             pm_array_node_t *cast = (pm_array_node_t *) node;
             pm_node_list_t *elements = &cast->elements;
 
-            for (size_t index = 0; index < elements->size; index++) {
-                PM_COMPILE(elements->nodes[index]);
+            // In the case that there is a splat node within the array,
+            // the array gets compiled slightly differently.
+            if (node->flags & PM_ARRAY_NODE_FLAGS_CONTAINS_SPLAT) {
+                if (elements->size == 1) {
+                    // If the only nodes is a SplatNode, we never
+                    // need to emit the newarray or concatarray
+                    // instructions
+                    PM_COMPILE_NOT_POPPED(elements->nodes[0]);
+                }
+                else {
+                    // We treat all sequences of non-splat elements as their
+                    // own arrays, followed by a newarray, and then continually
+                    // concat the arrays with the SplatNodes
+                    int new_array_size = 0;
+                    bool need_to_concat_array = false;
+                    for (size_t index = 0; index < elements->size; index++) {
+                        pm_node_t *array_element = elements->nodes[index];
+                        if (PM_NODE_TYPE_P(array_element, PM_SPLAT_NODE)) {
+                            pm_splat_node_t *splat_element = (pm_splat_node_t *)array_element;
+
+                            // If we already have non-splat elements, we need to emit a newarray
+                            // instruction
+                            if (new_array_size) {
+                                ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(new_array_size));
+
+                                // We don't want to emit a concat array in the case where
+                                // we're seeing our first splat, and already have elements
+                                if (need_to_concat_array) {
+                                    ADD_INSN(ret, &dummy_line_node, concatarray);
+                                }
+
+                                new_array_size = 0;
+                            }
+
+                            PM_COMPILE_NOT_POPPED(splat_element->expression);
+
+                            if (index > 0) {
+                                ADD_INSN(ret, &dummy_line_node, concatarray);
+                            }
+                            else {
+                                // If this is the first element, we need to splatarray
+                                ADD_INSN1(ret, &dummy_line_node, splatarray, Qtrue);
+                            }
+
+                            need_to_concat_array = true;
+                        }
+                        else {
+                            new_array_size++;
+                            PM_COMPILE_NOT_POPPED(array_element);
+                        }
+                    }
+
+                    if (new_array_size) {
+                        ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(new_array_size));
+                        if (need_to_concat_array) {
+                            ADD_INSN(ret, &dummy_line_node, concatarray);
+                        }
+                    }
+                }
+
+                PM_POP_IF_POPPED;
             }
+            else {
+                for (size_t index = 0; index < elements->size; index++) {
+                    PM_COMPILE(elements->nodes[index]);
+                }
 
-            if (!popped) {
-                ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(elements->size));
+                if (!popped) {
+                    ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(elements->size));
+                }
             }
         }
-
         return;
       }
       case PM_ASSOC_NODE: {
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index a79db57bb5..a01666ee54 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -488,6 +488,13 @@ module Prism
       assert_prism_eval("[1, 2, 3]")
       assert_prism_eval("%i[foo bar baz]")
       assert_prism_eval("%w[foo bar baz]")
+      assert_prism_eval("[*1..2]")
+      assert_prism_eval("[*1..2, 3, 4, *5..6, 7, 8]")
+      assert_prism_eval("[*1..2, 3, 4, *5..6, 7, 8, *9..11]")
+      assert_prism_eval("[0, *1..2, 3, 4, *5..6, 7, 8, *9..11]")
+      assert_prism_eval("[-1, true, 0, *1..2, 3, 4, *5..6, 7, 8, *9..11]")
+      assert_prism_eval("a = [1,2]; [0, *a, 3, 4, *5..6, 7, 8, *9..11]")
+      assert_prism_eval("[[*1..2], 3, *4..5]")
     end
 
     def test_AssocNode