make sure we're checking constant assumptions on PCells of the same function root

Author: timfel

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

@@ -70,6 +70,19 @@ public void setRef(Object ref) {
         this.ref = ref;
     }
 
+    /**
+     * Use this to pass in the effectivelyFinal assumption from a node that made it constant.
+     */
+    public void setRef(Object ref, Assumption constantAssumption) {
+        assert constantAssumption == effectivelyFinal;
+        if (constantAssumption.isValid()) {
+            if (this.ref != null) {
+                constantAssumption.invalidate();
+            }
+        }
+        this.ref = ref;
+    }
+
     public Assumption isEffectivelyFinalAssumption() {
         return effectivelyFinal;
     }

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

@@ -32,6 +32,7 @@
 import com.oracle.graal.python.nodes.frame.ReadLocalVariableNode;
 import com.oracle.graal.python.nodes.frame.WriteIdentifierNode;
 import com.oracle.graal.python.nodes.statement.StatementNode;
+import com.oracle.truffle.api.Assumption;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.NodeChild;
@@ -93,7 +94,17 @@ void doWriteCached(@SuppressWarnings("unused") PCell cell, Object value,
             doWriteGeneric(cachedCell, value);
         }
 
-        @Specialization(replaces = "doWriteCached")
+        @Specialization(guards = "cell.isEffectivelyFinalAssumption() == effectivelyFinalAssumption", limit = "1", assumptions = "effectivelyFinalAssumption")
+        void doWriteCachedAssumption(PCell cell, Object value,
+                        @SuppressWarnings("unused") @Cached("cell.isEffectivelyFinalAssumption()") Assumption effectivelyFinalAssumption) {
+            if (value == NO_VALUE) {
+                cell.clearRef();
+            } else {
+                cell.setRef(value, effectivelyFinalAssumption);
+            }
+        }
+
+        @Specialization(replaces = {"doWriteCached", "doWriteCachedAssumption"})
         void doWriteGeneric(PCell cell, Object value) {
             if (value == NO_VALUE) {
                 cell.clearRef();