Skip to content

Commit 17011c4

Browse files
committed
apply feedback PyRwLock -> PyMutex & remove AtomicCell<bool> lock field
1 parent aaf1805 commit 17011c4

File tree

1 file changed

+9
-14
lines changed

1 file changed

+9
-14
lines changed

vm/src/stdlib/itertools.rs

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1184,33 +1184,28 @@ mod decl {
11841184
#[derive(Debug)]
11851185
struct PyItertoolsTeeData {
11861186
iterable: PyIter,
1187-
values: PyRwLock<Vec<PyObjectRef>>,
1188-
locked: AtomicCell<bool>,
1187+
values: PyMutex<Vec<PyObjectRef>>,
11891188
}
11901189

11911190
impl PyItertoolsTeeData {
11921191
fn new(iterable: PyIter, _vm: &VirtualMachine) -> PyResult<PyRc<Self>> {
11931192
Ok(PyRc::new(Self {
11941193
iterable,
1195-
values: PyRwLock::new(vec![]),
1196-
locked: AtomicCell::new(false),
1194+
values: PyMutex::new(vec![]),
11971195
}))
11981196
}
11991197

12001198
fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult<PyIterReturn> {
1201-
if self.values.read().len() == index {
1202-
if self.locked.swap(true) {
1203-
return Err(vm.new_runtime_error("cannot re-enter the tee iterator"));
1204-
}
1205-
1206-
let result = self.iterable.next(vm);
1207-
self.locked.store(false);
1199+
let Some(mut values) = self.values.try_lock() else {
1200+
return Err(vm.new_runtime_error("cannot re-enter the tee iterator"));
1201+
};
12081202

1209-
let obj = raise_if_stop!(result?);
1210-
self.values.write().push(obj);
1203+
if values.len() == index {
1204+
let obj = raise_if_stop!(self.iterable.next(vm)?);
1205+
values.push(obj);
12111206
}
12121207

1213-
Ok(PyIterReturn::Return(self.values.read()[index].clone()))
1208+
Ok(PyIterReturn::Return(values[index].clone()))
12141209
}
12151210
}
12161211

0 commit comments

Comments
 (0)