Skip to content

Commit 91c14d1

Browse files
1ndahous31ndahous3
committed
Stream deque inputs and release evicted items outside locks
Assisted-by: Codex:GPT-6
1 parent f871707 commit 91c14d1

2 files changed

Lines changed: 219 additions & 79 deletions

File tree

‎crates/vm/src/stdlib/_collections.rs‎

Lines changed: 94 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ mod _collections {
1111
VirtualMachine, atomic_func,
1212
builtins::{
1313
IterStatus::{Active, Exhausted},
14-
PositionIterInternal, PyDict, PyGenericAlias, PyInt, PyStr, PyTuple, PyType, PyTypeRef,
15-
locked_step,
14+
PositionIterInternal, PyDict, PyGenericAlias, PyInt, PyList, PyStr, PyTuple, PyType,
15+
PyTypeRef, locked_step,
1616
},
1717
common::lock::{PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard},
1818
convert::ToPyObject,
@@ -90,8 +90,43 @@ mod _collections {
9090
self.deque.write()
9191
}
9292

93-
fn is_over_maxlen(&self, deque: &VecDeque<PyObjectRef>) -> bool {
94-
self.maxlen.is_some_and(|maxlen| deque.len() > maxlen)
93+
fn append_item(&self, item: PyObjectRef, left: bool, maxlen: Option<usize>) {
94+
let removed = {
95+
let mut deque = self.borrow_deque_mut();
96+
if left {
97+
deque.push_front(item);
98+
} else {
99+
deque.push_back(item);
100+
}
101+
self.state.fetch_add(1);
102+
// Trim after pushing, so that a `maxlen` of zero drops what just
103+
// arrived instead of popping from an empty deque and keeping it.
104+
if maxlen.is_some_and(|maxlen| deque.len() > maxlen) {
105+
if left {
106+
deque.pop_back()
107+
} else {
108+
deque.pop_front()
109+
}
110+
} else {
111+
None
112+
}
113+
};
114+
drop(removed);
115+
}
116+
117+
fn extend_items(
118+
&self,
119+
items: impl Iterator<Item = PyResult>,
120+
left: bool,
121+
maxlen: Option<usize>,
122+
) -> PyResult<()> {
123+
for item in items {
124+
let item = item?;
125+
if maxlen != Some(0) {
126+
self.append_item(item, left, maxlen);
127+
}
128+
}
129+
Ok(())
95130
}
96131
}
97132

@@ -110,24 +145,12 @@ mod _collections {
110145
impl PyDeque {
111146
#[pymethod]
112147
fn append(&self, item: PyObjectRef) {
113-
self.state.fetch_add(1);
114-
let mut deque = self.borrow_deque_mut();
115-
deque.push_back(item);
116-
// Trim after pushing, so that a `maxlen` of zero drops what just
117-
// arrived instead of popping from an empty deque and keeping it.
118-
if self.is_over_maxlen(&deque) {
119-
deque.pop_front();
120-
}
148+
self.append_item(item, false, self.maxlen);
121149
}
122150

123151
#[pymethod]
124152
fn appendleft(&self, item: PyObjectRef) {
125-
self.state.fetch_add(1);
126-
let mut deque = self.borrow_deque_mut();
127-
deque.push_front(item);
128-
if self.is_over_maxlen(&deque) {
129-
deque.pop_back();
130-
}
153+
self.append_item(item, true, self.maxlen);
131154
}
132155

133156
#[pymethod]
@@ -160,54 +183,66 @@ mod _collections {
160183

161184
#[pymethod]
162185
fn extend(&self, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
163-
self._extend(&iterable, vm)
186+
self._extend(&iterable, false, vm)
164187
}
165188

166-
fn _extend(&self, iter: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
167-
self.state.fetch_add(1);
168-
let max_len = self.maxlen;
169-
let mut elements: Vec<PyObjectRef> = iter.try_to_value(vm)?;
170-
if let Some(max_len) = max_len {
171-
if max_len > elements.len() {
189+
fn _extend(&self, iterable: &PyObject, left: bool, vm: &VirtualMachine) -> PyResult<()> {
190+
if iterable
191+
.downcast_ref::<Self>()
192+
.is_some_and(|other| core::ptr::eq(self, &**other))
193+
{
194+
// Snapshot self-extension through the iterator protocol so
195+
// subclass overrides of __iter__ are still respected.
196+
let elements: Vec<PyObjectRef> = iterable.try_to_value(vm)?;
197+
// Retain the snapshot until extension finishes, as a temporary list does.
198+
return self.extend_items(elements.iter().cloned().map(Ok), left, self.maxlen);
199+
}
200+
201+
let maxlen = self.maxlen;
202+
let class = iterable.class();
203+
if maxlen.is_none()
204+
&& (class.is(vm.ctx.types.list_type) || class.is(vm.ctx.types.tuple_type))
205+
{
206+
// Exact containers cannot run Python while being copied, and
207+
// an unbounded deque cannot evict items. Fill under one lock.
208+
let elements: Vec<PyObjectRef> = iterable.try_to_value(vm)?;
209+
let count = elements.len();
210+
if count != 0 {
172211
let mut deque = self.borrow_deque_mut();
173-
let drain_until = deque.len().saturating_sub(max_len - elements.len());
174-
deque.drain(..drain_until);
175-
} else {
176-
self.borrow_deque_mut().clear();
177-
elements.drain(..(elements.len() - max_len));
212+
deque.reserve(count);
213+
self.state.fetch_add(count);
214+
if left {
215+
for item in elements {
216+
deque.push_front(item);
217+
}
218+
} else {
219+
deque.extend(elements);
220+
}
178221
}
222+
return Ok(());
223+
}
224+
225+
if class.is(vm.ctx.types.tuple_type) {
226+
let tuple = iterable.downcast_ref::<PyTuple>().unwrap();
227+
self.extend_items(tuple.as_slice().iter().cloned().map(Ok), left, maxlen)
228+
} else if class.is(vm.ctx.types.list_type) {
229+
let list = iterable.downcast_ref::<PyList>().unwrap();
230+
let mut index = 0;
231+
let items = core::iter::from_fn(|| {
232+
// Evicted elements can mutate the source list in __del__.
233+
let item = list.borrow_vec().get(index).cloned();
234+
index += 1;
235+
item.map(Ok)
236+
});
237+
self.extend_items(items, left, maxlen)
238+
} else {
239+
self.extend_items(iterable.get_iter(vm)?.into_iter(vm), left, maxlen)
179240
}
180-
self.borrow_deque_mut().extend(elements);
181-
Ok(())
182241
}
183242

184243
#[pymethod]
185244
fn extendleft(&self, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
186-
let max_len = self.maxlen;
187-
let mut elements: Vec<PyObjectRef> = iterable.try_to_value(vm)?;
188-
elements.reverse();
189-
if let Some(max_len) = max_len {
190-
elements.truncate(max_len);
191-
}
192-
if elements.is_empty() {
193-
return Ok(());
194-
}
195-
let mut created = VecDeque::from(elements);
196-
let discarded = {
197-
let mut deque = self.borrow_deque_mut();
198-
let discarded = if let Some(max_len) = max_len {
199-
let retained = (max_len - created.len()).min(deque.len());
200-
deque.split_off(retained)
201-
} else {
202-
VecDeque::new()
203-
};
204-
created.append(&mut deque);
205-
core::mem::swap(&mut created, &mut deque);
206-
self.state.fetch_add(1);
207-
discarded
208-
};
209-
drop(discarded);
210-
Ok(())
245+
self._extend(&iterable, true, vm)
211246
}
212247

213248
#[pymethod]
@@ -557,27 +592,7 @@ mod _collections {
557592
let Some(iterable) = iterable.into_option() else {
558593
return Ok(());
559594
};
560-
let maxlen = zelf.maxlen;
561-
let mut elements: Vec<PyObjectRef> = iterable.try_to_value(vm)?;
562-
if let Some(maxlen) = maxlen {
563-
elements.drain(..elements.len().saturating_sub(maxlen));
564-
}
565-
if elements.is_empty() {
566-
return Ok(());
567-
}
568-
let discarded = {
569-
let mut deque = zelf.borrow_deque_mut();
570-
let discard_until = maxlen.map_or(0, |maxlen| {
571-
deque.len().saturating_sub(maxlen - elements.len())
572-
});
573-
let discarded: Vec<_> = deque.drain(..discard_until).collect();
574-
deque.extend(elements);
575-
zelf.state.fetch_add(1);
576-
discarded
577-
};
578-
drop(discarded);
579-
580-
Ok(())
595+
zelf._extend(&iterable, false, vm)
581596
}
582597
}
583598

@@ -626,7 +641,7 @@ mod _collections {
626641

627642
inplace_concat: atomic_func!(|seq, other, vm| {
628643
let zelf = PyDeque::sequence_downcast(seq);
629-
zelf._extend(other, vm)?;
644+
zelf._extend(other, false, vm)?;
630645
Ok(zelf.to_owned().into())
631646
}),
632647

‎extra_tests/snippets/stdlib_collections_deque.py‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -414,3 +414,128 @@ def __del__(self):
414414
test_deque_bulk_noops_preserve_iterators()
415415
test_deque_bulk_mutation_releases_lock_before_finalizing()
416416
test_deque_reinitialization_order_and_limit()
417+
418+
419+
def test_deque_streaming_preserves_prefix_on_error():
420+
for method in ("extend", "extendleft", "__iadd__", "__init__"):
421+
d = deque([10], maxlen=2)
422+
seen = []
423+
424+
def source():
425+
yield 1
426+
seen.append(list(d))
427+
yield 2
428+
raise ValueError("source failed")
429+
430+
if method == "__init__":
431+
assert_raises(ValueError, d.__init__, source(), 2)
432+
expected_seen = [1]
433+
else:
434+
assert_raises(ValueError, getattr(d, method), source())
435+
expected_seen = [1, 10] if method == "extendleft" else [10, 1]
436+
assert seen == [expected_seen], (method, seen)
437+
assert list(d) == ([2, 1] if method == "extendleft" else [1, 2])
438+
439+
440+
def test_deque_streaming_keeps_only_bounded_items():
441+
for operation in ("constructor", "extend", "extendleft"):
442+
for maxlen in (0, 3):
443+
444+
class Item:
445+
live = 0
446+
peak = 0
447+
448+
def __init__(self):
449+
Item.live += 1
450+
Item.peak = max(Item.peak, Item.live)
451+
452+
def __del__(self):
453+
Item.live -= 1
454+
455+
source = (Item() for _ in range(20))
456+
if operation == "constructor":
457+
d = deque(source, maxlen=maxlen)
458+
else:
459+
d = deque(maxlen=maxlen)
460+
getattr(d, operation)(source)
461+
assert Item.peak <= maxlen + 1, (operation, maxlen, Item.peak)
462+
assert Item.live == maxlen, (operation, maxlen, Item.live)
463+
del d
464+
assert Item.live == 0
465+
466+
467+
def test_deque_self_extension_uses_subclass_iterator():
468+
class CustomDeque(deque):
469+
def __iter__(self):
470+
return iter((7, 8))
471+
472+
for method in ("extend", "extendleft", "__iadd__"):
473+
d = CustomDeque([1, 2])
474+
getattr(d, method)(d)
475+
expected = [8, 7, 1, 2] if method == "extendleft" else [1, 2, 7, 8]
476+
assert list(deque.__iter__(d)) == expected, method
477+
478+
d = deque([1, 2])
479+
assert_raises(RuntimeError, d.extend, iter(d))
480+
assert list(d) == [1, 2, 1]
481+
482+
483+
def test_deque_extend_noops_preserve_iterators():
484+
d = deque([1, 2])
485+
iterator = iter(d)
486+
d.extend(iter(()))
487+
assert_raises(TypeError, d.extend, None)
488+
assert list(iterator) == [1, 2]
489+
490+
d = deque(maxlen=0)
491+
iterator = iter(d)
492+
d.extend(iter((1, 2)))
493+
assert list(iterator) == []
494+
495+
496+
def test_deque_self_extension_retains_snapshot_until_complete():
497+
for method in ("extend", "extendleft"):
498+
seen = []
499+
500+
class Item:
501+
def __init__(self, value):
502+
self.value = value
503+
504+
def __del__(self):
505+
seen.append([item.value for item in deque.__iter__(d)])
506+
507+
class CustomDeque(deque):
508+
def __iter__(self):
509+
return (Item(value) for value in range(3))
510+
511+
d = CustomDeque(maxlen=1)
512+
getattr(d, method)(d)
513+
assert seen == [[2], [2]], (method, seen)
514+
last = d[0]
515+
d.clear()
516+
del last
517+
518+
519+
def test_deque_eviction_can_mutate_source_list():
520+
for method in ("extend", "extendleft"):
521+
d = deque(maxlen=1)
522+
source = [1, 2]
523+
seen = []
524+
525+
class Item:
526+
def __del__(self):
527+
seen.append(list(d))
528+
source[1:] = [3, 4]
529+
530+
d.append(Item())
531+
getattr(d, method)(source)
532+
assert seen == [[1]], (method, seen)
533+
assert list(d) == [4], (method, d)
534+
535+
536+
test_deque_streaming_preserves_prefix_on_error()
537+
test_deque_streaming_keeps_only_bounded_items()
538+
test_deque_self_extension_uses_subclass_iterator()
539+
test_deque_extend_noops_preserve_iterators()
540+
test_deque_self_extension_retains_snapshot_until_complete()
541+
test_deque_eviction_can_mutate_source_list()

0 commit comments

Comments
 (0)