Optimize GetFrameLocalsNode

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/PythonLanguage.java

@@ -597,7 +597,7 @@ public Object execute(VirtualFrame frame) {
                 Object[] arguments = PArguments.create();
                 // escape?
                 PFrame pFrame = materializeFrameNode.execute(this, false, true, frame);
-                Object locals = getFrameLocalsNode.execute(frame, pFrame);
+                Object locals = getFrameLocalsNode.execute(pFrame);
                 PArguments.setSpecialArgument(arguments, locals);
                 PArguments.setGlobals(arguments, PArguments.getGlobals(frame));
                 boolean wasAcquired = gilNode.acquire();

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/BuiltinFunctions.java

@@ -861,8 +861,8 @@ private static void inheritGlobals(PFrame callerFrame, Object[] args) {
             PArguments.setGlobals(args, callerFrame.getGlobals());
         }
 
-        private static void inheritLocals(VirtualFrame frame, PFrame callerFrame, Object[] args, GetFrameLocalsNode getFrameLocalsNode) {
-            Object callerLocals = getFrameLocalsNode.execute(frame, callerFrame);
+        private static void inheritLocals(PFrame callerFrame, Object[] args, GetFrameLocalsNode getFrameLocalsNode) {
+            Object callerLocals = getFrameLocalsNode.execute(callerFrame);
             setCustomLocals(args, callerLocals);
         }
 
@@ -895,7 +895,7 @@ Object execInheritGlobalsInheritLocals(VirtualFrame frame, Object source, @Suppr
             PFrame callerFrame = readCallerFrameNode.executeWith(frame, 0);
             Object[] args = PArguments.create();
             inheritGlobals(callerFrame, args);
-            inheritLocals(frame, callerFrame, args, getFrameLocalsNode);
+            inheritLocals(callerFrame, args, getFrameLocalsNode);
 
             return invokeNode.execute(frame, getCt.execute(code), args);
         }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/frame/FrameBuiltins.java

@@ -232,7 +232,7 @@ Object getUpdating(VirtualFrame frame, PFrame self,
                 PFrame pyFrame = materializeNode.execute(false, true, frame);
                 assert pyFrame == self;
             }
-            return getFrameLocalsNode.execute(frame, self);
+            return getFrameLocalsNode.execute(self);
         }
     }
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/frame/PFrame.java

@@ -191,6 +191,10 @@ public Object getLocalsDict() {
         return localsDict;
     }
 
+    public boolean hasCustomLocals() {
+        return locals == null;
+    }
+
     public void setLocalsDict(Object dict) {
         localsDict = dict;
     }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/lib/PyDictSetItem.java

@@ -65,6 +65,10 @@ public static void executeUncached(PDict dict, Object key, Object item) {
         getUncached().execute(null, dict, key, item);
     }
 
+    public final void execute(PDict dict, TruffleString key, Object item) {
+        execute(null, dict, key, item);
+    }
+
     public abstract void execute(Frame frame, PDict dict, Object key, Object item);
 
     // We never need a frame for setting string keys

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/lib/PyEvalGetLocals.java

@@ -56,6 +56,6 @@ static Object getLocals(VirtualFrame frame,
                     @Cached ReadCallerFrameNode readCallerFrameNode,
                     @Cached GetFrameLocalsNode getFrameLocalsNode) {
         PFrame callerFrame = readCallerFrameNode.executeWith(frame, 0);
-        return getFrameLocalsNode.execute(frame, callerFrame);
+        return getFrameLocalsNode.execute(callerFrame);
     }
 }

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

@@ -40,21 +40,22 @@
  */
 package com.oracle.graal.python.nodes.frame;
 
-import com.oracle.graal.python.builtins.PythonBuiltinClassType;
 import com.oracle.graal.python.builtins.objects.cell.PCell;
+import com.oracle.graal.python.builtins.objects.dict.PDict;
 import com.oracle.graal.python.builtins.objects.frame.PFrame;
-import com.oracle.graal.python.lib.PyObjectDelItem;
-import com.oracle.graal.python.lib.PyObjectSetItem;
+import com.oracle.graal.python.lib.PyDictDelItem;
+import com.oracle.graal.python.lib.PyDictSetItem;
 import com.oracle.graal.python.nodes.bytecode.FrameInfo;
-import com.oracle.graal.python.nodes.object.IsBuiltinClassProfile;
-import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
+import com.oracle.truffle.api.dsl.Bind;
 import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Cached.Shared;
 import com.oracle.truffle.api.dsl.GenerateUncached;
 import com.oracle.truffle.api.dsl.NeverDefault;
 import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.frame.MaterializedFrame;
-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.strings.TruffleString;
 
@@ -65,51 +66,74 @@
  */
 @GenerateUncached
 public abstract class GetFrameLocalsNode extends Node {
-    public abstract Object execute(VirtualFrame frame, PFrame pyFrame);
+    public abstract Object execute(PFrame pyFrame);
 
-    // TODO cache fd, explode
-    @Specialization(guards = "hasLocals(pyFrame)")
-    static Object doLoop(VirtualFrame frame, PFrame pyFrame,
+    @Specialization(guards = "!pyFrame.hasCustomLocals()")
+    static Object doLoop(PFrame pyFrame,
                     @Cached PythonObjectFactory factory,
-                    @Cached PyObjectSetItem setItem,
-                    @Cached PyObjectDelItem delItem) {
-        // TODO do argcells need special handling?
+                    @Cached CopyLocalsToDict copyLocalsToDict) {
         MaterializedFrame locals = pyFrame.getLocals();
-        Object localsDict = pyFrame.getLocalsDict();
+        // It doesn't have custom locals, so it has to be a builtin dict or null
+        PDict localsDict = (PDict) pyFrame.getLocalsDict();
         if (localsDict == null) {
             localsDict = factory.createDict();
             pyFrame.setLocalsDict(localsDict);
         }
-        FrameInfo frameInfo = (FrameInfo) locals.getFrameDescriptor().getInfo();
-        for (int i = 0; i < frameInfo.getVariableCount(); i++) {
-            TruffleString name = frameInfo.getVariableName(i);
-            Object value = locals.getValue(i);
-            if (value instanceof PCell) {
-                value = ((PCell) value).getRef();
-            }
-            if (value == null) {
-                try {
-                    delItem.execute(frame, localsDict, name);
-                } catch (PException e) {
-                    // TODO
-                    e.expect(PythonBuiltinClassType.KeyError, IsBuiltinClassProfile.getUncached());
-                }
-            } else {
-                setItem.execute(frame, localsDict, name, value);
-            }
-        }
+        copyLocalsToDict.execute(locals, localsDict);
         return localsDict;
     }
 
-    @Specialization(guards = "!hasLocals(pyFrame)")
+    @Specialization(guards = "pyFrame.hasCustomLocals()")
     static Object doCustomLocals(PFrame pyFrame) {
         Object localsDict = pyFrame.getLocalsDict();
         assert localsDict != null;
         return localsDict;
     }
 
-    protected static boolean hasLocals(PFrame pyFrame) {
-        return pyFrame.getLocals() != null;
+    @GenerateUncached
+    abstract static class CopyLocalsToDict extends Node {
+        abstract void execute(MaterializedFrame locals, PDict dict);
+
+        @Specialization(guards = {"cachedFd == locals.getFrameDescriptor()", "count < 32"}, limit = "1")
+        @ExplodeLoop
+        void doCachedFd(MaterializedFrame locals, PDict dict,
+                        @SuppressWarnings("unused") @Cached("locals.getFrameDescriptor()") FrameDescriptor cachedFd,
+                        @Bind("getInfo(cachedFd)") FrameInfo info,
+                        @Bind("info.getVariableCount()") int count,
+                        @Shared("setItem") @Cached PyDictSetItem setItem,
+                        @Shared("delItem") @Cached PyDictDelItem delItem) {
+            for (int i = 0; i < count; i++) {
+                copyItem(locals, info, dict, setItem, delItem, i);
+            }
+        }
+
+        @Specialization(replaces = "doCachedFd")
+        void doGeneric(MaterializedFrame locals, PDict dict,
+                        @Shared("setItem") @Cached PyDictSetItem setItem,
+                        @Shared("delItem") @Cached PyDictDelItem delItem) {
+            FrameInfo info = getInfo(locals.getFrameDescriptor());
+            int count = info.getVariableCount();
+            for (int i = 0; i < count; i++) {
+                copyItem(locals, info, dict, setItem, delItem, i);
+            }
+        }
+
+        private static void copyItem(MaterializedFrame locals, FrameInfo info, PDict dict, PyDictSetItem setItem, PyDictDelItem delItem, int i) {
+            TruffleString name = info.getVariableName(i);
+            Object value = locals.getValue(i);
+            if (value instanceof PCell) {
+                value = ((PCell) value).getRef();
+            }
+            if (value == null) {
+                delItem.execute(dict, name);
+            } else {
+                setItem.execute(dict, name, value);
+            }
+        }
+
+        protected static FrameInfo getInfo(FrameDescriptor fd) {
+            return (FrameInfo) fd.getInfo();
+        }
     }
 
     @NeverDefault