[PRISM] Compile IndexOrWriteNode

2 files changed, 90 insertions, 0 deletions

diff --git a/prism_compile.c b/prism_compile.c
index fa5a7aafca..b1e30f7416 100644
--- a/prism_compile.c
+++ b/prism_compile.c
@@ -2652,6 +2652,75 @@ pm_compile_node(rb_iseq_t *iseq, const pm_node_t *node, LINK_ANCHOR *const ret,
         PM_COMPILE(cast->value);
         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);
+        return;
+      }
       case PM_INSTANCE_VARIABLE_AND_WRITE_NODE: {
         pm_instance_variable_and_write_node_t *instance_variable_and_write_node = (pm_instance_variable_and_write_node_t*) node;
 
diff --git a/test/ruby/test_compile_prism.rb b/test/ruby/test_compile_prism.rb
index a01666ee54..96aad1ab5b 100644
--- a/test/ruby/test_compile_prism.rb
+++ b/test/ruby/test_compile_prism.rb
@@ -238,6 +238,27 @@ module Prism
       assert_prism_eval("$pit = 1")
     end
 
+    def test_IndexOrWriteNode
+      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)
+            super(block.call(key), value)
+          end
+        end
+
+        hash = CustomHash.new
+
+        # Call the custom method with a block that modifies
+        # the key before assignment
+        hash["key", &(Proc.new { _1.upcase })] ||= "value"
+        hash
+      CODE
+    end
+
     def test_InstanceVariableAndWriteNode
       assert_prism_eval("@pit = 0; @pit &&= 1")
     end