Skip to content

Commit e1035c4

Browse files
committed
MutObjectSequenceOp find using Iteartor trait
to isolate control flow code from next() operation
1 parent 9c2847a commit e1035c4

File tree

3 files changed

+72
-49
lines changed

3 files changed

+72
-49
lines changed

vm/src/builtins/list.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,7 @@ impl PyList {
274274
) -> PyResult<usize> {
275275
let (start, stop) = range.saturate(self.len(), vm)?;
276276
let index = self.mut_index_range(vm, &needle, start..stop)?;
277-
if let Some(index) = index.into() {
277+
if let Some(index) = index {
278278
Ok(index)
279279
} else {
280280
Err(vm.new_value_error(format!("'{}' is not in list", needle.str(vm)?)))
@@ -301,7 +301,7 @@ impl PyList {
301301
fn remove(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
302302
let index = self.mut_index(vm, &needle)?;
303303

304-
if let Some(index) = index.into() {
304+
if let Some(index) = index {
305305
// defer delete out of borrow
306306
Ok(self.borrow_vec_mut().remove(index))
307307
} else {

vm/src/sequence.rs

Lines changed: 68 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,18 @@ use crate::{
55
use optional::Optioned;
66
use std::ops::Range;
77

8-
pub trait MutObjectSequenceOp<'a> {
9-
type Guard;
8+
pub trait MutObjectSequenceOp<'a>: Sized {
9+
type Guard: Sized;
1010

1111
fn do_get(index: usize, guard: &Self::Guard) -> Option<&PyObjectRef>;
1212
fn do_lock(&'a self) -> Self::Guard;
1313

1414
fn mut_count(&'a self, vm: &VirtualMachine, needle: &PyObject) -> PyResult<usize> {
15+
let mut find_iter = self._mut_find(vm, needle, 0..isize::MAX as usize);
1516
let mut count = 0;
16-
self._mut_iter_equal_skeleton::<_, false>(vm, needle, 0..isize::MAX as usize, || {
17-
count += 1
18-
})?;
17+
while (find_iter.next().transpose()?).is_some() {
18+
count += 1;
19+
}
1920
Ok(count)
2021
}
2122

@@ -24,68 +25,90 @@ pub trait MutObjectSequenceOp<'a> {
2425
vm: &VirtualMachine,
2526
needle: &PyObject,
2627
range: Range<usize>,
27-
) -> PyResult<Optioned<usize>> {
28-
self._mut_iter_equal_skeleton::<_, true>(vm, needle, range, || {})
28+
) -> PyResult<Option<usize>> {
29+
self._mut_find(vm, needle, range).next().transpose()
2930
}
3031

31-
fn mut_index(&'a self, vm: &VirtualMachine, needle: &PyObject) -> PyResult<Optioned<usize>> {
32+
fn mut_index(&'a self, vm: &VirtualMachine, needle: &PyObject) -> PyResult<Option<usize>> {
3233
self.mut_index_range(vm, needle, 0..isize::MAX as usize)
3334
}
3435

3536
fn mut_contains(&'a self, vm: &VirtualMachine, needle: &PyObject) -> PyResult<bool> {
3637
self.mut_index(vm, needle).map(|x| x.is_some())
3738
}
3839

39-
fn _mut_iter_equal_skeleton<F, const SHORT: bool>(
40+
fn _mut_find<'b, 'vm>(
4041
&'a self,
41-
vm: &VirtualMachine,
42-
needle: &PyObject,
42+
vm: &'vm VirtualMachine,
43+
needle: &'b PyObject,
4344
range: Range<usize>,
44-
mut f: F,
45-
) -> PyResult<Optioned<usize>>
46-
where
47-
F: FnMut(),
48-
{
49-
let mut borrower = None;
50-
let mut i = range.start;
51-
52-
let index = loop {
53-
if i >= range.end {
54-
break Optioned::<usize>::none();
55-
}
56-
let guard = if let Some(x) = borrower.take() {
57-
x
58-
} else {
59-
self.do_lock()
60-
};
45+
) -> MutObjectSequenceFindIter<'a, 'b, 'vm, Self> {
46+
MutObjectSequenceFindIter {
47+
seq: self,
48+
needle,
49+
pos: range.start,
50+
end: range.end,
51+
guard: None,
52+
vm,
53+
}
54+
}
55+
}
56+
57+
pub struct MutObjectSequenceFindIter<'a, 'b, 'vm, S: MutObjectSequenceOp<'a>> {
58+
// mutable fields
59+
pos: usize,
60+
guard: Option<S::Guard>,
61+
// immutable fields
62+
seq: &'a S,
63+
needle: &'b PyObject,
64+
end: usize,
65+
vm: &'vm VirtualMachine,
66+
}
6167

62-
let elem = if let Some(x) = Self::do_get(i, &guard) {
68+
impl<'a, 'b, 'vm, S: MutObjectSequenceOp<'a>> MutObjectSequenceFindIter<'a, 'b, 'vm, S> {
69+
#[inline]
70+
fn next_impl(&mut self) -> PyResult<Optioned<usize>> {
71+
loop {
72+
if self.pos >= self.end {
73+
return Ok(Optioned::none());
74+
}
75+
let guard = self.guard.take().unwrap_or_else(|| self.seq.do_lock());
76+
let elem = if let Some(x) = S::do_get(self.pos, &guard) {
6377
x
6478
} else {
65-
break Optioned::<usize>::none();
79+
return Ok(Optioned::none());
6680
};
6781

68-
if elem.is(needle) {
69-
f();
70-
if SHORT {
71-
break Optioned::<usize>::some(i);
72-
}
73-
borrower = Some(guard);
82+
let is_equal = if elem.is(self.needle) {
83+
self.guard = Some(guard);
84+
true
7485
} else {
7586
let elem = elem.clone();
7687
drop(guard);
88+
elem.rich_compare_bool(self.needle, PyComparisonOp::Eq, self.vm)?
89+
};
7790

78-
if elem.rich_compare_bool(needle, PyComparisonOp::Eq, vm)? {
79-
f();
80-
if SHORT {
81-
break Optioned::<usize>::some(i);
82-
}
83-
}
91+
if is_equal {
92+
break;
8493
}
85-
i += 1;
86-
};
8794

88-
Ok(index)
95+
self.pos += 1;
96+
}
97+
98+
let i = self.pos;
99+
self.pos += 1;
100+
Ok(Optioned::some(i))
101+
}
102+
}
103+
104+
impl<'a, 'b, 'vm, S: MutObjectSequenceOp<'a>> Iterator
105+
for MutObjectSequenceFindIter<'a, 'b, 'vm, S>
106+
{
107+
type Item = PyResult<usize>;
108+
109+
#[inline]
110+
fn next(&mut self) -> Option<Self::Item> {
111+
self.next_impl().map(Into::into).transpose()
89112
}
90113
}
91114

vm/src/stdlib/collections.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ mod _collections {
166166
let index = self.mut_index_range(vm, &needle, start..stop)?;
167167
if start_state != self.state.load() {
168168
Err(vm.new_runtime_error("deque mutated during iteration".to_owned()))
169-
} else if let Some(index) = index.into() {
169+
} else if let Some(index) = index {
170170
Ok(index)
171171
} else {
172172
Err(vm.new_value_error(
@@ -227,7 +227,7 @@ mod _collections {
227227

228228
if start_state != self.state.load() {
229229
Err(vm.new_index_error("deque mutated during remove().".to_owned()))
230-
} else if let Some(index) = index.into() {
230+
} else if let Some(index) = index {
231231
let mut deque = self.borrow_deque_mut();
232232
self.state.fetch_add(1);
233233
Ok(deque.remove(index).unwrap())

0 commit comments

Comments
 (0)