Skip to content

Commit 0addc7d

Browse files
coolreader18coolreader18
committed
Don't push something on the stack when starting a generator
1 parent c97d521 commit 0addc7d

4 files changed

Lines changed: 31 additions & 38 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: 9 additions & 20 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
}

‎vm/src/coroutine.rs‎

Lines changed: 10 additions & 18 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
}
@@ -70,28 +68,25 @@ impl Coro {
7068
if self.running.compare_exchange(false, true).is_err() {
7169
return Err(vm.new_value_error(format!("{} already executing", self.variant.name())));
7270
}
73-
let curr_exception_stack_len = vm.exceptions.borrow().len();
74-
vm.exceptions
75-
.borrow_mut()
76-
.append(&mut self.exceptions.lock());
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+
}
7777
let result = vm.with_frame(self.frame.clone(), func);
78-
std::mem::swap(
79-
&mut *self.exceptions.lock(),
80-
&mut vm
81-
.exceptions
82-
.borrow_mut()
83-
.split_off(curr_exception_stack_len),
84-
);
78+
self.exceptions
79+
.lock()
80+
.extend(vm.exceptions.borrow_mut().drain(curr_exception_stack_len..));
8581
self.running.store(false);
86-
self.started.store(true);
8782
result
8883
}
8984

9085
pub fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
9186
if self.closed.load() {
9287
return Err(vm.new_exception_empty(self.variant.stop_iteration(vm)));
9388
}
94-
let value = if self.started.load() {
89+
let value = if self.frame.lasti() > 0 {
9590
Some(value)
9691
} else if !vm.is_none(&value) {
9792
return Err(vm.new_type_error(format!(
@@ -161,9 +156,6 @@ impl Coro {
161156
}
162157
}
163158

164-
pub fn started(&self) -> bool {
165-
self.started.load()
166-
}
167159
pub fn running(&self) -> bool {
168160
self.running.load()
169161
}

‎vm/src/frame.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -813,6 +813,15 @@ impl ExecutingFrame<'_> {
813813
self.push_value(awaitable);
814814
Ok(None)
815815
}
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+
}
816825
bytecode::Instruction::ForIter { target } => self.execute_for_iter(vm, *target),
817826
bytecode::Instruction::MakeFunction(flags) => self.execute_make_function(vm, *flags),
818827
bytecode::Instruction::CallFunctionPositional { nargs } => {

0 commit comments

Comments
 (0)