thread safe and simpler PostProcessArgumentsNode and ArgumentStatePush

com.oracle.truffle.r.nodes/src/com/oracle/truffle/r/nodes/function/ArgumentStatePush.java

@@ -24,25 +24,33 @@ package com.oracle.truffle.r.nodes.function;
 
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
+import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.dsl.TypeSystemReference;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.profiles.ConditionProfile;
+import com.oracle.truffle.r.nodes.EmptyTypeSystemFlatLayout;
 import com.oracle.truffle.r.nodes.access.WriteLocalFrameVariableNode;
 import com.oracle.truffle.r.runtime.FastROptions;
 import com.oracle.truffle.r.runtime.RArguments;
+import com.oracle.truffle.r.runtime.RInternalError;
 import com.oracle.truffle.r.runtime.data.RFunction;
 import com.oracle.truffle.r.runtime.data.RLanguage;
 import com.oracle.truffle.r.runtime.data.RShareable;
+import com.oracle.truffle.r.runtime.data.RSharingAttributeStorage;
 import com.oracle.truffle.r.runtime.data.model.RAbstractContainer;
 
 /**
  * A {@link ArgumentStatePush} is used to bump up state transition for function arguments.
  */
+@TypeSystemReference(EmptyTypeSystemFlatLayout.class)
 public abstract class ArgumentStatePush extends Node {
 
     public abstract void executeObject(VirtualFrame frame, Object shareable);
 
+    @Child private WriteLocalFrameVariableNode write;
+
     private final ConditionProfile isRefCountUpdateable = ConditionProfile.createBinaryProfile();
 
     private final int index;
@@ -57,76 +65,59 @@ public abstract class ArgumentStatePush extends Node {
         this.index = index;
     }
 
-    public boolean refCounted() {
-        return mask > 0;
+    protected int createWriteArgMask(VirtualFrame frame, RShareable shareable) {
+        if (shareable instanceof RAbstractContainer) {
+            if (shareable instanceof RLanguage || ((RAbstractContainer) shareable).getLength() < REF_COUNT_SIZE_THRESHOLD) {
+                // don't decrement ref count for small objects or language objects- this
+                // is pretty conservative and can be further finessed
+                return -1;
+            }
+        }
+        RFunction fun = RArguments.getFunction(frame);
+        if (fun == null) {
+            return -1;
+        }
+        Object root = fun.getRootNode();
+        if (!(root instanceof FunctionDefinitionNode)) {
+            // root is RBuiltinRootNode
+            return -1;
+        }
+        FunctionDefinitionNode fdn = (FunctionDefinitionNode) root;
+        PostProcessArgumentsNode postProcessNode = fdn.getArgPostProcess();
+        if (postProcessNode == null) {
+            // arguments to this function are not to be reference counted
+            return -1;
+        }
+        if (!postProcessNode.updateBits(index)) {
+            // this will fail if the index is too big
+            return -1;
+        }
+        return 1 << index;
     }
 
     @Specialization
-    public void transitionState(VirtualFrame frame, RShareable shareable) {
+    public void transitionState(VirtualFrame frame, RSharingAttributeStorage shareable,
+                    @Cached("createWriteArgMask(frame, shareable)") int writeArgMask) {
         if (isRefCountUpdateable.profile(!shareable.isSharedPermanent())) {
             shareable.incRefCount();
-            if (!FastROptions.RefCountIncrementOnly.getBooleanValue()) {
-                if (mask == 0) {
+            if (writeArgMask != -1 && !FastROptions.RefCountIncrementOnly.getBooleanValue()) {
+                if (write == null) {
                     CompilerDirectives.transferToInterpreterAndInvalidate();
-                    if (shareable instanceof RAbstractContainer) {
-                        if (shareable instanceof RLanguage || ((RAbstractContainer) shareable).getLength() < REF_COUNT_SIZE_THRESHOLD) {
-                            // don't decrement ref count for small objects or language objects- this
-                            // is
-                            // pretty conservative and can be further finessed
-                            mask = -1;
-                            return;
-                        }
-                    }
-                    RFunction fun = RArguments.getFunction(frame);
-                    if (fun == null) {
-                        mask = -1;
-                        return;
-                    }
-                    Object root = fun.getRootNode();
-                    if (!(root instanceof FunctionDefinitionNode)) {
-                        // root is RBuiltinRootNode
-                        mask = -1;
-                        return;
-                    }
-                    FunctionDefinitionNode fdn = (FunctionDefinitionNode) root;
-                    PostProcessArgumentsNode postProcessNode = fdn.getArgPostProcess();
-                    if (postProcessNode == null) {
-                        // arguments to this function are not to be reference counted
-                        mask = -1;
-                        return;
-                    }
-                    // this is needed for when FunctionDefinitionNode is split by the Truffle
-                    // runtime
-                    postProcessNode = postProcessNode.getActualNode();
-                    if (index >= Math.min(postProcessNode.getLength(), MAX_COUNTED_ARGS)) {
-                        mask = -1;
-                        return;
-                    }
-                    mask = 1 << index;
-                    int transArgsBitSet = postProcessNode.transArgsBitSet;
-                    postProcessNode.transArgsBitSet = transArgsBitSet | mask;
-                    writeArgNode = insert(WriteLocalFrameVariableNode.createForRefCount(Integer.valueOf(mask)));
-                }
-                if (mask != -1) {
-                    writeArgNode.execute(frame, shareable);
+                    write = insert(WriteLocalFrameVariableNode.createForRefCount(Integer.valueOf(writeArgMask)));
                 }
+                write.execute(frame, shareable);
             }
         }
     }
 
+    @Specialization
+    public void transitionState(RShareable shareable) {
+        throw RInternalError.shouldNotReachHere("unexpected RShareable that is not RSharingAttributeStorage: " + shareable.getClass());
+    }
+
     @Specialization(guards = "!isShareable(o)")
-    public void transitionStateNonShareable(VirtualFrame frame, @SuppressWarnings("unused") Object o) {
-        if (mask > 0) {
-            // this argument used to be reference counted but is no longer
-            CompilerDirectives.transferToInterpreterAndInvalidate();
-            RFunction fun = RArguments.getFunction(frame);
-            assert fun != null;
-            FunctionDefinitionNode fdn = (FunctionDefinitionNode) fun.getRootNode();
-            assert fdn != null;
-            int transArgsBitSet = fdn.getArgPostProcess().transArgsBitSet;
-            fdn.getArgPostProcess().transArgsBitSet = transArgsBitSet & (~mask);
-            mask = -1;
-        }
+    public void transitionStateNonShareable(@SuppressWarnings("unused") Object o) {
+        // do nothing
     }
 
     protected boolean isShareable(Object o) {

com.oracle.truffle.r.nodes/src/com/oracle/truffle/r/nodes/function/PostProcessArgumentsNode.java

@@ -22,13 +22,11 @@
  */
 package com.oracle.truffle.r.nodes.function;
 
-import com.oracle.truffle.api.CompilerAsserts;
 import com.oracle.truffle.api.CompilerDirectives;
-import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
-import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.profiles.ConditionProfile;
+import com.oracle.truffle.api.utilities.AssumedValue;
 import com.oracle.truffle.r.nodes.access.variables.LocalReadVariableNode;
 import com.oracle.truffle.r.runtime.data.RNull;
 import com.oracle.truffle.r.runtime.data.RShareable;
@@ -41,28 +39,20 @@ import com.oracle.truffle.r.runtime.nodes.RNode;
 public final class PostProcessArgumentsNode extends RNode {
 
     @Children private final LocalReadVariableNode[] sequence;
-    @Child private PostProcessArgumentsNode nextOptPostProccessArgNode;
-    @CompilationFinal int transArgsBitSet;
 
-    // the first time this node is cloned (via FunctionDefinitionNode) it's from the trufflerizer -
-    // in this case we should not create the node chain chain in the clone operation (below) at the
-    // reference count meta data needs to be associated with this node and not its clone
-    @CompilationFinal private boolean createClone;
+    // stays the same during cloning
+    private final AssumedValue<Integer> transArgsBitSet;
+
+    private final ConditionProfile isNonNull = ConditionProfile.createBinaryProfile();
     private final ConditionProfile isRefCountUpdateable = ConditionProfile.createBinaryProfile();
 
-    private PostProcessArgumentsNode(LocalReadVariableNode[] sequence) {
-        this.sequence = sequence;
-        this.createClone = false;
-        this.transArgsBitSet = 0;
+    private PostProcessArgumentsNode(int length) {
+        this.sequence = new LocalReadVariableNode[Math.min(length, ArgumentStatePush.MAX_COUNTED_ARGS)];
+        this.transArgsBitSet = new AssumedValue<>("PostProcessArgumentsNode.transArgsBitSet", 0);
     }
 
     public static PostProcessArgumentsNode create(int length) {
-        int maxLength = Math.min(length, ArgumentStatePush.MAX_COUNTED_ARGS);
-        LocalReadVariableNode[] argReadNodes = new LocalReadVariableNode[maxLength];
-        for (int i = 0; i < maxLength; i++) {
-            argReadNodes[i] = LocalReadVariableNode.create(Integer.valueOf(1 << i), false);
-        }
-        return new PostProcessArgumentsNode(argReadNodes);
+        return new PostProcessArgumentsNode(length);
     }
 
     public int getLength() {
@@ -72,20 +62,20 @@ public final class PostProcessArgumentsNode extends RNode {
     @Override
     @ExplodeLoop
     public Object execute(VirtualFrame frame) {
-        if (transArgsBitSet > 0) {
+        int bits = transArgsBitSet.get();
+        if (bits != 0) {
             for (int i = 0; i < sequence.length; i++) {
                 int mask = 1 << i;
-                if ((transArgsBitSet & mask) == mask) {
-                    RShareable s = (RShareable) sequence[i].execute(frame);
-                    if (s == null) {
-                        // it happens rarely in compiled code, but if it does happen, stop
-                        // decrementing reference counts
+                if ((bits & mask) != 0) {
+                    if (sequence[i] == null) {
                         CompilerDirectives.transferToInterpreterAndInvalidate();
-                        transArgsBitSet = ArgumentStatePush.INVALID_INDEX;
-                        break;
+                        sequence[i] = insert(LocalReadVariableNode.create(Integer.valueOf(mask), false));
                     }
-                    if (isRefCountUpdateable.profile(!s.isSharedPermanent())) {
-                        s.decRefCount();
+                    RShareable s = (RShareable) sequence[i].execute(frame);
+                    if (isNonNull.profile(s != null)) {
+                        if (isRefCountUpdateable.profile(!s.isSharedPermanent())) {
+                            s.decRefCount();
+                        }
                     }
                 }
             }
@@ -93,38 +83,15 @@ public final class PostProcessArgumentsNode extends RNode {
         return RNull.instance;
     }
 
-    @Override
-    public Node deepCopy() {
-        CompilerAsserts.neverPartOfCompilation();
-        if (createClone) {
-            return deepCopyUnconditional();
-        } else {
-            this.createClone = true;
-            return this;
-        }
-    }
-
-    /*
-     * Deep copies are also made from other places than the trufflerizer, in which case we need to
-     * always create the node chain.
-     */
-    private PostProcessArgumentsNode deepCopyUnconditional() {
-        CompilerAsserts.neverPartOfCompilation();
-        PostProcessArgumentsNode copy = (PostProcessArgumentsNode) super.deepCopy();
-        nextOptPostProccessArgNode = insert(copy);
-        return copy;
-
-    }
-
-    PostProcessArgumentsNode getActualNode() {
-        CompilerAsserts.neverPartOfCompilation();
-        if (nextOptPostProccessArgNode == null) {
-            return this;
-        } else {
-            PostProcessArgumentsNode nextNode = nextOptPostProccessArgNode.getActualNode();
-            // shorten up the lookup chain
-            nextOptPostProccessArgNode = insert(nextNode);
-            return nextNode;
+    public boolean updateBits(int index) {
+        if (index < sequence.length) {
+            int bits = transArgsBitSet.get();
+            int newBits = bits | (1 << index);
+            if (newBits != bits) {
+                transArgsBitSet.set(newBits);
+            }
+            return true;
         }
+        return false;
     }
 }

com.oracle.truffle.r.nodes/src/com/oracle/truffle/r/nodes/function/WrapArgumentNode.java

@@ -25,7 +25,6 @@ package com.oracle.truffle.r.nodes.function;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.NodeCost;
 import com.oracle.truffle.api.nodes.UnexpectedResultException;
-import com.oracle.truffle.api.profiles.BranchProfile;
 import com.oracle.truffle.r.nodes.access.ConstantNode;
 import com.oracle.truffle.r.runtime.data.RMissing;
 import com.oracle.truffle.r.runtime.data.RNull;
@@ -44,8 +43,6 @@ public final class WrapArgumentNode extends WrapArgumentBaseNode {
 
     @Child private ArgumentStatePush argPushStateNode;
 
-    private final BranchProfile refCounted = BranchProfile.create();
-
     private WrapArgumentNode(RNode operand, boolean modeChange, int index) {
         super(operand, modeChange);
         this.modeChange = modeChange;
@@ -79,9 +76,6 @@ public final class WrapArgumentNode extends WrapArgumentBaseNode {
             if (rShareable != null) {
                 shareable.enter();
                 argPushStateNode.executeObject(frame, rShareable);
-            } else if (argPushStateNode.refCounted()) {
-                refCounted.enter();
-                argPushStateNode.executeObject(frame, RNull.instance);
             }
         }
         return result;