Skip to content

Commit 50d6363

Browse files
msimacekmsimacek
committed
Return generator resume target through an object
1 parent 0a0a8a0 commit 50d6363

File tree

6 files changed

+136
-85
lines changed

6 files changed

+136
-85
lines changed

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

Lines changed: 23 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@
6262
import com.oracle.graal.python.nodes.PRaiseNode;
6363
import com.oracle.graal.python.nodes.SpecialMethodNames;
6464
import com.oracle.graal.python.nodes.bytecode.FrameInfo;
65+
import com.oracle.graal.python.nodes.bytecode.GeneratorResult;
6566
import com.oracle.graal.python.nodes.call.CallTargetInvokeNode;
6667
import com.oracle.graal.python.nodes.call.GenericInvokeNode;
6768
import com.oracle.graal.python.nodes.call.special.LookupAndCallVarargsNode;
@@ -142,7 +143,7 @@ abstract static class ResumeGeneratorNode extends Node {
142143
public abstract Object execute(VirtualFrame frame, PGenerator self, Object sendValue);
143144

144145
@Specialization(guards = {"!self.usesBytecode()", "sameCallTarget(self.getCurrentCallTarget(), call.getCallTarget())"}, limit = "getCallSiteInlineCacheMaxDepth()")
145-
Object cachedAST(VirtualFrame frame, PGenerator self, Object sendValue,
146+
static Object cachedAST(VirtualFrame frame, PGenerator self, Object sendValue,
146147
@Cached("createDirectCall(self.getCurrentCallTarget())") CallTargetInvokeNode call) {
147148
self.setRunning(true);
148149
Object[] arguments = prepareArguments(self);
@@ -156,13 +157,13 @@ Object cachedAST(VirtualFrame frame, PGenerator self, Object sendValue,
156157
throw e;
157158
} finally {
158159
self.setRunning(false);
159-
self.setNextCallTarget(PythonLanguage.get(this));
160+
self.setNextCallTarget();
160161
}
161162
}
162163

163164
@Specialization(guards = "!self.usesBytecode()", replaces = "cachedAST")
164165
@Megamorphic
165-
Object genericAST(VirtualFrame frame, PGenerator self, Object sendValue,
166+
static Object genericAST(VirtualFrame frame, PGenerator self, Object sendValue,
166167
@Cached ConditionProfile hasFrameProfile,
167168
@Cached GenericInvokeNode call) {
168169
self.setRunning(true);
@@ -181,7 +182,7 @@ Object genericAST(VirtualFrame frame, PGenerator self, Object sendValue,
181182
throw e;
182183
} finally {
183184
self.setRunning(false);
184-
self.setNextCallTarget(PythonLanguage.get(this));
185+
self.setNextCallTarget();
185186
}
186187
}
187188

@@ -196,9 +197,9 @@ Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
196197
if (sendValue != null) {
197198
PArguments.setSpecialArgument(arguments, sendValue);
198199
}
199-
Object result;
200+
GeneratorResult result;
200201
try {
201-
result = call.execute(frame, null, null, null, arguments);
202+
result = (GeneratorResult) call.execute(frame, null, null, null, arguments);
202203
} catch (PException e) {
203204
self.markAsFinished();
204205
// PEP 479 - StopIteration raised from generator body needs to be wrapped in
@@ -208,18 +209,7 @@ Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
208209
} finally {
209210
self.setRunning(false);
210211
}
211-
if (returnProfile.profile(result == null)) {
212-
// Null result indicates a generator return
213-
Object returnValue = self.getReturnValue();
214-
if (returnValue != PNone.NONE) {
215-
throw raiseNode.raise(StopIteration, returnValue);
216-
} else {
217-
throw raiseNode.raise(StopIteration);
218-
}
219-
} else {
220-
self.setNextCallTarget(PythonLanguage.get(this));
221-
}
222-
return result;
212+
return handleResult(self, result, returnProfile, raiseNode);
223213
}
224214

225215
@Specialization(guards = "self.usesBytecode()", replaces = "cached")
@@ -235,12 +225,12 @@ Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
235225
if (sendValue != null) {
236226
PArguments.setSpecialArgument(arguments, sendValue);
237227
}
238-
Object result;
228+
GeneratorResult result;
239229
try {
240230
if (hasFrameProfile.profile(frame != null)) {
241-
result = call.execute(frame, self.getCurrentCallTarget(), arguments);
231+
result = (GeneratorResult) call.execute(frame, self.getCurrentCallTarget(), arguments);
242232
} else {
243-
result = call.execute(self.getCurrentCallTarget(), arguments);
233+
result = (GeneratorResult) call.execute(self.getCurrentCallTarget(), arguments);
244234
}
245235
} catch (PException e) {
246236
self.markAsFinished();
@@ -251,18 +241,22 @@ Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
251241
} finally {
252242
self.setRunning(false);
253243
}
254-
if (returnProfile.profile(result == null)) {
255-
// Null result indicates a generator return
256-
Object returnValue = self.getReturnValue();
244+
return handleResult(self, result, returnProfile, raiseNode);
245+
}
246+
247+
private Object handleResult(PGenerator self, GeneratorResult result, ConditionProfile returnProfile, PRaiseNode raiseNode) {
248+
if (returnProfile.profile(result.isReturn)) {
249+
self.markAsFinished();
250+
Object returnValue = result.value;
257251
if (returnValue != PNone.NONE) {
258252
throw raiseNode.raise(StopIteration, returnValue);
259253
} else {
260254
throw raiseNode.raise(StopIteration);
261255
}
262256
} else {
263-
self.setNextCallTarget(PythonLanguage.get(this));
257+
self.handleResult(PythonLanguage.get(this), result);
264258
}
265-
return result;
259+
return result.value;
266260
}
267261

268262
protected static CallTargetInvokeNode createDirectCall(CallTarget target) {
@@ -638,8 +632,9 @@ static Object getFrame(PGenerator self,
638632
}
639633
} else {
640634
FrameInfo info = (FrameInfo) generatorFrame.getFrameDescriptor().getInfo();
641-
frame.setLasti(info.getBci(generatorFrame));
642-
frame.setLine(info.getLineno(generatorFrame));
635+
int bci = self.getBci();
636+
frame.setLasti(bci);
637+
frame.setLine(info.getRootNode().bciToLine(bci));
643638
}
644639
return frame;
645640
}

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

Lines changed: 34 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import com.oracle.graal.python.builtins.objects.iterator.PIntRangeIterator;
3636
import com.oracle.graal.python.builtins.objects.object.PythonBuiltinObject;
3737
import com.oracle.graal.python.nodes.bytecode.FrameInfo;
38+
import com.oracle.graal.python.nodes.bytecode.GeneratorResult;
3839
import com.oracle.graal.python.nodes.bytecode.PBytecodeGeneratorRootNode;
3940
import com.oracle.graal.python.nodes.bytecode.PBytecodeRootNode;
4041
import com.oracle.graal.python.nodes.generator.AbstractYieldNode;
@@ -62,7 +63,7 @@ public final class PGenerator extends PythonBuiltinObject {
6263
* entry point into the generator: the first call, and continuation for each yield. Each AST can
6364
* then specialize towards which nodes are executed when starting from that particular entry
6465
* point. When yielding, the next index to the next call target to continue from is updated via
65-
* {@link #setNextCallTarget}.
66+
* {@link #setNextCallTarget} or {@link #handleResult}.
6667
*/
6768
@CompilationFinal(dimensions = 1) protected final RootCallTarget[] callTargets;
6869
protected final Object[] arguments;
@@ -167,25 +168,32 @@ private PGenerator(PythonLanguage lang, String name, String qualname, PBytecodeR
167168
this.generatorInfo = null;
168169
}
169170

170-
public void setNextCallTarget(PythonLanguage language) {
171-
if (usesBytecode()) {
172-
currentCallTarget = getBci();
173-
if (callTargets[currentCallTarget] == null) {
174-
CompilerDirectives.transferToInterpreterAndInvalidate();
175-
PBytecodeGeneratorRootNode rootNode = new PBytecodeGeneratorRootNode(language, bytecodeRootNode, currentCallTarget, getStackTop());
176-
callTargets[currentCallTarget] = rootNode.getCallTarget();
177-
}
178-
} else {
179-
currentCallTarget = PArguments.getControlDataFromGeneratorArguments(getArguments()).getLastYieldIndex();
171+
public void handleResult(PythonLanguage language, GeneratorResult result) {
172+
assert usesBytecode();
173+
if (result.isReturn) {
174+
markAsFinished();
175+
return;
176+
}
177+
currentCallTarget = result.resumeBci;
178+
if (callTargets[currentCallTarget] == null) {
179+
CompilerDirectives.transferToInterpreterAndInvalidate();
180+
PBytecodeGeneratorRootNode rootNode = new PBytecodeGeneratorRootNode(language, bytecodeRootNode, result.resumeBci, result.resumeStackTop);
181+
callTargets[currentCallTarget] = rootNode.getCallTarget();
180182
}
181183
}
182184

185+
public void setNextCallTarget() {
186+
assert !usesBytecode();
187+
currentCallTarget = PArguments.getControlDataFromGeneratorArguments(getArguments()).getLastYieldIndex();
188+
}
189+
183190
/**
184191
* Returns the call target that should be used the next time the generator is called. Each time
185192
* a generator call target returns through a yield, the generator should be updated with the
186-
* next yield index to use via {@link #setNextCallTarget}
193+
* next yield index to use via {@link #handleResult}
187194
*/
188195
public RootCallTarget getCurrentCallTarget() {
196+
assert !finished;
189197
return callTargets[currentCallTarget];
190198
}
191199

@@ -212,27 +220,30 @@ public Object getYieldFrom() {
212220
}
213221
return null;
214222
} else {
215-
if (running) {
223+
if (running || finished) {
216224
return null;
217225
}
218-
return frameInfo.getYieldFrom(PArguments.getGeneratorFrame(arguments));
226+
return frameInfo.getYieldFrom(PArguments.getGeneratorFrame(arguments), getBci(), getCurrentRootNode().getResumeStackTop());
219227
}
220228
}
221229

230+
private PBytecodeGeneratorRootNode getCurrentRootNode() {
231+
assert usesBytecode();
232+
return (PBytecodeGeneratorRootNode) getCurrentCallTarget().getRootNode();
233+
}
234+
222235
public boolean isStarted() {
223236
return currentCallTarget != 0 && !running;
224237
}
225238

226239
public int getBci() {
227-
return frameInfo.getBci(PArguments.getGeneratorFrame(arguments));
228-
}
229-
230-
public int getStackTop() {
231-
return frameInfo.getGeneratorStackTop(PArguments.getGeneratorFrame(arguments));
232-
}
233-
234-
public Object getReturnValue() {
235-
return frameInfo.getGeneratorReturnValue(PArguments.getGeneratorFrame(arguments));
240+
if (!isStarted()) {
241+
return -1;
242+
} else if (finished) {
243+
return bytecodeRootNode.getCodeUnit().code.length;
244+
} else {
245+
return getCurrentRootNode().getResumeBci();
246+
}
236247
}
237248

238249
public Object[] getArguments() {

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/FrameInfo.java‎

Lines changed: 2 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -63,25 +63,15 @@ public int getBci(Frame frame) {
6363
return -1;
6464
}
6565

66-
public int getGeneratorStackTop(Frame frame) {
67-
return frame.getInt(rootNode.generatorStackTopOffset);
68-
}
69-
70-
public Object getYieldFrom(Frame generatorFrame) {
71-
int bci = getBci(generatorFrame);
66+
public Object getYieldFrom(Frame generatorFrame, int bci, int stackTop) {
7267
/* Match the `yield from` bytecode pattern and get the object from stack */
73-
if (bci > 3 && rootNode.bytecode[bci - 3] == OpCodesConstants.SEND && rootNode.bytecode[bci - 1] == OpCodesConstants.YIELD_VALUE &&
68+
if (bci > 3 && bci < rootNode.bytecode.length && rootNode.bytecode[bci - 3] == OpCodesConstants.SEND && rootNode.bytecode[bci - 1] == OpCodesConstants.YIELD_VALUE &&
7469
rootNode.bytecode[bci] == OpCodesConstants.RESUME_YIELD) {
75-
int stackTop = generatorFrame.getInt(rootNode.generatorStackTopOffset);
7670
return generatorFrame.getObject(stackTop);
7771
}
7872
return null;
7973
}
8074

81-
public Object getGeneratorReturnValue(Frame frame) {
82-
return frame.getObject(rootNode.generatorReturnOffset);
83-
}
84-
8575
public int getLineno(Frame frame) {
8676
return bciToLine(getBci(frame));
8777
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/GeneratorResult.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
3+
* DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4+
*
5+
* The Universal Permissive License (UPL), Version 1.0
6+
*
7+
* Subject to the condition set forth below, permission is hereby granted to any
8+
* person obtaining a copy of this software, associated documentation and/or
9+
* data (collectively the "Software"), free of charge and under any and all
10+
* copyright rights in the Software, and any and all patent rights owned or
11+
* freely licensable by each licensor hereunder covering either (i) the
12+
* unmodified Software as contributed to or provided by such licensor, or (ii)
13+
* the Larger Works (as defined below), to deal in both
14+
*
15+
* (a) the Software, and
16+
*
17+
* (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
18+
* one is included with the Software each a "Larger Work" to which the Software
19+
* is contributed by such licensors),
20+
*
21+
* without restriction, including without limitation the rights to copy, create
22+
* derivative works of, display, perform, and distribute the Software and make,
23+
* use, sell, offer for sale, import, export, have made, and have sold the
24+
* Software and the Larger Work(s), and to sublicense the foregoing rights on
25+
* either these or other terms.
26+
*
27+
* This license is subject to the following condition:
28+
*
29+
* The above copyright notice and either this complete permission notice or at a
30+
* minimum a reference to the UPL must be included in all copies or substantial
31+
* portions of the Software.
32+
*
33+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
34+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
35+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
36+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
37+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
38+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
39+
* SOFTWARE.
40+
*/
41+
package com.oracle.graal.python.nodes.bytecode;
42+
43+
import com.oracle.truffle.api.CompilerDirectives.ValueType;
44+
45+
@ValueType
46+
public class GeneratorResult {
47+
public final int resumeBci;
48+
public final int resumeStackTop;
49+
public final boolean isReturn;
50+
public final Object value;
51+
52+
private GeneratorResult(int resumeBci, int resumeStackTop, boolean isReturn, Object value) {
53+
this.resumeBci = resumeBci;
54+
this.resumeStackTop = resumeStackTop;
55+
this.isReturn = isReturn;
56+
this.value = value;
57+
}
58+
59+
public static GeneratorResult createYield(int resumeBci, int resumeStackTop, Object value) {
60+
return new GeneratorResult(resumeBci, resumeStackTop, false, value);
61+
}
62+
63+
public static GeneratorResult createReturn(Object value) {
64+
return new GeneratorResult(-1, -1, true, value);
65+
}
66+
}

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/PBytecodeGeneratorRootNode.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,4 +272,12 @@ public boolean isPythonInternal() {
272272
public SourceSection getSourceSection() {
273273
return rootNode.getSourceSection();
274274
}
275+
276+
public int getResumeBci() {
277+
return resumeBci;
278+
}
279+
280+
public int getResumeStackTop() {
281+
return resumeStackTop;
282+
}
275283
}

0 commit comments

Comments
 (0)