Handle StopIteration from a generator

msimacek · msimacek · commit 52fdc80daa2e · 2022-05-25T16:46:10.000+02:00
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/GeneratorBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/GeneratorBuiltins.java
@@ -143,8 +143,7 @@ abstract static class ResumeGeneratorNode extends Node {
 
         @Specialization(guards = "sameCallTarget(self.getCurrentCallTarget(), call.getCallTarget())", limit = "getCallSiteInlineCacheMaxDepth()")
         Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
-                        @Cached("createDirectCall(self.getCurrentCallTarget())") CallTargetInvokeNode call,
-                        @Cached PRaiseNode raiseNode) {
+                        @Cached("createDirectCall(self.getCurrentCallTarget())") CallTargetInvokeNode call) {
             self.setRunning(true);
             Object[] arguments = prepareArguments(self);
             if (sendValue != null) {
@@ -153,9 +152,6 @@ Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
             Object result;
             try {
                 result = call.execute(frame, null, null, null, arguments);
-                if (result == null) {
-                    self.markAsFinished();
-                }
             } catch (PException e) {
                 self.markAsFinished();
                 throw e;
@@ -165,16 +161,14 @@ Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
                     self.setNextCallTarget(PythonLanguage.get(this));
                 }
             }
-            handleResult(self, result, raiseNode);
             return result;
         }
 
         @Specialization(replaces = "cached")
         @Megamorphic
         Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
                         @Cached ConditionProfile hasFrameProfile,
-                        @Cached GenericInvokeNode call,
-                        @Cached PRaiseNode raiseNode) {
+                        @Cached GenericInvokeNode call) {
             self.setRunning(true);
             Object[] arguments = prepareArguments(self);
             if (sendValue != null) {
@@ -187,9 +181,6 @@ Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
                 } else {
                     result = call.execute(self.getCurrentCallTarget(), arguments);
                 }
-                if (result == null) {
-                    self.markAsFinished();
-                }
             } catch (PException e) {
                 self.markAsFinished();
                 throw e;
@@ -199,21 +190,9 @@ Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
                     self.setNextCallTarget(PythonLanguage.get(this));
                 }
             }
-            handleResult(self, result, raiseNode);
             return result;
         }
 
-        private static void handleResult(PGenerator self, Object result, PRaiseNode raiseNode) {
-            if (result == null) {
-                Object returnValue = self.getReturnValue();
-                if (returnValue != PNone.NONE) {
-                    throw raiseNode.raise(StopIteration, returnValue);
-                } else {
-                    throw raiseNode.raise(StopIteration);
-                }
-            }
-        }
-
         protected static CallTargetInvokeNode createDirectCall(CallTarget target) {
             return CallTargetInvokeNode.create(target, false, true);
         }
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/PGenerator.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/generator/PGenerator.java
@@ -182,7 +182,7 @@ public void setNextCallTarget(PythonLanguage language) {
     /**
      * Returns the call target that should be used the next time the generator is called. Each time
      * a generator call target returns through a yield, the generator should be updated with the
-     * next yield index to use via {@link #setNextCallTarget()}
+     * next yield index to use via {@link #setNextCallTarget}
      */
     public RootCallTarget getCurrentCallTarget() {
         return callTargets[currentCallTarget];
@@ -230,10 +230,6 @@ public int getStackTop() {
         return frameInfo.getGeneratorStackTop(PArguments.getGeneratorFrame(arguments));
     }
 
-    public Object getReturnValue() {
-        return frameInfo.getGeneratorReturnValue(PArguments.getGeneratorFrame(arguments));
-    }
-
     public Object[] getArguments() {
         return arguments;
     }
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorFunctionRootNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorFunctionRootNode.java
@@ -69,7 +69,7 @@ public PBytecodeGeneratorFunctionRootNode(PythonLanguage language, FrameDescript
         this.originalName = originalName;
         // TODO compress somehow
         this.callTargets = new RootCallTarget[rootNode.getCodeUnit().code.length];
-        this.callTargets[0] = rootNode.getCallTarget();
+        this.callTargets[0] = new PBytecodeGeneratorRootNode(language, rootNode, 0, rootNode.getInitialStackTop()).getCallTarget();
     }
 
     @Override
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorRootNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorRootNode.java
@@ -40,13 +40,25 @@
  */
 package com.oracle.graal.python.nodes.bytecode;
 
+import static com.oracle.graal.python.builtins.PythonBuiltinClassType.RuntimeError;
+import static com.oracle.graal.python.runtime.exception.PythonErrorType.StopIteration;
+
 import com.oracle.graal.python.PythonLanguage;
+import com.oracle.graal.python.builtins.objects.PNone;
+import com.oracle.graal.python.builtins.objects.function.PArguments;
 import com.oracle.graal.python.builtins.objects.function.Signature;
+import com.oracle.graal.python.nodes.ErrorMessages;
+import com.oracle.graal.python.nodes.PRaiseNode;
 import com.oracle.graal.python.nodes.PRootNode;
+import com.oracle.graal.python.nodes.object.IsBuiltinClassProfile;
 import com.oracle.graal.python.runtime.ExecutionContext;
+import com.oracle.graal.python.runtime.exception.PException;
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import com.oracle.truffle.api.frame.MaterializedFrame;
 import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.profiles.ConditionProfile;
 import com.oracle.truffle.api.source.SourceSection;
 
 public class PBytecodeGeneratorRootNode extends PRootNode {
@@ -55,6 +67,10 @@ public class PBytecodeGeneratorRootNode extends PRootNode {
     private final int resumeStackTop;
 
     @Child private ExecutionContext.CalleeContext calleeContext = ExecutionContext.CalleeContext.create();
+    @Child private IsBuiltinClassProfile errorProfile;
+    @Child private PRaiseNode raise = PRaiseNode.create();
+
+    private final ConditionProfile returnProfile = ConditionProfile.create();
 
     @TruffleBoundary
     public PBytecodeGeneratorRootNode(PythonLanguage language, PBytecodeRootNode rootNode, int resumeBci, int resumeStackTop) {
@@ -67,11 +83,29 @@ public PBytecodeGeneratorRootNode(PythonLanguage language, PBytecodeRootNode roo
     @Override
     public Object execute(VirtualFrame frame) {
         calleeContext.enter(frame);
+        Object result;
         try {
-            return rootNode.executeFromBci(frame, resumeBci, resumeStackTop);
+            result = rootNode.executeFromBci(frame, resumeBci, resumeStackTop);
+        } catch (PException pe) {
+            // PEP 479 - StopIteration raised from generator body needs to be wrapped in
+            // RuntimeError
+            pe.expectStopIteration(getErrorProfile());
+            throw raise.raise(RuntimeError, pe.setCatchingFrameAndGetEscapedException(frame, this), ErrorMessages.GENERATOR_RAISED_STOPITER);
         } finally {
             calleeContext.exit(frame, this);
         }
+        if (returnProfile.profile(result == null)) {
+            // Null result indicates a generator return
+            MaterializedFrame generatorFrame = PArguments.getGeneratorFrame(frame);
+            PBytecodeRootNode.FrameInfo info = (PBytecodeRootNode.FrameInfo) generatorFrame.getFrameDescriptor().getInfo();
+            Object returnValue = info.getGeneratorReturnValue(generatorFrame);
+            if (returnValue != PNone.NONE) {
+                throw raise.raise(StopIteration, returnValue);
+            } else {
+                throw raise.raise(StopIteration);
+            }
+        }
+        return result;
     }
 
     @Override
@@ -99,4 +133,12 @@ public boolean isPythonInternal() {
     public SourceSection getSourceSection() {
         return rootNode.getSourceSection();
     }
+
+    private IsBuiltinClassProfile getErrorProfile() {
+        if (errorProfile == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            errorProfile = insert(IsBuiltinClassProfile.create());
+        }
+        return errorProfile;
+    }
 }
diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeRootNode.java
@@ -770,7 +770,7 @@ public void createGeneratorFrame(Object[] arguments) {
         PArguments.setGeneratorFrameLocals(generatorFrameArguments, factory.createDictLocals(generatorFrame));
         // This is just for inspection, unstarted generators should have bci == -1
         generatorFrame.setInt(bcioffset, -1);
-        generatorFrame.setInt(generatorStackTopOffset, stackoffset - 1);
+        generatorFrame.setInt(generatorStackTopOffset, -1);
         copyArgsAndCells(generatorFrame, arguments);
     }
 
@@ -787,6 +787,10 @@ private void copyArgsAndCells(Frame localFrame, Object[] arguments) {
         initFreeVars(localFrame, arguments);
     }
 
+    int getInitialStackTop() {
+        return stackoffset - 1;
+    }
+
     @Override
     public Object execute(VirtualFrame virtualFrame) {
         calleeContext.enter(virtualFrame);
@@ -795,7 +799,7 @@ public Object execute(VirtualFrame virtualFrame) {
                 copyArgsAndCells(virtualFrame, virtualFrame.getArguments());
             }
 
-            return executeFromBci(virtualFrame, 0, stackoffset - 1);
+            return executeFromBci(virtualFrame, 0, getInitialStackTop());
         } finally {
             calleeContext.exit(virtualFrame, this);
         }

Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ public PBytecodeGeneratorFunctionRootNode(PythonLanguage language, FrameDescript`
`69`	`69`	`this.originalName = originalName;`
`70`	`70`	`// TODO compress somehow`
`71`	`71`	`this.callTargets = new RootCallTarget[rootNode.getCodeUnit().code.length];`
`72`		`- this.callTargets[0] = rootNode.getCallTarget();`
	`72`	`+ this.callTargets[0] = new PBytecodeGeneratorRootNode(language, rootNode, 0, rootNode.getInitialStackTop()).getCallTarget();`
`73`	`73`	`}`
`74`	`74`
`75`	`75`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -770,7 +770,7 @@ public void createGeneratorFrame(Object[] arguments) {`
`770`	`770`	`PArguments.setGeneratorFrameLocals(generatorFrameArguments, factory.createDictLocals(generatorFrame));`
`771`	`771`	`// This is just for inspection, unstarted generators should have bci == -1`
`772`	`772`	`generatorFrame.setInt(bcioffset, -1);`
`773`		`- generatorFrame.setInt(generatorStackTopOffset, stackoffset - 1);`
	`773`	`+ generatorFrame.setInt(generatorStackTopOffset, -1);`
`774`	`774`	`copyArgsAndCells(generatorFrame, arguments);`
`775`	`775`	`}`
`776`	`776`
`@@ -787,6 +787,10 @@ private void copyArgsAndCells(Frame localFrame, Object[] arguments) {`
`787`	`787`	`initFreeVars(localFrame, arguments);`
`788`	`788`	`}`
`789`	`789`
	`790`	`+ int getInitialStackTop() {`
	`791`	`+ return stackoffset - 1;`
	`792`	`+ }`
	`793`	`+`
`790`	`794`	`@Override`
`791`	`795`	`public Object execute(VirtualFrame virtualFrame) {`
`792`	`796`	`calleeContext.enter(virtualFrame);`
`@@ -795,7 +799,7 @@ public Object execute(VirtualFrame virtualFrame) {`
`795`	`799`	`copyArgsAndCells(virtualFrame, virtualFrame.getArguments());`
`796`	`800`	`}`
`797`	`801`
`798`		`- return executeFromBci(virtualFrame, 0, stackoffset - 1);`
	`802`	`+ return executeFromBci(virtualFrame, 0, getInitialStackTop());`
`799`	`803`	`} finally {`
`800`	`804`	`calleeContext.exit(virtualFrame, this);`
`801`	`805`	`}`