Fix generator code introspection

Author: msimacek

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/code/PCode.java

@@ -54,12 +54,13 @@
 import com.oracle.graal.python.builtins.objects.str.StringUtils;
 import com.oracle.graal.python.builtins.objects.tuple.PTuple;
 import com.oracle.graal.python.nodes.ModuleRootNode;
-import com.oracle.graal.python.nodes.bytecode.PBytecodeRootNode;
 import com.oracle.graal.python.nodes.PClosureFunctionRootNode;
 import com.oracle.graal.python.nodes.PClosureRootNode;
 import com.oracle.graal.python.nodes.PRootNode;
 import com.oracle.graal.python.nodes.argument.ReadVarArgsNode;
 import com.oracle.graal.python.nodes.argument.ReadVarKeywordsNode;
+import com.oracle.graal.python.nodes.bytecode.PBytecodeGeneratorFunctionRootNode;
+import com.oracle.graal.python.nodes.bytecode.PBytecodeRootNode;
 import com.oracle.graal.python.nodes.expression.ExpressionNode;
 import com.oracle.graal.python.nodes.frame.GlobalNode;
 import com.oracle.graal.python.nodes.frame.PythonFrame;
@@ -359,7 +360,13 @@ public boolean visit(Node node) {
     }
 
     private static RootNode rootNodeForExtraction(RootNode rootNode) {
-        return (rootNode instanceof GeneratorFunctionRootNode) ? ((GeneratorFunctionRootNode) rootNode).getFunctionRootNode() : rootNode;
+        if (rootNode instanceof GeneratorFunctionRootNode) {
+            return ((GeneratorFunctionRootNode) rootNode).getFunctionRootNode();
+        }
+        if (rootNode instanceof PBytecodeGeneratorFunctionRootNode) {
+            return ((PBytecodeGeneratorFunctionRootNode) rootNode).getBytecodeRootNode();
+        }
+        return rootNode;
     }
 
     @TruffleBoundary

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/GeneratorBuiltins.java

@@ -46,7 +46,6 @@
 import com.oracle.graal.python.builtins.PythonBuiltins;
 import com.oracle.graal.python.builtins.modules.BuiltinFunctions;
 import com.oracle.graal.python.builtins.objects.PNone;
-import com.oracle.graal.python.builtins.objects.code.PCode;
 import com.oracle.graal.python.builtins.objects.common.SequenceNodes.GetObjectArrayNode;
 import com.oracle.graal.python.builtins.objects.dict.PDict;
 import com.oracle.graal.python.builtins.objects.exception.PBaseException;
@@ -572,12 +571,7 @@ public abstract static class GetCodeNode extends PythonUnaryBuiltinNode {
         @Specialization
         Object getCode(PGenerator self,
                         @Cached ConditionProfile hasCodeProfile) {
-            PCode code = self.getCode();
-            if (hasCodeProfile.profile(code == null)) {
-                code = factory().createCode(self.getCurrentCallTarget());
-                self.setCode(code);
-            }
-            return code;
+            return self.getOrCreateCode(hasCodeProfile, factory());
         }
     }
 

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/PGenerator.java

@@ -53,6 +53,7 @@
 import com.oracle.truffle.api.frame.MaterializedFrame;
 import com.oracle.truffle.api.library.ExportMessage;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
+import com.oracle.truffle.api.profiles.ConditionProfile;
 
 public final class PGenerator extends PythonBuiltinObject {
 
@@ -276,14 +277,19 @@ public String toString() {
         return "<generator object " + name + " at " + hashCode() + ">";
     }
 
-    public PCode getCode() {
+    public PCode getOrCreateCode(ConditionProfile hasCodeProfile, PythonObjectFactory factory) {
+        if (hasCodeProfile.profile(code == null)) {
+            RootCallTarget callTarget;
+            if (usesBytecode()) {
+                callTarget = bytecodeRootNode.getCallTarget();
+            } else {
+                callTarget = callTargets[0];
+            }
+            code = factory.createCode(callTarget);
+        }
         return code;
     }
 
-    public void setCode(PCode code) {
-        this.code = code;
-    }
-
     public boolean isRunning() {
         return running;
     }

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

@@ -107,4 +107,8 @@ public boolean isPythonInternal() {
     public SourceSection getSourceSection() {
         return rootNode.getSourceSection();
     }
+
+    public PBytecodeRootNode getBytecodeRootNode() {
+        return rootNode;
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/object/IsNode.java

@@ -54,6 +54,7 @@
 import com.oracle.graal.python.builtins.objects.str.PString;
 import com.oracle.graal.python.builtins.objects.str.StringNodes;
 import com.oracle.graal.python.builtins.objects.type.PythonBuiltinClass;
+import com.oracle.graal.python.nodes.bytecode.PBytecodeGeneratorFunctionRootNode;
 import com.oracle.graal.python.nodes.expression.BinaryOp;
 import com.oracle.graal.python.nodes.generator.GeneratorFunctionRootNode;
 import com.oracle.graal.python.runtime.PythonContext;
@@ -218,9 +219,13 @@ static boolean doCode(PCode left, PCode right,
                 RootNode rightRootNode = rightCt.getRootNode();
                 if (leftRootNode instanceof GeneratorFunctionRootNode) {
                     leftRootNode = ((GeneratorFunctionRootNode) leftRootNode).getFunctionRootNode();
+                } else if (leftRootNode instanceof PBytecodeGeneratorFunctionRootNode) {
+                    leftRootNode = ((PBytecodeGeneratorFunctionRootNode) leftRootNode).getBytecodeRootNode();
                 }
                 if (rightRootNode instanceof GeneratorFunctionRootNode) {
                     rightRootNode = ((GeneratorFunctionRootNode) rightRootNode).getFunctionRootNode();
+                } else if (rightRootNode instanceof PBytecodeGeneratorFunctionRootNode) {
+                    rightRootNode = ((PBytecodeGeneratorFunctionRootNode) rightRootNode).getBytecodeRootNode();
                 }
                 return leftRootNode == rightRootNode;
             } else {