[Truffle] Insert new CachedYieldDispatchNodes at the front, like norm…

…al DispatchNodes.
jruby · Mar 9, 2015 · 504f4a8 · 504f4a8
1 parent d9ac26b
commit 504f4a8
Showing 1 changed file with 13 additions and 9 deletions.
diff --git a/truffle/src/main/java/org/jruby/truffle/nodes/yield/UninitializedYieldDispatchNode.java b/truffle/src/main/java/org/jruby/truffle/nodes/yield/UninitializedYieldDispatchNode.java
@@ -38,15 +38,15 @@ protected boolean guard(RubyProc block) {
     @Override
     public Object dispatchWithSelfAndBlock(VirtualFrame frame, final RubyProc block, Object self, RubyProc modifiedBlock, Object... argumentsObjects) {
         CompilerDirectives.transferToInterpreterAndInvalidate();
-        final UninitializedYieldDispatchNode currentNode = this;
 
         final YieldDispatchNode dispatch = atomic(new Callable<YieldDispatchNode>() {
             @Override
             public YieldDispatchNode call() {
                 // First try to see if we did not a miss a specialization added by another thread.
-                final YieldDispatchHeadNode dispatchHead = (YieldDispatchHeadNode) NodeUtil.getNthParent(currentNode, depth + 1);
+                final YieldDispatchHeadNode dispatchHead = getHeadNode();
+                final YieldDispatchNode first = dispatchHead.getDispatch();
 
-                YieldDispatchNode lookupDispatch = dispatchHead.getDispatch();
+                YieldDispatchNode lookupDispatch = first;
                 while (lookupDispatch != null) {
                     if (lookupDispatch.guard(block)) {
                         // This one worked, no need to rewrite anything.
@@ -55,20 +55,24 @@ public YieldDispatchNode call() {
                     lookupDispatch = lookupDispatch.getNext();
                 }
 
+                final YieldDispatchNode newDispatchNode;
                 if (depth < DispatchNode.DISPATCH_POLYMORPHIC_MAX) {
                     depth++;
-                    CachedYieldDispatchNode cachedDispatch = new CachedYieldDispatchNode(getContext(), block, currentNode);
-                    replace(cachedDispatch);
-                    return cachedDispatch;
+                    newDispatchNode = new CachedYieldDispatchNode(getContext(), block, first);
                 } else {
-                    GeneralYieldDispatchNode newGeneralYield = new GeneralYieldDispatchNode(getContext());
-                    dispatchHead.getDispatch().replace(newGeneralYield);
-                    return newGeneralYield;
+                    newDispatchNode = new GeneralYieldDispatchNode(getContext());
                 }
+
+                first.replace(newDispatchNode);
+                return first;
             }
         });
 
         return dispatch.dispatchWithSelfAndBlock(frame, block, self, modifiedBlock, argumentsObjects);
     }
 
+    private YieldDispatchHeadNode getHeadNode() {
+        return NodeUtil.findParent(this, YieldDispatchHeadNode.class);
+    }
+
 }