Skip to content

Commit 559123c

Browse files
coolreader18coolreader18
committed
Don't push something on the stack when starting a generator
1 parent 2123df6 commit 559123c

5 files changed

Lines changed: 47 additions & 50 deletions

File tree

‎bytecode/src/lib.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,7 @@ pub enum Instruction {
378378
},
379379
GetAIter,
380380
GetANext,
381+
EndAsyncFor,
381382

382383
/// Reverse order evaluation in MapAdd
383384
/// required to support named expressions of Python 3.8 in dict comprehension
@@ -902,6 +903,7 @@ impl Instruction {
902903
SetupAsyncWith { .. } => 0,
903904
GetAIter => 0,
904905
GetANext => 1,
906+
EndAsyncFor => -1,
905907
}
906908
}
907909

@@ -1049,6 +1051,7 @@ impl Instruction {
10491051
GetAwaitable => w!(GetAwaitable),
10501052
GetAIter => w!(GetAIter),
10511053
GetANext => w!(GetANext),
1054+
EndAsyncFor => w!(EndAsyncFor),
10521055
MapAdd { i } => w!(MapAdd, i),
10531056
}
10541057
}

‎compiler/src/compile.rs‎

Lines changed: 12 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1376,9 +1376,8 @@ impl Compiler {
13761376
// The thing iterated:
13771377
self.compile_expression(iter)?;
13781378

1379-
if is_async {
1379+
let check_asynciter_block = if is_async {
13801380
let check_asynciter_block = self.new_block();
1381-
let body_block = self.new_block();
13821381

13831382
self.emit(Instruction::GetAIter);
13841383

@@ -1391,21 +1390,8 @@ impl Compiler {
13911390
self.emit(Instruction::YieldFrom);
13921391
self.compile_store(target)?;
13931392
self.emit(Instruction::PopBlock);
1394-
self.emit(Instruction::Jump { target: body_block });
13951393

1396-
self.switch_to_block(check_asynciter_block);
1397-
self.emit(Instruction::Duplicate);
1398-
let stopasynciter = self.name("StopAsyncIteration");
1399-
self.emit(Instruction::LoadGlobal(stopasynciter));
1400-
self.emit(Instruction::CompareOperation {
1401-
op: bytecode::ComparisonOperator::ExceptionMatch,
1402-
});
1403-
self.emit(Instruction::JumpIfTrue { target: else_block });
1404-
self.emit(Instruction::Raise {
1405-
kind: bytecode::RaiseKind::Reraise,
1406-
});
1407-
1408-
self.switch_to_block(body_block);
1394+
Some(check_asynciter_block)
14091395
} else {
14101396
// Retrieve Iterator
14111397
self.emit(Instruction::GetIter);
@@ -1415,24 +1401,27 @@ impl Compiler {
14151401

14161402
// Start of loop iteration, set targets:
14171403
self.compile_store(target)?;
1418-
}
1404+
None
1405+
};
14191406

14201407
let was_in_loop = self.ctx.loop_data;
14211408
self.ctx.loop_data = Some((for_block, after_block));
14221409
self.compile_statements(body)?;
14231410
self.ctx.loop_data = was_in_loop;
14241411
self.emit(Instruction::Jump { target: for_block });
14251412

1413+
if let Some(check_asynciter_block) = check_asynciter_block {
1414+
self.switch_to_block(check_asynciter_block);
1415+
self.emit(Instruction::EndAsyncFor);
1416+
}
1417+
14261418
self.switch_to_block(else_block);
14271419
self.emit(Instruction::PopBlock);
14281420
if let Some(orelse) = orelse {
14291421
self.compile_statements(orelse)?;
14301422
}
14311423

14321424
self.switch_to_block(after_block);
1433-
if is_async {
1434-
self.emit(Instruction::Pop);
1435-
}
14361425

14371426
Ok(())
14381427
}
@@ -2169,14 +2158,13 @@ impl Compiler {
21692158
);
21702159
let arg0 = self.varname(".0");
21712160

2161+
// if this is a generator expression, we need to insert a LoadConst(None) before we return;
2162+
// the other kinds load their collection types below
21722163
let is_genexpr = matches!(kind, ast::ComprehensionKind::GeneratorExpression { .. });
21732164

21742165
// Create empty object of proper type:
21752166
match kind {
2176-
ast::ComprehensionKind::GeneratorExpression { .. } => {
2177-
// pop the None that was passed to kickstart the generator
2178-
self.emit(Instruction::Pop);
2179-
}
2167+
ast::ComprehensionKind::GeneratorExpression { .. } => {}
21802168
ast::ComprehensionKind::List { .. } => {
21812169
self.emit(Instruction::BuildList {
21822170
size: 0,

‎compiler/src/ir.rs‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -165,10 +165,7 @@ impl CodeInfo {
165165
let mut maxdepth = 0u32;
166166
let mut stack = Vec::with_capacity(self.blocks.len());
167167
let mut startdepths = vec![u32::MAX; self.blocks.len()];
168-
// if it's a generator or a coroutine, it starts with something already on the stack
169-
startdepths[0] =
170-
self.flags
171-
.intersects(CodeFlags::IS_GENERATOR | CodeFlags::IS_COROUTINE) as u32;
168+
startdepths[0] = 0;
172169
stack.push(Label(0));
173170
'process_blocks: while let Some(block) = stack.pop() {
174171
let mut depth = startdepths[block.0 as usize];

‎vm/src/coroutine.rs‎

Lines changed: 18 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ pub struct Coro {
3838
pub closed: AtomicCell<bool>,
3939
running: AtomicCell<bool>,
4040
exceptions: PyMutex<Vec<PyBaseExceptionRef>>,
41-
started: AtomicCell<bool>,
4241
variant: Variant,
4342
name: PyMutex<PyStrRef>,
4443
}
@@ -50,7 +49,6 @@ impl Coro {
5049
closed: AtomicCell::new(false),
5150
running: AtomicCell::new(false),
5251
exceptions: PyMutex::new(vec![]),
53-
started: AtomicCell::new(false),
5452
variant,
5553
name: PyMutex::new(name),
5654
}
@@ -67,34 +65,37 @@ impl Coro {
6765
where
6866
F: FnOnce(FrameRef) -> PyResult<ExecutionResult>,
6967
{
70-
self.running.store(true);
71-
let curr_exception_stack_len = vm.exceptions.borrow().len();
72-
vm.exceptions
73-
.borrow_mut()
74-
.append(&mut self.exceptions.lock());
68+
if self.running.compare_exchange(false, true).is_err() {
69+
return Err(vm.new_value_error(format!("{} already executing", self.variant.name())));
70+
}
71+
let curr_exception_stack_len;
72+
{
73+
let mut vm_excs = vm.exceptions.borrow_mut();
74+
curr_exception_stack_len = vm_excs.len();
75+
vm_excs.append(&mut self.exceptions.lock());
76+
}
7577
let result = vm.with_frame(self.frame.clone(), func);
76-
std::mem::swap(
77-
&mut *self.exceptions.lock(),
78-
&mut vm
79-
.exceptions
80-
.borrow_mut()
81-
.split_off(curr_exception_stack_len),
82-
);
78+
self.exceptions
79+
.lock()
80+
.extend(vm.exceptions.borrow_mut().drain(curr_exception_stack_len..));
8381
self.running.store(false);
84-
self.started.store(true);
8582
result
8683
}
8784

8885
pub fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
8986
if self.closed.load() {
9087
return Err(vm.new_exception_empty(self.variant.stop_iteration(vm)));
9188
}
92-
if !self.started.load() && !vm.is_none(&value) {
89+
let value = if self.frame.lasti() > 0 {
90+
Some(value)
91+
} else if !vm.is_none(&value) {
9392
return Err(vm.new_type_error(format!(
9493
"can't send non-None value to a just-started {}",
9594
self.variant.name()
9695
)));
97-
}
96+
} else {
97+
None
98+
};
9899
let result = self.run_with_context(vm, |f| f.resume(value, vm));
99100
self.maybe_close(&result);
100101
match result {
@@ -155,9 +156,6 @@ impl Coro {
155156
}
156157
}
157158

158-
pub fn started(&self) -> bool {
159-
self.started.load()
160-
}
161159
pub fn running(&self) -> bool {
162160
self.running.load()
163161
}

‎vm/src/frame.rs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -254,11 +254,13 @@ impl FrameRef {
254254

255255
pub(crate) fn resume(
256256
&self,
257-
value: PyObjectRef,
257+
value: Option<PyObjectRef>,
258258
vm: &VirtualMachine,
259259
) -> PyResult<ExecutionResult> {
260260
self.with_exec(|mut exec| {
261-
exec.push_value(value);
261+
if let Some(value) = value {
262+
exec.push_value(value)
263+
}
262264
exec.run(vm)
263265
})
264266
}
@@ -811,6 +813,15 @@ impl ExecutingFrame<'_> {
811813
self.push_value(awaitable);
812814
Ok(None)
813815
}
816+
bytecode::Instruction::EndAsyncFor => {
817+
let exc = self.pop_value();
818+
if exc.isinstance(&vm.ctx.exceptions.stop_async_iteration) {
819+
vm.pop_exception().expect("Should have exception in stack");
820+
Ok(None)
821+
} else {
822+
Err(exc.downcast().unwrap())
823+
}
824+
}
814825
bytecode::Instruction::ForIter { target } => self.execute_for_iter(vm, *target),
815826
bytecode::Instruction::MakeFunction(flags) => self.execute_make_function(vm, *flags),
816827
bytecode::Instruction::CallFunctionPositional { nargs } => {

0 commit comments

Comments
 (0)