Sync locals back after a tracing call

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java

@@ -174,6 +174,7 @@
 import com.oracle.graal.python.nodes.expression.UnaryOpNode;
 import com.oracle.graal.python.nodes.frame.DeleteGlobalNode;
 import com.oracle.graal.python.nodes.frame.DeleteGlobalNodeGen;
+import com.oracle.graal.python.nodes.frame.GetFrameLocalsNode;
 import com.oracle.graal.python.nodes.frame.MaterializeFrameNode;
 import com.oracle.graal.python.nodes.frame.ReadFromLocalsNode;
 import com.oracle.graal.python.nodes.frame.ReadFromLocalsNodeGen;
@@ -2962,7 +2963,10 @@ private void invokeTraceFunction(VirtualFrame virtualFrame, Object arg, PythonCo
             if (line != -1) {
                 pyFrame.setLineLock(line);
             }
+            // Force locals dict sync, so that we can sync them back later
+            GetFrameLocalsNode.getUncached().execute(pyFrame);
             Object result = CallTernaryMethodNode.getUncached().execute(null, traceFn, pyFrame, event.pythonName, nonNullArg);
+            syncLocalsBack(virtualFrame, pyFrame);
             Object realResult = result == PNone.NONE ? null : result;
             pyFrame.setLocalTraceFun(realResult);
         } catch (Throwable e) {
@@ -2976,6 +2980,36 @@ private void invokeTraceFunction(VirtualFrame virtualFrame, Object arg, PythonCo
         }
     }
 
+    // PyFrame_LocalsToFast
+    private void syncLocalsBack(VirtualFrame virtualFrame, PFrame pyFrame) {
+        Frame localFrame = virtualFrame;
+        if (co.isGeneratorOrCoroutine()) {
+            localFrame = PArguments.getGeneratorFrame(virtualFrame);
+        }
+        Object localsDict = pyFrame.getLocalsDict();
+        copyLocalsArray(localFrame, localsDict, varnames, 0, false);
+        copyLocalsArray(localFrame, localsDict, cellvars, celloffset, true);
+        copyLocalsArray(localFrame, localsDict, freevars, freeoffset, true);
+    }
+
+    private static void copyLocalsArray(Frame localFrame, Object localsDict, TruffleString[] namesArray, int offset, boolean deref) {
+        for (int i = 0; i < namesArray.length; i++) {
+            TruffleString varname = namesArray[i];
+            Object value = null;
+            try {
+                value = PyObjectGetItem.getUncached().execute(null, localsDict, varname);
+            } catch (AbstractTruffleException e) {
+                // CPython ignores all exceptions
+            }
+            if (deref) {
+                PCell cell = (PCell) localFrame.getObject(offset + i);
+                cell.setRef(value);
+            } else {
+                localFrame.setObject(offset + i, value);
+            }
+        }
+    }
+
     private void profileCEvent(VirtualFrame virtualFrame, Object callable, PythonContext.ProfileEvent event, MutableLoopData mutableData, boolean profilingEnabled) {
         if (profilingEnabled) {
             profileCEvent(virtualFrame, callable, event, mutableData);
@@ -3012,7 +3046,10 @@ private void invokeProfileFunction(VirtualFrame virtualFrame, Object arg, Python
         }
 
         try {
+            // Force locals dict sync, so that we can sync them back later
+            GetFrameLocalsNode.getUncached().execute(pyFrame);
             Object result = CallTernaryMethodNode.getUncached().execute(null, profileFun, pyFrame, event.name, arg == null ? PNone.NONE : arg);
+            syncLocalsBack(virtualFrame, pyFrame);
             Object realResult = result == PNone.NONE ? null : result;
             pyFrame.setLocalTraceFun(realResult);
         } catch (Throwable e) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/frame/GetFrameLocalsNode.java

@@ -140,4 +140,8 @@ protected static FrameInfo getInfo(FrameDescriptor fd) {
     public static GetFrameLocalsNode create() {
         return GetFrameLocalsNodeGen.create();
     }
+
+    public static GetFrameLocalsNode getUncached() {
+        return GetFrameLocalsNodeGen.getUncached();
+    }
 }