Track deque mutations and release removed items outside locks (#8847)

Assisted-by: Codex:GPT-6

crates/vm/src/stdlib/_collections.rs

@@ -183,21 +183,27 @@ mod _collections {
             let max_len = self.maxlen;
             let mut elements: Vec<PyObjectRef> = iterable.try_to_value(vm)?;
             elements.reverse();
-
             if let Some(max_len) = max_len {
-                if max_len > elements.len() {
-                    let mut deque = self.borrow_deque_mut();
-                    let truncate_until = max_len - elements.len();
-                    deque.truncate(truncate_until);
-                } else {
-                    self.borrow_deque_mut().clear();
-                    elements.truncate(max_len);
-                }
+                elements.truncate(max_len);
+            }
+            if elements.is_empty() {
+                return Ok(());
             }
             let mut created = VecDeque::from(elements);
-            let mut borrowed = self.borrow_deque_mut();
-            created.append(&mut borrowed);
-            core::mem::swap(&mut created, &mut borrowed);
+            let discarded = {
+                let mut deque = self.borrow_deque_mut();
+                let discarded = if let Some(max_len) = max_len {
+                    let retained = (max_len - created.len()).min(deque.len());
+                    deque.split_off(retained)
+                } else {
+                    VecDeque::new()
+                };
+                created.append(&mut deque);
+                core::mem::swap(&mut created, &mut deque);
+                self.state.fetch_add(1);
+                discarded
+            };
+            drop(discarded);
             Ok(())
         }
 
@@ -268,19 +274,37 @@ mod _collections {
         }
 
         #[pymethod]
-        fn remove(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        fn remove(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
             let start_state = self.state.load();
-            let index = self.mut_index(vm, &value)?;
-
-            if start_state != self.state.load() {
-                Err(vm.new_index_error("deque mutated during remove()."))
-            } else if let Some(index) = index.into() {
-                let mut deque = self.borrow_deque_mut();
-                self.state.fetch_add(1);
-                Ok(deque.remove(index).unwrap())
-            } else {
-                Err(vm.new_value_error("deque.remove(x): x not in deque"))
+            let len = self.borrow_deque().len();
+            let mutated = || vm.new_index_error("deque mutated during iteration");
+            for index in 0..len {
+                let item = self
+                    .borrow_deque()
+                    .get(index)
+                    .cloned()
+                    .ok_or_else(mutated)?;
+                let equal = item.rich_compare_bool(&value, PyComparisonOp::Eq, vm)?;
+                // Releasing the comparison reference can also run a finalizer.
+                drop(item);
+                if start_state != self.state.load() {
+                    return Err(mutated());
+                }
+                if equal {
+                    let removed = {
+                        let mut deque = self.borrow_deque_mut();
+                        if start_state != self.state.load() {
+                            return Err(mutated());
+                        }
+                        let removed = deque.remove(index).ok_or_else(mutated)?;
+                        self.state.fetch_add(1);
+                        removed
+                    };
+                    drop(removed);
+                    return Ok(());
+                }
             }
+            Err(vm.new_value_error("deque.remove(x): x not in deque"))
         }
 
         #[pymethod]
@@ -333,10 +357,18 @@ mod _collections {
         }
 
         fn __delitem__(&self, idx: isize, vm: &VirtualMachine) -> PyResult<()> {
-            let mut deque = self.borrow_deque_mut();
-            idx.wrapped_at(deque.len())
-                .and_then(|i| deque.remove(i).map(drop))
-                .ok_or_else(|| vm.new_index_error("deque index out of range"))
+            let removed = {
+                let mut deque = self.borrow_deque_mut();
+                let removed = idx
+                    .wrapped_at(deque.len())
+                    .and_then(|i| deque.remove(i))
+                    .ok_or_else(|| vm.new_index_error("deque index out of range"))?;
+                self.state.fetch_add(1);
+                removed
+            };
+            // Finalizers can mutate the deque, so release the lock first.
+            drop(removed);
+            Ok(())
         }
 
         fn __contains__(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<bool> {
@@ -381,8 +413,16 @@ mod _collections {
         }
 
         fn __imul__(zelf: PyRef<Self>, n: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
-            let mul_deque = zelf._mul(n, vm)?;
-            *zelf.borrow_deque_mut() = mul_deque;
+            if n == 1 || zelf.borrow_deque().is_empty() {
+                return Ok(zelf);
+            }
+            let mut mul_deque = zelf._mul(n, vm)?;
+            {
+                let mut deque = zelf.borrow_deque_mut();
+                core::mem::swap(&mut *deque, &mut mul_deque);
+                zelf.state.fetch_add(1);
+            }
+            drop(mul_deque);
             Ok(zelf)
         }
 
@@ -490,25 +530,11 @@ mod _collections {
                 None
             };
 
-            // retrieve elements first to not to make too huge lock
-            let elements = iterable
-                .into_option()
-                .map(|iter| {
-                    let mut elements: Vec<PyObjectRef> = iter.try_to_value(vm)?;
-                    if let Some(maxlen) = maxlen {
-                        elements.drain(..elements.len().saturating_sub(maxlen));
-                    }
-                    Ok(elements)
-                })
-                .transpose()?;
-
             // SAFETY: This is hacky part for read-only field
             // Because `maxlen` is only mutated from __init__. We can abuse the lock of deque to ensure this is locked enough.
             // If we make a single lock of deque not only for extend but also for setting maxlen, it will be safe.
-            {
+            let removed = {
                 let mut deque = zelf.borrow_deque_mut();
-                // Clear any previous data present.
-                deque.clear();
                 unsafe {
                     // `maxlen` is better to be defined as UnsafeCell in common practice,
                     // but then more type works without any safety benefits
@@ -516,10 +542,37 @@ mod _collections {
                         &zelf.maxlen as *const _ as *const core::cell::UnsafeCell<Option<usize>>;
                     *(*unsafe_maxlen).get() = maxlen;
                 }
-                if let Some(elements) = elements {
-                    deque.extend(elements);
+                let removed = core::mem::take(&mut *deque);
+                if !removed.is_empty() {
+                    zelf.state.fetch_add(1);
                 }
+                removed
+            };
+            // Release old elements before iterating; their finalizers can add new ones.
+            drop(removed);
+
+            let Some(iterable) = iterable.into_option() else {
+                return Ok(());
+            };
+            let maxlen = zelf.maxlen;
+            let mut elements: Vec<PyObjectRef> = iterable.try_to_value(vm)?;
+            if let Some(maxlen) = maxlen {
+                elements.drain(..elements.len().saturating_sub(maxlen));
             }
+            if elements.is_empty() {
+                return Ok(());
+            }
+            let discarded = {
+                let mut deque = zelf.borrow_deque_mut();
+                let discard_until = maxlen.map_or(0, |maxlen| {
+                    deque.len().saturating_sub(maxlen - elements.len())
+                });
+                let discarded: Vec<_> = deque.drain(..discard_until).collect();
+                deque.extend(elements);
+                zelf.state.fetch_add(1);
+                discarded
+            };
+            drop(discarded);
 
             Ok(())
         }

extra_tests/snippets/stdlib_collections_deque.py

@@ -151,3 +151,262 @@ d.append("d")
 assert list(d) == ["b", "c", "d"]
 d.appendleft("z")
 assert list(d) == ["z", "b", "c"]
+
+
+def assert_deque_error(error_type, message, function, *args):
+    with assert_raises(error_type) as caught:
+        function(*args)
+    assert type(caught.exception) is error_type, type(caught.exception)
+    assert str(caught.exception) == message, str(caught.exception)
+
+
+def test_deque_remove_detects_deletion_during_comparison():
+    for match in (True, False):
+        d = deque()
+
+        class Item:
+            def __eq__(self, other):
+                del d[0]
+                return match
+
+        d.append(Item())
+        assert_deque_error(IndexError, "deque mutated during iteration", d.remove, None)
+        assert not d
+
+
+def test_deque_remove_stops_after_comparison_appends():
+    d = deque()
+
+    class Item:
+        calls = 0
+
+        def __eq__(self, other):
+            self.calls += 1
+            assert self.calls == 1, "remove continued after mutation"
+            d.append(self)
+            return False
+
+    item = Item()
+    d.append(item)
+    assert_deque_error(IndexError, "deque mutated during iteration", d.remove, None)
+    assert item.calls == 1 and len(d) == 2
+
+
+def test_deque_remove_preserves_identity_first_match_and_return_value():
+    class Identity:
+        def __eq__(self, other):
+            raise AssertionError("identity should not compare")
+
+    class Unequal:
+        calls = 0
+
+        def __eq__(self, other):
+            self.calls += 1
+            return False
+
+    identity, unequal = Identity(), Unequal()
+    d = deque([identity, unequal, 1, 2, 1])
+    assert d.remove(identity) is None
+    assert d.remove(1) is None
+    assert unequal.calls == 1
+    assert list(d) == [unequal, 2, 1]
+    assert_deque_error(ValueError, "deque.remove(x): x not in deque", d.remove, 3)
+    assert list(d) == [unequal, 2, 1]
+
+
+def test_deque_deletion_invalidates_iterators_only_after_success():
+    for factory in (iter, reversed):
+        for index in (0, 1, -1):
+            d = deque([1, 2, 3])
+            iterator = factory(d)
+            del d[index]
+            assert_deque_error(
+                RuntimeError, "deque mutated during iteration", next, iterator
+            )
+        d = deque([1, 2, 3])
+        iterator = factory(d)
+
+        def delete(index):
+            del d[index]
+
+        assert_deque_error(IndexError, "deque index out of range", delete, 10)
+        assert list(iterator) == list(factory(d))
+        iterator = factory(d)
+        assert d.remove(2) is None
+        assert_deque_error(
+            RuntimeError, "deque mutated during iteration", next, iterator
+        )
+
+
+def test_deque_remove_preserves_comparison_error_after_mutation():
+    d = deque()
+
+    class Item:
+        def __eq__(self, other):
+            del d[0]
+            raise LookupError("comparison failed")
+
+    d.append(Item())
+    assert_deque_error(LookupError, "comparison failed", d.remove, None)
+    assert not d
+
+
+def test_deque_deletion_releases_lock_before_finalizing_item():
+    d = deque()
+
+    class Item:
+        def __del__(self):
+            d.append("finalized")
+
+    d.append(Item())
+    del d[0]
+    assert list(d) == ["finalized"]
+
+
+def test_deque_remove_detects_mutation_from_compared_item_finalizer():
+    for match in (True, False):
+        d = deque()
+
+        class Item:
+            def __eq__(self, other):
+                d[0] = "replacement"
+                return match
+
+            def __del__(self):
+                d.append("finalized")
+
+        d.append(Item())
+        assert_deque_error(IndexError, "deque mutated during iteration", d.remove, None)
+        assert list(d) == ["replacement", "finalized"]
+
+
+test_deque_remove_detects_deletion_during_comparison()
+test_deque_remove_stops_after_comparison_appends()
+test_deque_remove_preserves_identity_first_match_and_return_value()
+test_deque_deletion_invalidates_iterators_only_after_success()
+test_deque_remove_preserves_comparison_error_after_mutation()
+test_deque_deletion_releases_lock_before_finalizing_item()
+test_deque_remove_detects_mutation_from_compared_item_finalizer()
+
+
+def test_deque_remove_detects_bulk_mutation():
+    operations = (
+        lambda d: d.extendleft([1]),
+        lambda d: d.__imul__(0),
+        lambda d: d.__imul__(2),
+        lambda d: d.__init__([1]),
+    )
+    for operation in operations:
+        for match in (True, False):
+            d = deque()
+
+            class Item:
+                calls = 0
+
+                def __eq__(self, other):
+                    self.calls += 1
+                    assert self.calls == 1, "remove continued after mutation"
+                    operation(d)
+                    return match
+
+            d.append(Item())
+            assert_deque_error(
+                IndexError, "deque mutated during iteration", d.remove, None
+            )
+
+
+def test_deque_bulk_mutation_invalidates_iterators():
+    operations = (
+        lambda d: d.extendleft([9]),
+        lambda d: d.__imul__(0),
+        lambda d: d.__imul__(2),
+        lambda d: d.__init__([9]),
+    )
+    for factory in (iter, reversed):
+        for maxlen in (None, 1, 3):
+            for operation in operations:
+                d = deque([1, 2], maxlen=maxlen)
+                iterator = factory(d)
+                operation(d)
+                assert_deque_error(
+                    RuntimeError, "deque mutated during iteration", next, iterator
+                )
+
+
+def test_deque_bulk_noops_preserve_iterators():
+    for factory in (iter, reversed):
+        for items in ([], [1, 2]):
+            for operation in (lambda d: d.extendleft([]), lambda d: d.__imul__(1)):
+                d = deque(items)
+                iterator = factory(d)
+                operation(d)
+                assert list(iterator) == list(factory(d))
+        for count in (-1, 0, 2):
+            d = deque()
+            iterator = factory(d)
+            d *= count
+            assert list(iterator) == []
+        d = deque(maxlen=0)
+        iterator = factory(d)
+        d.extendleft([1, 2])
+        assert list(iterator) == []
+        d = deque()
+        iterator = factory(d)
+        d.__init__([])
+        assert list(iterator) == []
+
+
+def test_deque_bulk_mutation_releases_lock_before_finalizing():
+    operations = (
+        (lambda d: d.extendleft([9]), [9], ["finalized"]),
+        (lambda d: d.__imul__(0), [], ["finalized"]),
+        (lambda d: d.__init__([9]), [], ["finalized", 9]),
+    )
+    for operation, observed, expected in operations:
+        d = deque(maxlen=1)
+        snapshots = []
+
+        class Item:
+            def __del__(self):
+                snapshots.append(list(d))
+                d.append("finalized")
+
+        d.append(Item())
+        operation(d)
+        assert snapshots == [observed], snapshots
+        assert list(d) == expected, d
+
+
+def test_deque_reinitialization_order_and_limit():
+    d = deque([1, 2, 3])
+    d.__init__(d)
+    assert not d
+    d = deque([1, 2, 3])
+    iterator = iter(d)
+    assert_raises(ValueError, d.__init__, [], -1)
+    assert list(iterator) == [1, 2, 3]
+    iterator = iter(d)
+    assert_raises(TypeError, d.__init__, 1)
+    assert not d
+    assert_deque_error(RuntimeError, "deque mutated during iteration", next, iterator)
+
+    for maxlen, expected in ((1, [9]), (2, ["finalized", 9])):
+        d = deque()
+        snapshots = []
+
+        class Item:
+            def __del__(self):
+                snapshots.append((list(d), d.maxlen))
+                d.append("finalized")
+
+        d.append(Item())
+        d.__init__([9], maxlen)
+        assert snapshots == [([], maxlen)], snapshots
+        assert list(d) == expected, d
+
+
+test_deque_remove_detects_bulk_mutation()
+test_deque_bulk_mutation_invalidates_iterators()
+test_deque_bulk_noops_preserve_iterators()
+test_deque_bulk_mutation_releases_lock_before_finalizing()
+test_deque_reinitialization_order_and_limit()