[GR-17457] Avoid executing arg nodes multiple times with op element assignment.

PullRequest: truffleruby/3372

Author: aardvark179

CHANGELOG.md

@@ -47,6 +47,7 @@ Bug fixes:
 * Fix `/#{...}/o` to evaluate only once per context when splitting happens (@eregon).
 * Fix `Kernel#sprintf` formatting of floats to be like CRuby (@aardvark179).
 * Fix `Process.egid=` to accept `String`s (#2615, @ngtban)
+* Fix optional assignment to only evaluate index arguments once (#2658, @aardvark179).
 
 Compatibility:
 

spec/ruby/language/optional_assignments_spec.rb

@@ -300,6 +300,44 @@ def []=(k, v)
         (@b[:k] ||= 12).should == 12
       end
 
+      it 'correctly handles a splatted argument for the index' do
+        (@b[*[:k]] ||= 12).should == 12
+      end
+
+      it "evaluates the index precisely once" do
+        ary = [:x, :y]
+        @a[:x] = 15
+        @a[ary.pop] ||= 25
+        ary.should == [:x]
+        @a.should == { x: 15, y: 25 }
+      end
+
+      it "evaluates the index arguments in the correct order" do
+        ary = Class.new(Array) do
+          def [](x, y)
+            super(x + 3 * y)
+          end
+
+          def []=(x, y, value)
+            super(x + 3 * y, value)
+          end
+        end.new
+        ary[0, 0] = 1
+        ary[1, 0] = 1
+        ary[2, 0] = nil
+        ary[3, 0] = 1
+        ary[4, 0] = 1
+        ary[5, 0] = 1
+        ary[6, 0] = nil
+
+        foo = [0, 2]
+
+        ary[foo.pop, foo.pop] ||= 2
+
+        ary[2, 0].should == 2
+        ary[6, 0].should == nil
+      end
+
       it 'returns the assigned value, not the result of the []= method with +=' do
         @b[:k] = 17
         (@b[:k] += 12).should == 29

src/main/java/org/truffleruby/parser/BodyTranslator.java

@@ -2474,58 +2474,98 @@ public RubyNode visitOpElementAsgnNode(OpElementAsgnParseNode node) {
                 tempName,
                 0,
                 node.getReceiverNode());
+        final ArrayList<ValueFromNode> argValues = argsToTemp(node);
 
         final String op = node.getOperatorName();
         final boolean logicalOperation = op.equals("&&") || op.equals("||");
 
         if (logicalOperation) {
-            final ParseNode write = write(node, readReceiverFromTemp, value);
-            final ParseNode operation = operation(node, readReceiverFromTemp, op, write);
+            final ParseNode write = write(node, readReceiverFromTemp, argValues, value);
+            final ParseNode operation = operation(node, readReceiverFromTemp, argValues, op, write);
 
-            return block(node, writeReceiverToTemp, operation);
+            return block(node, writeReceiverToTemp, argValues, operation);
         } else {
-            final ParseNode operation = operation(node, readReceiverFromTemp, op, value);
-            final ParseNode write = write(node, readReceiverFromTemp, operation);
+            final ParseNode operation = operation(node, readReceiverFromTemp, argValues, op, value);
+            final ParseNode write = write(node, readReceiverFromTemp, argValues, operation);
 
-            return block(node, writeReceiverToTemp, write);
+            return block(node, writeReceiverToTemp, argValues, write);
         }
     }
 
-    private RubyNode block(OpElementAsgnParseNode node, ParseNode writeReceiverToTemp, ParseNode main) {
+    private RubyNode block(OpElementAsgnParseNode node, ParseNode writeReceiverToTemp,
+            ArrayList<ValueFromNode> argValues, ParseNode main) {
         final BlockParseNode block = new BlockParseNode(node.getPosition());
         block.add(writeReceiverToTemp);
         block.add(main);
 
-        final RubyNode ret = block.accept(this);
+        /* prepareAndThen is going to take an argument, and the action that comes after it, and return a node that does
+         * both of those things. We start off with ret being the block (our final action) and so the first node we
+         * should produce is one that evaluates the last argument, and then the block. The final value of ret should be
+         * a node that evaluates the first argument, and then any other arguments, and then the block. So, we must go
+         * through the argument list in reverse order. */
+        RubyNode ret = block.accept(this);
+        var listIterator = argValues.listIterator(argValues.size());
+        while (listIterator.hasPrevious()) {
+            ret = listIterator.previous().prepareAndThen(node.getPosition(), ret);
+        }
         return addNewlineIfNeeded(node, ret);
     }
 
-    private ParseNode write(OpElementAsgnParseNode node, ParseNode readReceiverFromTemp, ParseNode value) {
-        final ParseNode readArguments = node.getArgsNode();
+    private ParseNode write(OpElementAsgnParseNode node, ParseNode readReceiverFromTemp,
+            ArrayList<ValueFromNode> argValues, ParseNode value) {
         final ParseNode writeArguments;
         // Like ParserSupport#arg_add, but copy the first node
-        if (readArguments instanceof ArrayParseNode) {
+        if (node.getArgsNode() instanceof ArrayParseNode) {
             final ArrayParseNode readArgsCopy = new ArrayParseNode(node.getPosition());
-            readArgsCopy.addAll((ArrayParseNode) readArguments).add(value);
+            for (var arg : argValues) {
+                readArgsCopy.add(arg.get(node.getPosition()));
+            }
+            readArgsCopy.add(value);
             writeArguments = readArgsCopy;
         } else {
-            writeArguments = new ArgsPushParseNode(node.getPosition(), readArguments, value);
+            writeArguments = new ArgsPushParseNode(node.getPosition(), argValues.get(0).get(node.getPosition()), value);
         }
 
         return new AttrAssignParseNode(node.getPosition(), readReceiverFromTemp, "[]=", writeArguments, false);
     }
 
+    private ArrayList<ValueFromNode> argsToTemp(OpElementAsgnParseNode node) {
+        ArrayList<ValueFromNode> argValues = new ArrayList<>();
+
+        final ParseNode readArguments = node.getArgsNode();
+        if (readArguments instanceof ArrayParseNode) {
+            for (ParseNode child : ((ArrayParseNode) readArguments).children()) {
+                argValues.add(ValueFromNode.valueFromNode(this, child));
+            }
+        } else {
+            argValues.add(ValueFromNode.valueFromNode(this, readArguments));
+        }
+
+        return argValues;
+    }
+
     private ParseNode operation(
             OpElementAsgnParseNode node,
             ParseNode readReceiverFromTemp,
+            ArrayList<ValueFromNode> argValues,
             String op,
             ParseNode right) {
+        ParseNode readArguments;
+        if (node.getArgsNode() instanceof ArrayParseNode) {
+            final ArrayParseNode readArgsArray = new ArrayParseNode(node.getPosition());
+            for (var arg : argValues) {
+                readArgsArray.add(arg.get(node.getPosition()));
+            }
+            readArguments = readArgsArray;
+        } else {
+            readArguments = argValues.get(0).get(node.getPosition());
+        }
 
         final ParseNode read = new CallParseNode(
                 node.getPosition(),
                 readReceiverFromTemp,
                 "[]",
-                node.getArgsNode(),
+                readArguments,
                 null);
         ParseNode operation;
         switch (op) {

src/main/java/org/truffleruby/parser/ValueFromNode.java

@@ -15,6 +15,7 @@
 import org.truffleruby.parser.ast.LocalVarParseNode;
 import org.truffleruby.parser.ast.ParseNode;
 import org.truffleruby.parser.ast.SelfParseNode;
+import org.truffleruby.parser.ast.SplatParseNode;
 
 import java.util.Arrays;
 
@@ -80,9 +81,30 @@ public ParseNode get(SourceIndexLength sourceSection) {
 
     }
 
+    class ValueFromSplatNode implements ValueFromNode {
+
+        private final ValueFromNode value;
+
+        public ValueFromSplatNode(BodyTranslator translator, SplatParseNode node) {
+            value = valueFromNode(translator, node.getValue());
+        }
+
+        @Override
+        public RubyNode prepareAndThen(SourceIndexLength sourceSection, RubyNode subsequent) {
+            return value.prepareAndThen(sourceSection, subsequent);
+        }
+
+        @Override
+        public ParseNode get(SourceIndexLength sourceSection) {
+            return new SplatParseNode(sourceSection, value.get(sourceSection));
+        }
+    }
+
     static ValueFromNode valueFromNode(BodyTranslator translator, ParseNode node) {
         if (node instanceof SelfParseNode) {
             return new ValueFromSelfNode();
+        } else if (node instanceof SplatParseNode) {
+            return new ValueFromSplatNode(translator, (SplatParseNode) node);
         } else {
             return new ValueFromEffectNode(translator, node);
         }