Skip to content

Commit d3b713d

Browse files
msimacekmsimacek
committed
Refactor generator next/send/throw
1 parent 3e3102e commit d3b713d

File tree

3 files changed

+82
-98
lines changed

3 files changed

+82
-98
lines changed

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

Lines changed: 78 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -130,18 +130,64 @@ private static Object[] prepareArguments(PGenerator self) {
130130
return arguments;
131131
}
132132

133-
private static Object resumeGenerator(PGenerator self, Object sendValue) {
134-
try {
133+
public static void checkResumable(PythonBuiltinBaseNode node, PGenerator self) {
134+
if (self.isFinished()) {
135+
throw node.raise(StopIteration);
136+
}
137+
if (self.isRunning()) {
138+
throw node.raise(ValueError, ErrorMessages.GENERATOR_ALREADY_EXECUTING);
139+
}
140+
}
141+
142+
@ImportStatic({PGuards.class, PythonOptions.class})
143+
@ReportPolymorphism
144+
abstract static class ResumeGeneratorNode extends Node {
145+
public abstract Object execute(VirtualFrame frame, PGenerator self, Object sendValue);
146+
147+
@Specialization(guards = "sameCallTarget(self.getCurrentCallTarget(), call.getCallTarget())", limit = "getCallSiteInlineCacheMaxDepth()")
148+
public Object cached(VirtualFrame frame, PGenerator self, Object sendValue,
149+
@Cached("createDirectCall(self.getCurrentCallTarget())") CallTargetInvokeNode call) {
150+
self.setRunning(true);
151+
Object[] arguments = prepareArguments(self);
152+
if (sendValue != null) {
153+
PArguments.setSpecialArgument(arguments, sendValue);
154+
}
155+
try {
156+
return call.execute(frame, null, null, arguments);
157+
} catch (PException e) {
158+
self.markAsFinished();
159+
throw e;
160+
} finally {
161+
self.setRunning(false);
162+
self.setNextCallTarget();
163+
}
164+
}
165+
166+
@Specialization(replaces = "cached")
167+
public Object generic(VirtualFrame frame, PGenerator self, Object sendValue,
168+
@Cached GenericInvokeNode call) {
135169
self.setRunning(true);
136-
PArguments.setSpecialArgument(self.getArguments(), sendValue);
137-
return self.getCurrentCallTarget().call(prepareArguments(self));
138-
} catch (PException e) {
139-
self.markAsFinished();
140-
throw e;
141-
} finally {
142-
self.setRunning(false);
143-
self.setNextCallTarget();
144-
PArguments.setSpecialArgument(self.getArguments(), null);
170+
Object[] arguments = prepareArguments(self);
171+
if (sendValue != null) {
172+
PArguments.setSpecialArgument(arguments, sendValue);
173+
}
174+
try {
175+
return call.execute(frame, self.getCurrentCallTarget(), arguments);
176+
} catch (PException e) {
177+
self.markAsFinished();
178+
throw e;
179+
} finally {
180+
self.setRunning(false);
181+
self.setNextCallTarget();
182+
}
183+
}
184+
185+
protected static CallTargetInvokeNode createDirectCall(CallTarget target) {
186+
return CallTargetInvokeNode.create(target, false, true);
187+
}
188+
189+
protected static boolean sameCallTarget(RootCallTarget target1, CallTarget target2) {
190+
return target1 == target2;
145191
}
146192
}
147193

@@ -204,65 +250,12 @@ public Object iter(PGenerator self) {
204250

205251
@Builtin(name = __NEXT__, minNumOfPositionalArgs = 1)
206252
@GenerateNodeFactory
207-
@ReportPolymorphism
208253
public abstract static class NextNode extends PythonUnaryBuiltinNode {
209-
210-
protected static CallTargetInvokeNode createDirectCall(CallTarget target) {
211-
return CallTargetInvokeNode.create(target, false, true);
212-
}
213-
214-
protected static GenericInvokeNode createIndirectCall() {
215-
return GenericInvokeNode.create();
216-
}
217-
218-
protected static boolean sameCallTarget(RootCallTarget target1, CallTarget target2) {
219-
return target1 == target2;
220-
}
221-
222-
@Specialization(guards = "sameCallTarget(self.getCurrentCallTarget(), call.getCallTarget())", limit = "getCallSiteInlineCacheMaxDepth()")
223-
public Object nextCached(VirtualFrame frame, PGenerator self,
224-
@Cached("createDirectCall(self.getCurrentCallTarget())") CallTargetInvokeNode call,
225-
@Cached BranchProfile alreadyRunning) {
226-
if (self.isFinished()) {
227-
throw raise(StopIteration);
228-
}
229-
if (self.isRunning()) {
230-
alreadyRunning.enter();
231-
throw raise(ValueError, ErrorMessages.GENERATOR_ALREADY_EXECUTING);
232-
}
233-
try {
234-
self.setRunning(true);
235-
return call.execute(frame, null, null, prepareArguments(self));
236-
} catch (PException e) {
237-
self.markAsFinished();
238-
throw e;
239-
} finally {
240-
self.setRunning(false);
241-
self.setNextCallTarget();
242-
}
243-
}
244-
245-
@Specialization(replaces = "nextCached")
254+
@Specialization
246255
public Object next(VirtualFrame frame, PGenerator self,
247-
@Cached("createIndirectCall()") GenericInvokeNode call,
248-
@Cached BranchProfile alreadyRunning) {
249-
if (self.isFinished()) {
250-
throw raise(StopIteration);
251-
}
252-
if (self.isRunning()) {
253-
alreadyRunning.enter();
254-
throw raise(ValueError, ErrorMessages.GENERATOR_ALREADY_EXECUTING);
255-
}
256-
try {
257-
self.setRunning(true);
258-
return call.execute(frame, self.getCurrentCallTarget(), prepareArguments(self));
259-
} catch (PException e) {
260-
self.markAsFinished();
261-
throw e;
262-
} finally {
263-
self.setRunning(false);
264-
self.setNextCallTarget();
265-
}
256+
@Cached ResumeGeneratorNode resumeGeneratorNode) {
257+
checkResumable(this, self);
258+
return resumeGeneratorNode.execute(frame, self, null);
266259
}
267260
}
268261

@@ -271,13 +264,10 @@ public Object next(VirtualFrame frame, PGenerator self,
271264
public abstract static class SendNode extends PythonBuiltinNode {
272265

273266
@Specialization
274-
public Object send(PGenerator self, Object value,
275-
@Cached BranchProfile alreadyRunning) {
276-
if (self.isRunning()) {
277-
alreadyRunning.enter();
278-
throw raise(ValueError, ErrorMessages.GENERATOR_ALREADY_EXECUTING);
279-
}
280-
return resumeGenerator(self, value);
267+
public Object send(VirtualFrame frame, PGenerator self, Object value,
268+
@Cached ResumeGeneratorNode resumeGeneratorNode) {
269+
checkResumable(this, self);
270+
return resumeGeneratorNode.execute(frame, self, value);
281271
}
282272
}
283273

@@ -393,38 +383,32 @@ private void checkExceptionClass(Object type) {
393383
@Specialization
394384
Object sendThrow(VirtualFrame frame, PGenerator self, Object typ, Object val, @SuppressWarnings("unused") PNone tb,
395385
@Cached PrepareExceptionNode prepareExceptionNode,
396-
@Cached BranchProfile alreadyRunning,
386+
@Cached ResumeGeneratorNode resumeGeneratorNode,
397387
@Shared("language") @CachedLanguage PythonLanguage language) {
398-
if (self.isRunning()) {
399-
alreadyRunning.enter();
400-
throw raise(ValueError, ErrorMessages.GENERATOR_ALREADY_EXECUTING);
401-
}
388+
checkResumable(this, self);
402389
PBaseException instance = prepareExceptionNode.execute(frame, typ, val);
403-
return doThrow(self, instance, language);
390+
return doThrow(frame, resumeGeneratorNode, self, instance, language);
404391
}
405392

406393
@Specialization
407394
Object sendThrow(VirtualFrame frame, PGenerator self, Object typ, Object val, PTraceback tb,
408395
@Cached PrepareExceptionNode prepareExceptionNode,
409-
@Cached BranchProfile alreadyRunning,
396+
@Cached ResumeGeneratorNode resumeGeneratorNode,
410397
@Shared("language") @CachedLanguage PythonLanguage language) {
411-
if (self.isRunning()) {
412-
alreadyRunning.enter();
413-
throw raise(ValueError, ErrorMessages.GENERATOR_ALREADY_EXECUTING);
414-
}
398+
checkResumable(this, self);
415399
PBaseException instance = prepareExceptionNode.execute(frame, typ, val);
416400
instance.setTraceback(tb);
417-
return doThrow(self, instance, language);
401+
return doThrow(frame, resumeGeneratorNode, self, instance, language);
418402
}
419403

420-
private Object doThrow(PGenerator self, PBaseException instance, PythonLanguage language) {
404+
private Object doThrow(VirtualFrame frame, ResumeGeneratorNode resumeGeneratorNode, PGenerator self, PBaseException instance, PythonLanguage language) {
421405
instance.setContext(null); // Will be filled when caught
422406
if (self.isStarted()) {
423407
instance.ensureReified();
424408
// Pass it to the generator where it will be thrown by the last yield, the location
425409
// will be filled there
426410
PException pException = PException.fromObject(instance, null, PythonOptions.isPExceptionWithJavaStacktrace(language));
427-
return resumeGenerator(self, pException);
411+
return resumeGeneratorNode.execute(frame, self, pException);
428412
} else {
429413
// Unstarted generator, we cannot pass the exception into the generator as there is
430414
// nothing that would handle it.
@@ -470,24 +454,24 @@ private GetTracebackNode ensureGetTracebackNode() {
470454
@GenerateNodeFactory
471455
public abstract static class CloseNode extends PythonUnaryBuiltinNode {
472456
@Specialization
473-
Object close(PGenerator self,
457+
Object close(VirtualFrame frame, PGenerator self,
474458
@CachedLibrary(limit = "3") PythonObjectLibrary lib,
475459
@Cached IsBuiltinClassProfile isGeneratorExit,
476460
@Cached IsBuiltinClassProfile isStopIteration,
477-
@Cached BranchProfile alreadyRunning,
461+
@Cached ResumeGeneratorNode resumeGeneratorNode,
462+
@Cached ConditionProfile isStartedPorfile,
478463
@CachedLanguage PythonLanguage language) {
479464
if (self.isRunning()) {
480-
alreadyRunning.enter();
481465
throw raise(ValueError, ErrorMessages.GENERATOR_ALREADY_EXECUTING);
482466
}
483-
if (self.isStarted()) {
467+
if (isStartedPorfile.profile(self.isStarted())) {
484468
PBaseException pythonException = factory().createBaseException(GeneratorExit);
485469
// Pass it to the generator where it will be thrown by the last yield, the location
486470
// will be filled there
487471
boolean withJavaStacktrace = PythonOptions.isPExceptionWithJavaStacktrace(language);
488472
PException pException = PException.fromObject(pythonException, null, withJavaStacktrace);
489473
try {
490-
resumeGenerator(self, pException);
474+
resumeGeneratorNode.execute(frame, self, pException);
491475
} catch (PException pe) {
492476
if (isGeneratorExit.profileException(pe, GeneratorExit, lib) || isStopIteration.profileException(pe, StopIteration, lib)) {
493477
// This is the "success" path

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,10 @@ public static PGenerator create(String name, String qualname, RootCallTarget[] c
8282
MaterializedFrame generatorFrame = Truffle.getRuntime().createMaterializedFrame(generatorFrameArguments, frameDescriptor);
8383
PArguments.setGeneratorFrame(arguments, generatorFrame);
8484
PArguments.setControlData(arguments, generatorArgs);
85+
// The invoking node will set these two to the correct value only when the callee requests
86+
// it, otherwise they stay at the initial value, which we must set to null here
87+
PArguments.setException(arguments, null);
88+
PArguments.setCallerFrameInfo(arguments, null);
8589
PArguments.setCurrentFrameInfo(generatorFrameArguments, new PFrame.Reference(null));
8690
// set generator closure to the generator frame locals
8791
CompilerAsserts.partialEvaluationConstant(cellSlots);

‎graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/function/GeneratorExpressionNode.java

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@
3232
import com.oracle.graal.python.parser.DefinitionCellSlots;
3333
import com.oracle.graal.python.parser.ExecutionCellSlots;
3434
import com.oracle.graal.python.parser.GeneratorInfo;
35-
import com.oracle.graal.python.runtime.exception.PException;
3635
import com.oracle.graal.python.runtime.object.PythonObjectFactory;
3736
import com.oracle.truffle.api.CompilerAsserts;
3837
import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
@@ -110,9 +109,6 @@ public Object execute(VirtualFrame frame) {
110109
}
111110
PArguments.setGlobals(arguments, PArguments.getGlobals(frame));
112111

113-
// The generator doesn't capture the currently handled exception at creation time.
114-
PArguments.setException(arguments, PException.NO_EXCEPTION);
115-
116112
if (callTargets == null) {
117113
callTargets = GeneratorFunctionRootNode.createYieldTargets(callTarget);
118114
}

0 commit comments

Comments
 (0)