ImportStarNode: use fast get custom locals access

Author: cosminbasca

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

@@ -115,4 +115,8 @@ static ConditionProfile[] getFlag() {
     static ConditionProfile[] getFlagUncached() {
         return DISABLED;
     }
+
+    public static GetCurrentFrameRef create() {
+        return GetCurrentFrameRefNodeGen.create();
+    }
 }

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

@@ -44,6 +44,7 @@
 import com.oracle.graal.python.builtins.objects.frame.PFrame;
 import com.oracle.graal.python.builtins.objects.frame.PFrame.Reference;
 import com.oracle.graal.python.builtins.objects.function.PArguments;
+import com.oracle.graal.python.builtins.objects.object.PythonObject;
 import com.oracle.graal.python.nodes.function.ClassBodyRootNode;
 import com.oracle.graal.python.runtime.object.PythonObjectFactory;
 import com.oracle.truffle.api.dsl.Cached;
@@ -52,6 +53,7 @@
 import com.oracle.truffle.api.frame.Frame;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.Node;
+import com.oracle.truffle.api.profiles.ConditionProfile;
 
 /**
  * Read the locals from the passed frame, updating them from the frame if that is needed. This does
@@ -101,4 +103,21 @@ Object frameToUpdate(PFrame frame,
     public static ReadLocalsNode create() {
         return ReadLocalsNodeGen.create();
     }
+
+    /**
+     * Fast access to custom locals if a custom dict was provided by the caller and the python frame does not escape.
+     *
+     * @param frame the current frame
+     * @param havePyFrame profile if we do have a python frame (PFrame)
+     * @param haveCustomLocals profile if the python frame has a custom set locals dict
+     * @return the custom locals (if found) or the globals
+     */
+    public static PythonObject fastGetCustomLocalsOrGlobals(VirtualFrame frame, ConditionProfile havePyFrame, ConditionProfile haveCustomLocals) {
+        PFrame pyFrame = PArguments.getCurrentFrameInfo(frame).getPyFrame();
+        if (havePyFrame.profile(pyFrame != null) && haveCustomLocals.profile(pyFrame.getLocals(null) != null)) {
+            return  (PythonObject) pyFrame.getLocals(null);
+        }
+        // return the globals
+        return PArguments.getGlobals(frame);
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/statement/ImportStarNode.java

@@ -25,9 +25,10 @@
  */
 package com.oracle.graal.python.nodes.statement;
 
+import static com.oracle.graal.python.nodes.frame.ReadLocalsNode.fastGetCustomLocalsOrGlobals;
+
 import com.oracle.graal.python.builtins.PythonBuiltinClassType;
 import com.oracle.graal.python.builtins.objects.dict.PDict;
-import com.oracle.graal.python.builtins.objects.frame.PFrame;
 import com.oracle.graal.python.builtins.objects.function.PArguments;
 import com.oracle.graal.python.builtins.objects.mappingproxy.PMappingproxy;
 import com.oracle.graal.python.builtins.objects.module.PythonModule;
@@ -58,6 +59,7 @@
 
 public class ImportStarNode extends AbstractImportNode {
     private final ConditionProfile javaImport = ConditionProfile.createBinaryProfile();
+    private final ConditionProfile havePyFrame = ConditionProfile.createBinaryProfile();
     private final ConditionProfile haveCustomLocals = ConditionProfile.createBinaryProfile();
 
     @Child private SetItemNode dictWriteNode;
@@ -106,18 +108,8 @@ public ImportStarNode(String moduleName, int level) {
 
     @Override
     public void executeVoid(VirtualFrame frame) {
-        PythonObject globals = PArguments.getGlobals(frame);
-        Object importedModule = importModule(frame, moduleName, globals, new String[]{"*"}, level);
-        PythonObject locals = globals;
-
-        Object customLocals = PArguments.getCustomLocals(frame);
-        if (haveCustomLocals.profile(customLocals != null)) {
-            assert customLocals instanceof PFrame.Reference;
-            PFrame pyFrame = ((PFrame.Reference) customLocals).getPyFrame();
-            if (pyFrame != null) {
-                locals = (PythonObject) pyFrame.getLocals(factory());
-            }
-        }
+        Object importedModule = importModule(frame, moduleName, PArguments.getGlobals(frame), new String[]{"*"}, level);
+        PythonObject locals = fastGetCustomLocalsOrGlobals(frame, havePyFrame, haveCustomLocals);
 
         if (javaImport.profile(emulateJython() && getContext().getEnv().isHostObject(importedModule))) {
             try {