[PRISM] Compile IndexAndWriteNode

2 files changed, 268 insertions, 224 deletions

diff --git a/prism_compile.c b/prism_compile.c
index b1e30f7416..9f9296900b 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -836,6 +836,245 @@ pm_compile_call_and_or_write_node(bool and_node, pm_node_t *receiver, pm_node_t
     return;
 }
 
+static int
+pm_setup_args(pm_arguments_node_t *arguments_node, int *flags, struct rb_callinfo_kwarg **kw_arg, rb_iseq_t *iseq, LINK_ANCHOR *const ret, const uint8_t *src, bool popped, pm_scope_node_t *scope_node, NODE dummy_line_node, pm_parser_t *parser)
+{
+    int orig_argc = 0;
+    if (arguments_node == NULL) {
+        if (*flags & VM_CALL_FCALL) {
+            *flags |= VM_CALL_VCALL;
+        }
+    }
+    else {
+        pm_node_list_t arguments_node_list = arguments_node->arguments;
+
+        bool has_keyword_splat = (arguments_node->base.flags & PM_ARGUMENTS_NODE_FLAGS_CONTAINS_KEYWORD_SPLAT);
+        bool has_splat = false;
+
+        // We count the number of elements post the splat node that are not keyword elements to
+        // eventually pass as an argument to newarray
+        int post_splat_counter = 0;
+
+        for (size_t index = 0; index < arguments_node_list.size; index++) {
+            pm_node_t *argument = arguments_node_list.nodes[index];
+
+            switch (PM_NODE_TYPE(argument)) {
+                // A keyword hash node contains all keyword arguments as AssocNodes and AssocSplatNodes
+              case PM_KEYWORD_HASH_NODE: {
+                  pm_keyword_hash_node_t *keyword_arg = (pm_keyword_hash_node_t *)argument;
+                  size_t len = keyword_arg->elements.size;
+
+                  if (has_keyword_splat) {
+                      int cur_hash_size = 0;
+                      orig_argc++;
+
+                      bool new_hash_emitted = false;
+                      for (size_t i = 0; i < len; i++) {
+                          pm_node_t *cur_node = keyword_arg->elements.nodes[i];
+
+                          pm_node_type_t cur_type = PM_NODE_TYPE(cur_node);
+
+                          switch (PM_NODE_TYPE(cur_node)) {
+                            case PM_ASSOC_NODE: {
+                                pm_assoc_node_t *assoc = (pm_assoc_node_t *)cur_node;
+
+                                PM_COMPILE_NOT_POPPED(assoc->key);
+                                PM_COMPILE_NOT_POPPED(assoc->value);
+                                cur_hash_size++;
+
+                                // If we're at the last keyword arg, or the last assoc node of this "set",
+                                // then we want to either construct a newhash or merge onto previous hashes
+                                if (i == (len - 1) || !PM_NODE_TYPE_P(keyword_arg->elements.nodes[i + 1], cur_type)) {
+                                    if (new_hash_emitted) {
+                                        ADD_SEND(ret, &dummy_line_node, id_core_hash_merge_ptr, INT2FIX(cur_hash_size * 2 + 1));
+                                    }
+                                    else {
+                                        ADD_INSN1(ret, &dummy_line_node, newhash, INT2FIX(cur_hash_size * 2));
+                                        cur_hash_size = 0;
+                                        new_hash_emitted = true;
+                                    }
+                                }
+
+                                break;
+                            }
+                            case PM_ASSOC_SPLAT_NODE: {
+                                if (len > 1) {
+                                    ADD_INSN1(ret, &dummy_line_node, putspecialobject, INT2FIX(VM_SPECIAL_OBJECT_VMCORE));
+                                    if (i == 0) {
+                                        ADD_INSN1(ret, &dummy_line_node, newhash, INT2FIX(0));
+                                        new_hash_emitted = true;
+                                    }
+                                    else {
+                                        PM_SWAP;
+                                    }
+                                }
+
+                                pm_assoc_splat_node_t *assoc_splat = (pm_assoc_splat_node_t *)cur_node;
+                                PM_COMPILE_NOT_POPPED(assoc_splat->value);
+
+                                *flags |= VM_CALL_KW_SPLAT | VM_CALL_KW_SPLAT_MUT;
+
+                                if (len > 1) {
+                                    ADD_SEND(ret, &dummy_line_node, id_core_hash_merge_kwd, INT2FIX(2));
+                                }
+
+                                if ((i < len - 1) && !PM_NODE_TYPE_P(keyword_arg->elements.nodes[i + 1], cur_type)) {
+                                    ADD_INSN1(ret, &dummy_line_node, putspecialobject, INT2FIX(VM_SPECIAL_OBJECT_VMCORE));
+                                    PM_SWAP;
+                                }
+
+                                cur_hash_size = 0;
+                                break;
+                            }
+                            default: {
+                                rb_bug("Unknown type");
+                            }
+                          }
+                      }
+                      break;
+                  }
+                  else {
+                      *kw_arg = rb_xmalloc_mul_add(len, sizeof(VALUE), sizeof(struct rb_callinfo_kwarg));
+                      *flags = VM_CALL_KWARG;
+                      (*kw_arg)->keyword_len = (int) len;
+
+                      // TODO: Method callers like `foo(a => b)`
+                      for (size_t i = 0; i < len; i++) {
+                          pm_assoc_node_t *assoc = (pm_assoc_node_t *)keyword_arg->elements.nodes[i];
+                          (*kw_arg)->keywords[i] = pm_static_literal_value(assoc->key, scope_node, parser);
+                          PM_COMPILE_NOT_POPPED(assoc->value);
+                      }
+                  }
+                  break;
+              }
+              case PM_SPLAT_NODE: {
+                  *flags |= VM_CALL_ARGS_SPLAT;
+                  pm_splat_node_t *splat_node = (pm_splat_node_t *)argument;
+                  if (splat_node->expression) {
+                      orig_argc++;
+                      PM_COMPILE_NOT_POPPED(splat_node->expression);
+                  }
+
+                  ADD_INSN1(ret, &dummy_line_node, splatarray, popped ? Qfalse : Qtrue);
+
+                  has_splat = true;
+                  post_splat_counter = 0;
+
+                  break;
+              }
+              case PM_FORWARDING_ARGUMENTS_NODE: {
+                  orig_argc++;
+                  *flags |= VM_CALL_ARGS_BLOCKARG | VM_CALL_ARGS_SPLAT;
+                  ADD_GETLOCAL(ret, &dummy_line_node, 3, 0);
+                  ADD_INSN1(ret, &dummy_line_node, splatarray, RBOOL(arguments_node_list.size > 1));
+                  ADD_INSN2(ret, &dummy_line_node, getblockparamproxy, INT2FIX(4), INT2FIX(0));
+                  break;
+              }
+              default: {
+                  orig_argc++;
+                  post_splat_counter++;
+                  PM_COMPILE_NOT_POPPED(argument);
+
+                  if (has_splat) {
+                      // If the next node starts the keyword section of parameters
+                      if ((index < arguments_node_list.size - 1) && PM_NODE_TYPE_P(arguments_node_list.nodes[index + 1], PM_KEYWORD_HASH_NODE)) {
+
+                          ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(post_splat_counter));
+                          ADD_INSN1(ret, &dummy_line_node, splatarray, Qfalse);
+                          ADD_INSN(ret, &dummy_line_node, concatarray);
+                      }
+                      // If it's the final node
+                      else if (index == arguments_node_list.size - 1) {
+                          ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(post_splat_counter));
+                          ADD_INSN(ret, &dummy_line_node, concatarray);
+                      }
+                  }
+              }
+            }
+        }
+    }
+    return orig_argc;
+}
+
+static void
+pm_compile_index_and_or_write_node(bool and_node, pm_node_t *receiver, pm_node_t *value, pm_arguments_node_t *arguments, pm_node_t *block, LINK_ANCHOR *const ret, rb_iseq_t *iseq, int lineno, const uint8_t * src, bool popped, pm_scope_node_t *scope_node, pm_parser_t *parser)
+{
+    NODE dummy_line_node = generate_dummy_line_node(lineno, lineno);
+    PM_PUTNIL_UNLESS_POPPED;
+
+    PM_COMPILE_NOT_POPPED(receiver);
+
+    int flag = 0;
+    struct rb_callinfo_kwarg *keywords = NULL;
+    int argc_int = 0;
+
+    if (arguments) {
+        argc_int = pm_setup_args(arguments, &flag, &keywords, iseq, ret, src, popped, scope_node, dummy_line_node, parser);
+    }
+
+    VALUE argc = INT2FIX(argc_int);
+    int boff = 0;
+
+    if (block) {
+        PM_COMPILE_NOT_POPPED(block);
+        flag |= VM_CALL_ARGS_BLOCKARG;
+        boff = 1;
+    }
+
+    ADD_INSN1(ret, &dummy_line_node, dupn, FIXNUM_INC(argc, 1 + boff));
+
+    ADD_SEND_WITH_FLAG(ret, &dummy_line_node, idAREF, argc, INT2FIX(flag));
+
+    LABEL *label = NEW_LABEL(lineno);
+    LABEL *lfin = NEW_LABEL(lineno);
+
+    PM_DUP;
+
+    if (and_node) {
+        ADD_INSNL(ret, &dummy_line_node, branchunless, label);
+    }
+    else {
+        // ornode
+        ADD_INSNL(ret, &dummy_line_node, branchif, label);
+    }
+
+    PM_POP;
+
+    PM_COMPILE_NOT_POPPED(value);
+
+    if (!popped) {
+        ADD_INSN1(ret, &dummy_line_node, setn, FIXNUM_INC(argc, 2 + boff));
+    }
+    if (flag & VM_CALL_ARGS_SPLAT) {
+        ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(1));
+        if (boff > 0) {
+            ADD_INSN1(ret, &dummy_line_node, dupn, INT2FIX(3));
+            ADD_INSN(ret, &dummy_line_node, swap);
+            ADD_INSN(ret, &dummy_line_node, pop);
+        }
+        ADD_INSN(ret, &dummy_line_node, concatarray);
+        if (boff > 0) {
+            ADD_INSN1(ret, &dummy_line_node, setn, INT2FIX(3));
+            PM_POP;
+        }
+        ADD_SEND_WITH_FLAG(ret, &dummy_line_node, idASET, argc, INT2FIX(flag));
+    }
+    else {
+        if (boff > 0)
+            ADD_INSN(ret, &dummy_line_node, swap);
+        ADD_SEND_WITH_FLAG(ret, &dummy_line_node, idASET, FIXNUM_INC(argc, 1), INT2FIX(flag));
+    }
+    PM_POP;
+    ADD_INSNL(ret, &dummy_line_node, jump, lfin);
+    ADD_LABEL(ret, label);
+    if (!popped) {
+        ADD_INSN1(ret, &dummy_line_node, setn, FIXNUM_INC(argc, 2 + boff));
+    }
+    ADD_INSN1(ret, &dummy_line_node, adjuststack, FIXNUM_INC(argc, 2 + boff));
+    ADD_LABEL(ret, lfin);
+    return;
+}
+
 /**
  * In order to properly compile multiple-assignment, some preprocessing needs to
  * be performed in the case of call or constant path targets. This is when they
@@ -1312,166 +1551,6 @@ pm_compile_defined_expr(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *con
     ADD_LABEL(ret, lfinish[0]);
 }
 
-static int
-pm_setup_args(pm_arguments_node_t *arguments_node, int *flags, struct rb_callinfo_kwarg **kw_arg, rb_iseq_t *iseq, LINK_ANCHOR *const ret, const uint8_t *src, bool popped, pm_scope_node_t *scope_node, NODE dummy_line_node, pm_parser_t *parser)
-{
-    int orig_argc = 0;
-    if (arguments_node == NULL) {
-        if (*flags & VM_CALL_FCALL) {
-            *flags |= VM_CALL_VCALL;
-        }
-    }
-    else {
-        pm_node_list_t arguments_node_list = arguments_node->arguments;
-
-        bool has_keyword_splat = (arguments_node->base.flags & PM_ARGUMENTS_NODE_FLAGS_CONTAINS_KEYWORD_SPLAT);
-        bool has_splat = false;
-
-        // We count the number of elements post the splat node that are not keyword elements to
-        // eventually pass as an argument to newarray
-        int post_splat_counter = 0;
-
-        for (size_t index = 0; index < arguments_node_list.size; index++) {
-            pm_node_t *argument = arguments_node_list.nodes[index];
-
-            switch (PM_NODE_TYPE(argument)) {
-                // A keyword hash node contains all keyword arguments as AssocNodes and AssocSplatNodes
-              case PM_KEYWORD_HASH_NODE: {
-                  pm_keyword_hash_node_t *keyword_arg = (pm_keyword_hash_node_t *)argument;
-                  size_t len = keyword_arg->elements.size;
-
-                  if (has_keyword_splat) {
-                      int cur_hash_size = 0;
-                      orig_argc++;
-
-                      bool new_hash_emitted = false;
-                      for (size_t i = 0; i < len; i++) {
-                          pm_node_t *cur_node = keyword_arg->elements.nodes[i];
-
-                          pm_node_type_t cur_type = PM_NODE_TYPE(cur_node);
-
-                          switch (PM_NODE_TYPE(cur_node)) {
-                            case PM_ASSOC_NODE: {
-                                pm_assoc_node_t *assoc = (pm_assoc_node_t *)cur_node;
-
-                                PM_COMPILE_NOT_POPPED(assoc->key);
-                                PM_COMPILE_NOT_POPPED(assoc->value);
-                                cur_hash_size++;
-
-                                // If we're at the last keyword arg, or the last assoc node of this "set",
-                                // then we want to either construct a newhash or merge onto previous hashes
-                                if (i == (len - 1) || !PM_NODE_TYPE_P(keyword_arg->elements.nodes[i + 1], cur_type)) {
-                                    if (new_hash_emitted) {
-                                        ADD_SEND(ret, &dummy_line_node, id_core_hash_merge_ptr, INT2FIX(cur_hash_size * 2 + 1));
-                                    }
-                                    else {
-                                        ADD_INSN1(ret, &dummy_line_node, newhash, INT2FIX(cur_hash_size * 2));
-                                        cur_hash_size = 0;
-                                        new_hash_emitted = true;
-                                    }
-                                }
-
-                                break;
-                            }
-                            case PM_ASSOC_SPLAT_NODE: {
-                                if (len > 1) {
-                                    ADD_INSN1(ret, &dummy_line_node, putspecialobject, INT2FIX(VM_SPECIAL_OBJECT_VMCORE));
-                                    if (i == 0) {
-                                        ADD_INSN1(ret, &dummy_line_node, newhash, INT2FIX(0));
-                                        new_hash_emitted = true;
-                                    }
-                                    else {
-                                        PM_SWAP;
-                                    }
-                                }
-
-                                pm_assoc_splat_node_t *assoc_splat = (pm_assoc_splat_node_t *)cur_node;
-                                PM_COMPILE_NOT_POPPED(assoc_splat->value);
-
-                                *flags |= VM_CALL_KW_SPLAT | VM_CALL_KW_SPLAT_MUT;
-
-                                if (len > 1) {
-                                    ADD_SEND(ret, &dummy_line_node, id_core_hash_merge_kwd, INT2FIX(2));
-                                }
-
-                                if ((i < len - 1) && !PM_NODE_TYPE_P(keyword_arg->elements.nodes[i + 1], cur_type)) {
-                                    ADD_INSN1(ret, &dummy_line_node, putspecialobject, INT2FIX(VM_SPECIAL_OBJECT_VMCORE));
-                                    PM_SWAP;
-                                }
-
-                                cur_hash_size = 0;
-                                break;
-                            }
-                            default: {
-                                rb_bug("Unknown type");
-                            }
-                          }
-                      }
-                      break;
-                  }
-                  else {
-                      *kw_arg = rb_xmalloc_mul_add(len, sizeof(VALUE), sizeof(struct rb_callinfo_kwarg));
-                      *flags = VM_CALL_KWARG;
-                      (*kw_arg)->keyword_len = (int) len;
-
-                      // TODO: Method callers like `foo(a => b)`
-                      for (size_t i = 0; i < len; i++) {
-                          pm_assoc_node_t *assoc = (pm_assoc_node_t *)keyword_arg->elements.nodes[i];
-                          (*kw_arg)->keywords[i] = pm_static_literal_value(assoc->key, scope_node, parser);
-                          PM_COMPILE_NOT_POPPED(assoc->value);
-                      }
-                  }
-                  break;
-              }
-              case PM_SPLAT_NODE: {
-                  *flags |= VM_CALL_ARGS_SPLAT;
-                  pm_splat_node_t *splat_node = (pm_splat_node_t *)argument;
-                  if (splat_node->expression) {
-                      orig_argc++;
-                      PM_COMPILE_NOT_POPPED(splat_node->expression);
-                  }
-
-                  ADD_INSN1(ret, &dummy_line_node, splatarray, popped ? Qfalse : Qtrue);
-
-                  has_splat = true;
-                  post_splat_counter = 0;
-
-                  break;
-              }
-              case PM_FORWARDING_ARGUMENTS_NODE: {
-                  orig_argc++;
-                  *flags |= VM_CALL_ARGS_BLOCKARG | VM_CALL_ARGS_SPLAT;
-                  ADD_GETLOCAL(ret, &dummy_line_node, 3, 0);
-                  ADD_INSN1(ret, &dummy_line_node, splatarray, RBOOL(arguments_node_list.size > 1));
-                  ADD_INSN2(ret, &dummy_line_node, getblockparamproxy, INT2FIX(4), INT2FIX(0));
-                  break;
-              }
-              default: {
-                  orig_argc++;
-                  post_splat_counter++;
-                  PM_COMPILE_NOT_POPPED(argument);
-
-                  if (has_splat) {
-                      // If the next node starts the keyword section of parameters
-                      if ((index < arguments_node_list.size - 1) && PM_NODE_TYPE_P(arguments_node_list.nodes[index + 1], PM_KEYWORD_HASH_NODE)) {
-
-                          ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(post_splat_counter));
-                          ADD_INSN1(ret, &dummy_line_node, splatarray, Qfalse);
-                          ADD_INSN(ret, &dummy_line_node, concatarray);
-                      }
-                      // If it's the final node
-                      else if (index == arguments_node_list.size - 1) {
-                          ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(post_splat_counter));
-                          ADD_INSN(ret, &dummy_line_node, concatarray);
-                      }
-                  }
-              }
-            }
-        }
-    }
-    return orig_argc;
-}
-
 /*
  * Compiles a prism node into instruction sequences
  *
@@ -2652,73 +2731,16 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         PM_COMPILE(cast->value);
         return;
       }
+      case PM_INDEX_AND_WRITE_NODE: {
+        pm_index_and_write_node_t *index_and_write_node = (pm_index_and_write_node_t *)node;
+
+        pm_compile_index_and_or_write_node(true, index_and_write_node->receiver, index_and_write_node->value, index_and_write_node->arguments, index_and_write_node->block, ret, iseq, lineno, src, popped, scope_node, parser);
+        return;
+      }
       case PM_INDEX_OR_WRITE_NODE: {
         pm_index_or_write_node_t *index_or_write_node = (pm_index_or_write_node_t *)node;
 
-        PM_PUTNIL_UNLESS_POPPED;
-
-        PM_COMPILE_NOT_POPPED(index_or_write_node->receiver);
-
-        int flag = 0;
-        struct rb_callinfo_kwarg *keywords = NULL;
-        int argc_int = 0;
-
-        if (index_or_write_node->arguments) {
-            argc_int = pm_setup_args(index_or_write_node->arguments, &flag, &keywords, iseq, ret, src, popped, scope_node, dummy_line_node, parser);
-        }
-
-        VALUE argc = INT2FIX(argc_int);
-        int boff = 0;
-
-        if (index_or_write_node->block) {
-            PM_COMPILE_NOT_POPPED(index_or_write_node->block);
-            flag |= VM_CALL_ARGS_BLOCKARG;
-            boff = 1;
-        }
-
-        ADD_INSN1(ret, &dummy_line_node, dupn, FIXNUM_INC(argc, 1 + boff));
-
-        ADD_SEND_WITH_FLAG(ret, &dummy_line_node, idAREF, argc, INT2FIX(flag));
-
-        LABEL *label = NEW_LABEL(lineno);
-        LABEL *lfin = NEW_LABEL(lineno);
-
-        PM_DUP;
-        ADD_INSNL(ret, &dummy_line_node, branchif, label);
-        PM_POP;
-
-        PM_COMPILE_NOT_POPPED(index_or_write_node->value);
-
-        if (!popped) {
-            ADD_INSN1(ret, &dummy_line_node, setn, FIXNUM_INC(argc, 2 + boff));
-        }
-        if (flag & VM_CALL_ARGS_SPLAT) {
-            ADD_INSN1(ret, &dummy_line_node, newarray, INT2FIX(1));
-            if (boff > 0) {
-                ADD_INSN1(ret, &dummy_line_node, dupn, INT2FIX(3));
-                ADD_INSN(ret, &dummy_line_node, swap);
-                ADD_INSN(ret, &dummy_line_node, pop);
-            }
-            ADD_INSN(ret, &dummy_line_node, concatarray);
-            if (boff > 0) {
-                ADD_INSN1(ret, &dummy_line_node, setn, INT2FIX(3));
-                PM_POP;
-            }
-            ADD_SEND_WITH_FLAG(ret, &dummy_line_node, idASET, argc, INT2FIX(flag));
-        }
-        else {
-            if (boff > 0)
-                ADD_INSN(ret, &dummy_line_node, swap);
-            ADD_SEND_WITH_FLAG(ret, &dummy_line_node, idASET, FIXNUM_INC(argc, 1), INT2FIX(flag));
-        }
-        PM_POP;
-        ADD_INSNL(ret, &dummy_line_node, jump, lfin);
-        ADD_LABEL(ret, label);
-        if (!popped) {
-            ADD_INSN1(ret, &dummy_line_node, setn, FIXNUM_INC(argc, 2 + boff));
-        }
-        ADD_INSN1(ret, &dummy_line_node, adjuststack, FIXNUM_INC(argc, 2 + boff));
-        ADD_LABEL(ret, lfin);
+        pm_compile_index_and_or_write_node(false, index_or_write_node->receiver, index_or_write_node->value, index_or_write_node->arguments, index_or_write_node->block, ret, iseq, lineno, src, popped, scope_node, parser);
         return;
       }
       case PM_INSTANCE_VARIABLE_AND_WRITE_NODE: {
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index 96aad1ab5b..2ee078673f 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -238,6 +238,28 @@ module Prism
       assert_prism_eval("$pit = 1")
     end
 
+    def test_IndexAndWriteNode
+      assert_prism_eval("[0][0] &&= 1")
+      assert_prism_eval("[nil][0] &&= 1")
+
+      # Testing `[]` with a block passed in
+      assert_prism_eval(<<-CODE)
+        class CustomHash < Hash
+          def []=(key, value, &block)
+            block ? super(block.call(key), value) : super(key, value)
+          end
+        end
+
+        hash = CustomHash.new
+
+        # Call the custom method with a block that modifies
+        # the key before assignment
+        hash["KEY"] = "test"
+        hash["key", &(Proc.new { _1.upcase })] &&= "value"
+        hash
+      CODE
+    end
+
     def test_IndexOrWriteNode
       assert_prism_eval("[0][0] ||= 1")
       assert_prism_eval("[nil][0] ||= 1")