use std::collections::BTreeMap;
use std::rc::Rc;
use crate::value::{VmError, VmValue};
impl super::super::Vm {
pub(super) fn execute_iter_init(&mut self) -> Result<(), VmError> {
let iterable = self.pop()?;
match iterable {
VmValue::List(items) => {
self.iterators
.push(super::super::IterState::Vec { items, idx: 0 });
}
VmValue::Dict(map) => {
let keys = map.keys().cloned().collect();
self.iterators.push(super::super::IterState::Dict {
entries: map,
keys,
idx: 0,
});
}
VmValue::Set(items) => {
self.iterators
.push(super::super::IterState::Vec { items, idx: 0 });
}
VmValue::Channel(ch) => {
self.iterators.push(super::super::IterState::Channel {
receiver: ch.receiver.clone(),
closed: ch.closed.clone(),
});
}
VmValue::Generator(gen) => {
self.iterators
.push(super::super::IterState::Generator { gen });
}
VmValue::Range(r) => {
let stop = if r.inclusive {
r.end.saturating_add(1)
} else {
r.end
};
let next = r.start;
self.iterators
.push(super::super::IterState::Range { next, stop });
}
VmValue::Iter(handle) => {
self.iterators
.push(super::super::IterState::VmIter { handle });
}
_ => {
self.iterators.push(super::super::IterState::Vec {
items: Rc::new(Vec::new()),
idx: 0,
});
}
}
Ok(())
}
pub(super) async fn execute_iter_next(&mut self) -> Result<(), VmError> {
let frame = self.frames.last_mut().unwrap();
let target = frame.chunk.read_u16(frame.ip) as usize;
frame.ip += 2;
let vm_iter_handle = match self.iterators.last() {
Some(super::super::IterState::VmIter { handle }) => Some(handle.clone()),
_ => None,
};
if let Some(handle) = vm_iter_handle {
let next_val = crate::vm::iter::next_handle(&handle, self).await?;
match next_val {
Some(v) => self.stack.push(v),
None => {
self.iterators.pop();
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
}
}
} else if let Some(state) = self.iterators.last_mut() {
match state {
super::super::IterState::Vec { items, idx } => {
if *idx < items.len() {
let item = items[*idx].clone();
*idx += 1;
self.stack.push(item);
} else {
self.iterators.pop();
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
}
}
super::super::IterState::Dict { entries, keys, idx } => {
if *idx < keys.len() {
let key = &keys[*idx];
let value = entries.get(key).cloned().unwrap_or(VmValue::Nil);
*idx += 1;
self.stack.push(VmValue::Dict(Rc::new(BTreeMap::from([
("key".to_string(), VmValue::String(Rc::from(key.as_str()))),
("value".to_string(), value),
]))));
} else {
self.iterators.pop();
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
}
}
super::super::IterState::Channel { receiver, closed } => {
let rx = receiver.clone();
let is_closed = closed.load(std::sync::atomic::Ordering::Relaxed);
let mut guard = rx.lock().await;
let item = if is_closed {
guard.try_recv().ok()
} else {
guard.recv().await
};
match item {
Some(val) => {
self.stack.push(val);
}
None => {
drop(guard);
self.iterators.pop();
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
}
}
}
super::super::IterState::Range { next, stop } => {
if *next < *stop {
let v = *next;
*next += 1;
self.stack.push(VmValue::Int(v));
} else {
self.iterators.pop();
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
}
}
super::super::IterState::Generator { gen } => {
if gen.done.get() {
self.iterators.pop();
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
} else {
let rx = gen.receiver.clone();
let mut guard = rx.lock().await;
match guard.recv().await {
Some(val) => {
self.stack.push(val);
}
None => {
gen.done.set(true);
drop(guard);
self.iterators.pop();
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
}
}
}
}
super::super::IterState::VmIter { .. } => {
unreachable!("VmIter state handled before this match");
}
}
} else {
let frame = self.frames.last_mut().unwrap();
frame.ip = target;
}
Ok(())
}
pub(super) fn execute_pop_iterator(&mut self) {
self.iterators.pop();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::vm::{IterState, Vm};
fn run_iter_init_test(test: impl std::future::Future<Output = ()>) {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(test);
}
#[test]
fn iter_init_list_keeps_shared_backing_store() {
run_iter_init_test(async {
let items = Rc::new(vec![VmValue::Int(1), VmValue::Int(2)]);
let mut vm = Vm::new();
vm.stack.push(VmValue::List(items.clone()));
vm.execute_iter_init().unwrap();
match vm.iterators.last().unwrap() {
IterState::Vec {
items: iter_items,
idx,
} => {
assert!(Rc::ptr_eq(&items, iter_items));
assert_eq!(*idx, 0);
}
_ => panic!("expected vec iterator state"),
}
});
}
#[test]
fn iter_init_set_keeps_shared_backing_store() {
run_iter_init_test(async {
let items = Rc::new(vec![VmValue::Int(1), VmValue::Int(2)]);
let mut vm = Vm::new();
vm.stack.push(VmValue::Set(items.clone()));
vm.execute_iter_init().unwrap();
match vm.iterators.last().unwrap() {
IterState::Vec {
items: iter_items,
idx,
} => {
assert!(Rc::ptr_eq(&items, iter_items));
assert_eq!(*idx, 0);
}
_ => panic!("expected vec iterator state"),
}
});
}
#[test]
fn iter_init_dict_keeps_shared_entries_and_snapshots_keys() {
run_iter_init_test(async {
let entries = Rc::new(BTreeMap::from([
("a".to_string(), VmValue::Int(1)),
("b".to_string(), VmValue::Int(2)),
]));
let mut vm = Vm::new();
vm.stack.push(VmValue::Dict(entries.clone()));
vm.execute_iter_init().unwrap();
match vm.iterators.last().unwrap() {
IterState::Dict {
entries: iter_entries,
keys,
idx,
} => {
assert!(Rc::ptr_eq(&entries, iter_entries));
assert_eq!(keys.as_slice(), ["a".to_string(), "b".to_string()]);
assert_eq!(*idx, 0);
}
_ => panic!("expected dict iterator state"),
}
});
}
}