try to assume cells are effectively final

Author: timfel

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cell/PCell.java

@@ -47,10 +47,14 @@
 import com.oracle.graal.python.builtins.PythonBuiltinClassType;
 import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
+import com.oracle.truffle.api.Assumption;
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
+import com.oracle.truffle.api.Truffle;
 
 public class PCell extends PythonBuiltinObject {
-    private Object ref;
+    private final Assumption effectivelyFinal = Truffle.getRuntime().createAssumption("cell is effectively final");
+    @CompilationFinal private Object ref;
 
     public PCell() {
         super(PythonLanguage.getCore().lookupType(PythonBuiltinClassType.PCell));
@@ -61,13 +65,22 @@ public Object getRef() {
     }
 
     public void clearRef() {
-        ref = null;
+        setRef(null);
     }
 
     public void setRef(Object ref) {
+        if (effectivelyFinal.isValid()) {
+            if (this.ref != null) {
+                effectivelyFinal.invalidate();
+            }
+        }
         this.ref = ref;
     }
 
+    public Assumption isEffectivelyFinalAssumption() {
+        return effectivelyFinal;
+    }
+
     public Object getPythonRef() {
         if (ref == null) {
             return PNone.NONE;

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/cell/ReadLocalCellNode.java

@@ -29,29 +29,32 @@
 import static com.oracle.graal.python.runtime.exception.PythonErrorType.UnboundLocalError;
 
 import com.oracle.graal.python.builtins.objects.cell.PCell;
+import com.oracle.graal.python.nodes.PBaseNode;
 import com.oracle.graal.python.nodes.PNode;
+import com.oracle.graal.python.nodes.cell.ReadLocalCellNodeGen.ReadFromCellNodeGen;
 import com.oracle.graal.python.nodes.frame.ReadLocalNode;
 import com.oracle.graal.python.nodes.frame.ReadLocalVariableNode;
 import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Fallback;
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.FrameSlot;
 import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.NodeCost;
 import com.oracle.truffle.api.nodes.NodeInfo;
 import com.oracle.truffle.api.profiles.ValueProfile;
 
 @NodeInfo(shortName = "read_cell")
 public abstract class ReadLocalCellNode extends PNode implements ReadLocalNode {
     @Child private PNode readLocal;
-
+    @Child private ReadFromCellNode readCell;
     private final FrameSlot frameSlot;
-    private final boolean isFreeVar;
 
     ReadLocalCellNode(FrameSlot frameSlot, boolean isFreeVar) {
         this.frameSlot = frameSlot;
         this.readLocal = ReadLocalVariableNode.create(frameSlot);
-        this.isFreeVar = isFreeVar;
+        this.readCell = ReadFromCellNodeGen.create(isFreeVar, frameSlot.getIdentifier());
     }
 
     public static PNode create(FrameSlot frameSlot, boolean isFreeVar) {
@@ -63,25 +66,56 @@ public PNode makeWriteNode(PNode rhs) {
         return WriteLocalCellNode.create(frameSlot, rhs);
     }
 
-    @Specialization
-    Object readObject(VirtualFrame frame,
-                    @Cached("createClassProfile()") ValueProfile refTypeProfile) {
-        Object cell = readLocal.execute(frame);
-        if (cell instanceof PCell) {
-            Object ref = refTypeProfile.profile(((PCell) cell).getRef());
+    static abstract class ReadFromCellNode extends PBaseNode {
+        private final boolean isFreeVar;
+        private final Object identifier;
+
+        public ReadFromCellNode(boolean isFreeVar, Object identifier) {
+            this.isFreeVar = isFreeVar;
+            this.identifier = identifier;
+        }
+
+        abstract Object execute(Object cell);
+
+        @Specialization(guards = "cachedCell == cell", limit = "1", assumptions = "cachedCell.isEffectivelyFinalAssumption()")
+        Object readFinal(@SuppressWarnings("unused") PCell cell,
+                        @SuppressWarnings("unused") @Cached("cell") PCell cachedCell,
+                        @Cached("cell.getRef()") Object ref) {
+            return ref;
+        }
+
+        @Specialization(replaces = "readFinal")
+        Object read(PCell cell,
+                        @Cached("createClassProfile()") ValueProfile refTypeProfile) {
+            Object ref = refTypeProfile.profile(cell.getRef());
             if (ref != null) {
                 return ref;
+            } else {
+                throw raiseUnbound();
             }
-            throw raiseUnbound();
         }
-        CompilerDirectives.transferToInterpreter();
-        throw new IllegalStateException("Expected a cell, got: " + cell.toString() + " instead.");
-    }
 
-    private PException raiseUnbound() {
-        if (isFreeVar) {
-            return raise(NameError, "free variable '%s' referenced before assignment in enclosing scope", frameSlot.getIdentifier());
+        @Fallback
+        Object read(Object cell) {
+            CompilerDirectives.transferToInterpreter();
+            throw new IllegalStateException("Expected a cell, got: " + cell.toString() + " instead.");
         }
-        return raise(UnboundLocalError, "local variable '%s' referenced before assignment", frameSlot.getIdentifier());
+
+        private PException raiseUnbound() {
+            if (isFreeVar) {
+                return raise(NameError, "free variable '%s' referenced before assignment in enclosing scope", identifier);
+            }
+            return raise(UnboundLocalError, "local variable '%s' referenced before assignment", identifier);
+        }
+    }
+
+    @Specialization
+    Object readObject(VirtualFrame frame) {
+        return readCell.execute(readLocal.execute(frame));
+    }
+
+    @Override
+    public NodeCost getCost() {
+        return NodeCost.NONE;
     }
 }