AST-inline Array#[]=

Author: eregon

src/main/java/org/truffleruby/core/array/ArrayNodes.java

@@ -430,7 +430,7 @@ protected Object fallbackTernary(RubyArray array, Object start, Object length, O
         protected int normalize(int arraySize, int index,
                 ConditionProfile negativeDenormalizedIndex, BranchProfile negativeNormalizedIndex) {
             if (negativeDenormalizedIndex.profile(index < 0)) {
-                index = arraySize + index;
+                index += arraySize;
                 if (index < 0) {
                     negativeNormalizedIndex.enter();
                     throw new RaiseException(

src/main/java/org/truffleruby/core/inlined/CoreMethods.java

@@ -89,6 +89,7 @@ public class CoreMethods {
     public final InternalMethod STRING_EQUAL;
     public final InternalMethod SYMBOL_TO_PROC;
     public final InternalMethod ARRAY_INDEX_GET;
+    public final InternalMethod ARRAY_INDEX_SET;
 
     public CoreMethods(RubyContext context) {
         this.context = context;
@@ -153,6 +154,7 @@ public CoreMethods(RubyContext context) {
         STRING_EQUAL = getMethod(stringClass, "==");
         SYMBOL_TO_PROC = getMethod(symbolClass, "to_proc");
         ARRAY_INDEX_GET = getMethod(arrayClass, "[]");
+        ARRAY_INDEX_SET = getMethod(arrayClass, "[]=");
     }
 
     private Assumption registerAssumption(RubyClass klass, String methodName) {
@@ -249,6 +251,12 @@ public RubyNode createCallNode(RubyCallNodeParameters callParameters, Translator
                     return InlinedKindOfNodeGen.create(context, callParameters, self, args[0]);
                 default:
             }
+        } else if (n == 3) {
+            switch (callParameters.getMethodName()) {
+                case "[]=":
+                    return InlinedIndexSetNodeGen.create(context, callParameters, self, args[0], args[1]);
+                default:
+            }
         }
 
         return new RubyCallNode(callParameters);

src/main/java/org/truffleruby/core/inlined/InlinedIndexSetNode.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2020 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.core.inlined;
+
+import com.oracle.truffle.api.dsl.Bind;
+import org.truffleruby.RubyContext;
+import org.truffleruby.core.array.ArrayWriteNormalizedNode;
+import org.truffleruby.core.array.RubyArray;
+import org.truffleruby.language.dispatch.RubyCallNodeParameters;
+import org.truffleruby.language.methods.LookupMethodOnSelfNode;
+
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.profiles.ConditionProfile;
+
+public abstract class InlinedIndexSetNode extends TernaryInlinedOperationNode {
+
+    protected static final String METHOD = "[]=";
+
+    public InlinedIndexSetNode(RubyContext context, RubyCallNodeParameters callNodeParameters) {
+        super(context, callNodeParameters);
+    }
+
+    @Specialization(
+            guards = {
+                    "lookupNode.lookupProtected(frame, array, METHOD) == coreMethods().ARRAY_INDEX_SET",
+                    "normalizedIndex >= 0" },
+            assumptions = "assumptions",
+            limit = "1")
+    protected Object arrayWrite(VirtualFrame frame, RubyArray array, int index, Object value,
+            @Cached LookupMethodOnSelfNode lookupNode,
+            @Cached ConditionProfile denormalized,
+            @Bind("normalize(array, index, denormalized)") int normalizedIndex,
+            @Cached ArrayWriteNormalizedNode writeNode) {
+        return writeNode.executeWrite(array, normalizedIndex, value);
+    }
+
+    @Specialization
+    protected Object fallback(VirtualFrame frame, Object a, Object b, Object c) {
+        return rewriteAndCall(frame, a, b, c);
+    }
+
+    protected int normalize(RubyArray array, int index, ConditionProfile denormalized) {
+        if (denormalized.profile(index < 0)) {
+            index += array.size;
+        }
+        return index;
+    }
+
+}

src/main/java/org/truffleruby/core/inlined/TernaryInlinedOperationNode.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright (c) 2020 Oracle and/or its affiliates. All rights reserved. This
+ * code is released under a tri EPL/GPL/LGPL license. You can use it,
+ * redistribute it and/or modify it under the terms of the:
+ *
+ * Eclipse Public License version 2.0, or
+ * GNU General Public License version 2, or
+ * GNU Lesser General Public License version 2.1.
+ */
+package org.truffleruby.core.inlined;
+
+import org.truffleruby.RubyContext;
+import org.truffleruby.language.RubyNode;
+import org.truffleruby.language.dispatch.RubyCallNodeParameters;
+
+import com.oracle.truffle.api.Assumption;
+import com.oracle.truffle.api.dsl.NodeChild;
+
+@NodeChild(value = "receiver", type = RubyNode.class)
+@NodeChild(value = "operand1", type = RubyNode.class)
+@NodeChild(value = "operand2", type = RubyNode.class)
+public abstract class TernaryInlinedOperationNode extends InlinedOperationNode {
+
+    public TernaryInlinedOperationNode(
+            RubyContext context,
+            RubyCallNodeParameters callNodeParameters,
+            Assumption... assumptions) {
+        super(context, callNodeParameters, assumptions);
+    }
+
+    protected abstract RubyNode getReceiver();
+
+    protected abstract RubyNode getOperand1();
+
+    protected abstract RubyNode getOperand2();
+
+    @Override
+    protected RubyNode getReceiverNode() {
+        return getReceiver();
+    }
+
+    @Override
+    protected RubyNode[] getArgumentNodes() {
+        return new RubyNode[]{ getOperand1(), getOperand2() };
+    }
+
+}