Skip to content

Align native sequence iterator state and lifecycle #9701

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 1 commit into from
Nov 5, 2026
Merged

Align native sequence iterator state and lifecycle #9701

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 16 additions & 17 deletions crates/vm/src/builtins/bytearray.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
//! Implementation of the python bytearray object.
use super::native_iter::{NativeIterKind, NativeSequenceIter};
use super::{
PositionIterInternal, PyBytes, PyGenericAlias, PyInt, PyStrRef, PyTuple, PyTupleRef, PyType,
PyTypeRef, iter::builtins_iter, locked_next,
PyBytes, PyGenericAlias, PyInt, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef,
iter::builtins_iter,
};
use crate::{
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
Expand Down Expand Up @@ -1404,7 +1405,7 @@ impl AsNumber for PyByteArray {
impl Iterable for PyByteArray {
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
Ok(PyByteArrayIterator {
internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
internal: NativeSequenceIter::new(zelf, NativeIterKind::ByteArray, 0),
}
.into_pyobject(vm))
}
Expand All @@ -1419,10 +1420,10 @@ impl Representable for PyByteArray {
}
}

#[pyclass(module = false, name = "bytearray_iterator")]
#[pyclass(module = false, name = "bytearray_iterator", traverse)]
#[derive(Debug)]
pub(crate) struct PyByteArrayIterator {
internal: PyMutex<PositionIterInternal<PyByteArrayRef>>,
internal: NativeSequenceIter<PyByteArrayRef>,
}

impl PyPayload for PyByteArrayIterator {
Expand All @@ -1436,36 +1437,34 @@ impl PyPayload for PyByteArrayIterator {
impl Py<PyByteArrayIterator> {
#[pymethod]
fn __length_hint__(&self) -> usize {
self.internal.lock().length_hint(|obj| obj.__len__())
self.internal.length_hint(|obj| obj.__len__())
}
#[pymethod]
fn __reduce__(&self, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
let func = builtins_iter(vm)?;
Ok(self.internal.lock().reduce(
Ok(self.internal.reduce(
func,
|x| x.clone().into(),
|x| x.into(),
|vm| vm.ctx.empty_tuple.clone().into(),
vm,
))
}

#[pymethod]
fn __setstate__(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
self.internal
.lock()
.set_state(&object, |obj, pos| pos.min(obj.__len__()), vm)
self.internal.set_state(&object, |obj| obj.__len__(), vm)
}
}

impl SelfIter for PyByteArrayIterator {}
impl IterNext for PyByteArrayIterator {
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
locked_next(&zelf.internal, |bytearray, pos| {
let buf = bytearray.borrow_buf();
Ok(PyIterReturn::from_result(
buf.get(pos).map(|&x| vm.new_pyobj(x)).ok_or(None),
))
})
let byte = zelf
.internal
.next(|bytearray, pos| bytearray.borrow_buf().get(pos).copied());
Ok(PyIterReturn::from_result(
byte.map(|byte| vm.new_pyobj(byte)).ok_or(None),
))
}
}

Expand Down
37 changes: 17 additions & 20 deletions crates/vm/src/builtins/bytes.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use super::native_iter::{NativeIterKind, NativeSequenceIter};
use super::{
PositionIterInternal, PyDictRef, PyGenericAlias, PyStrRef, PyTuple, PyTupleRef, PyType,
PyTypeRef, iter::builtins_iter, locked_next,
PyDictRef, PyGenericAlias, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef,
iter::builtins_iter,
};
use crate::common::lock::LazyLock;
use crate::{
Expand All @@ -15,7 +16,7 @@ use crate::{
ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
},
class::{PyClassDef, PyClassImpl},
common::{hash::PyHash, lock::PyMutex},
common::hash::PyHash,
convert::{ToPyObject, ToPyResult},
function::{ArgBytesLike, ArgIndex, FuncArgs, OptionalArg, PyComparisonValue},
protocol::{
Expand Down Expand Up @@ -759,7 +760,7 @@ impl Comparable for PyBytes {
impl Iterable for PyBytes {
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
Ok(PyBytesIterator {
internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
internal: NativeSequenceIter::new(zelf, NativeIterKind::Sequence, 0),
}
.into_pyobject(vm))
}
Expand All @@ -772,10 +773,10 @@ impl Representable for PyBytes {
}
}

#[pyclass(module = false, name = "bytes_iterator")]
#[pyclass(module = false, name = "bytes_iterator", traverse)]
#[derive(Debug)]
pub(crate) struct PyBytesIterator {
internal: PyMutex<PositionIterInternal<PyBytesRef>>,
internal: NativeSequenceIter<PyBytesRef>,
}

impl PyPayload for PyBytesIterator {
Expand All @@ -789,15 +790,15 @@ impl PyPayload for PyBytesIterator {
impl Py<PyBytesIterator> {
#[pymethod]
fn __length_hint__(&self) -> usize {
self.internal.lock().length_hint(|obj| obj.as_bytes().len())
self.internal.length_hint(|obj| obj.as_bytes().len())
}

#[pymethod]
fn __reduce__(&self, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
let func = builtins_iter(vm)?;
Ok(self.internal.lock().reduce(
Ok(self.internal.reduce(
func,
|x| x.clone().into(),
|x| x.into(),
|vm| vm.ctx.empty_tuple.clone().into(),
vm,
))
Expand All @@ -806,23 +807,19 @@ impl Py<PyBytesIterator> {
#[pymethod]
fn __setstate__(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
self.internal
.lock()
.set_state(&object, |obj, pos| pos.min(obj.as_bytes().len()), vm)
.set_state(&object, |obj| obj.as_bytes().len(), vm)
}
}

impl SelfIter for PyBytesIterator {}
impl IterNext for PyBytesIterator {
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
locked_next(&zelf.internal, |bytes, pos| {
Ok(PyIterReturn::from_result(
bytes
.as_bytes()
.get(pos)
.map(|&x| vm.new_pyobj(x))
.ok_or(None),
))
})
let byte = zelf
.internal
.next(|bytes, pos| bytes.as_bytes().get(pos).copied());
Ok(PyIterReturn::from_result(
byte.map(|byte| vm.new_pyobj(byte)).ok_or(None),
))
}
}

Expand Down
61 changes: 24 additions & 37 deletions crates/vm/src/builtins/list.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
use super::native_iter::{NativeIterKind, NativeSequenceIter};
use super::{
PositionIterInternal, PyGenericAlias, PyTupleRef, PyType, PyTypeRef,
PyGenericAlias, PyTupleRef, PyType, PyTypeRef,
iter::{builtins_iter, builtins_reversed},
locked_next, locked_rev_next,
};
use crate::atomic_func;
use crate::common::lock::{
PyMappedRwLockReadGuard, PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard,
PyMappedRwLockReadGuard, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard,
};
use crate::object::{Traverse, TraverseFn};
use crate::{
Expand Down Expand Up @@ -521,9 +521,9 @@ impl Py<PyList> {

#[pymethod]
fn __reversed__(zelf: PyRef<PyList>) -> PyListReverseIterator {
let position = zelf.__len__().saturating_sub(1);
let position = zelf.__len__() as isize - 1;
PyListReverseIterator {
internal: PyMutex::new(PositionIterInternal::new(zelf, position)),
internal: NativeSequenceIter::new(zelf, NativeIterKind::ListReverse, position),
}
}

Expand Down Expand Up @@ -767,7 +767,7 @@ impl AsSequence for PyList {
impl Iterable for PyList {
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
Ok(PyListIterator {
internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
internal: NativeSequenceIter::new(zelf, NativeIterKind::List, 0),
}
.into_pyobject(vm))
}
Expand Down Expand Up @@ -1075,7 +1075,7 @@ fn do_sort(
#[pyclass(module = false, name = "list_iterator", traverse)]
#[derive(Debug)]
pub(crate) struct PyListIterator {
internal: PyMutex<PositionIterInternal<PyListRef>>,
internal: NativeSequenceIter<PyListRef>,
}

impl PyPayload for PyListIterator {
Expand All @@ -1089,22 +1089,20 @@ impl PyPayload for PyListIterator {
impl Py<PyListIterator> {
#[pymethod]
fn __length_hint__(&self) -> usize {
self.internal.lock().length_hint(|obj| obj.__len__())
self.internal.length_hint(|obj| obj.__len__())
}

#[pymethod]
fn __setstate__(&self, object: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
self.internal
.lock()
.set_state(&object, |obj, pos| pos.min(obj.__len__()), vm)
self.internal.set_state(&object, |obj| obj.__len__(), vm)
}

#[pymethod]
fn __reduce__(&self, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
let func = builtins_iter(vm)?;
Ok(self.internal.lock().reduce(
Ok(self.internal.reduce(
func,
|x| x.clone().into(),
|x| x.into(),
|vm| vm.ctx.new_list(Vec::new()).into(),
vm,
))
Expand All @@ -1114,32 +1112,22 @@ impl Py<PyListIterator> {
impl PyListIterator {
/// Fast path for FOR_ITER specialization.
pub(crate) fn fast_next(&self) -> Option<PyObjectRef> {
locked_next(&self.internal, |list, pos| {
let vec = list.borrow_vec();
Ok(PyIterReturn::from_result(vec.get(pos).cloned().ok_or(None)))
})
.ok()
.and_then(|r| match r {
PyIterReturn::Return(v) => Some(v),
PyIterReturn::StopIteration(_) => None,
})
self.internal
.next(|list, pos| list.borrow_vec().get(pos).cloned())
}
}

impl SelfIter for PyListIterator {}
impl IterNext for PyListIterator {
fn next(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
locked_next(&zelf.internal, |list, pos| {
let vec = list.borrow_vec();
Ok(PyIterReturn::from_result(vec.get(pos).cloned().ok_or(None)))
})
Ok(PyIterReturn::from_result(zelf.fast_next().ok_or(None)))
}
}

#[pyclass(module = false, name = "list_reverseiterator", traverse)]
#[derive(Debug)]
pub(crate) struct PyListReverseIterator {
internal: PyMutex<PositionIterInternal<PyListRef>>,
internal: NativeSequenceIter<PyListRef>,
}

impl PyPayload for PyListReverseIterator {
Expand All @@ -1153,22 +1141,20 @@ impl PyPayload for PyListReverseIterator {
impl Py<PyListReverseIterator> {
#[pymethod]
fn __length_hint__(&self) -> usize {
self.internal.lock().rev_length_hint(|obj| obj.__len__())
self.internal.length_hint(|obj| obj.__len__())
}

#[pymethod]
fn __setstate__(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
self.internal
.lock()
.set_state(&state, |obj, pos| pos.min(obj.__len__()), vm)
self.internal.set_state(&state, |obj| obj.__len__(), vm)
}

#[pymethod]
fn __reduce__(&self, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
let func = builtins_reversed(vm)?;
Ok(self.internal.lock().reduce(
Ok(self.internal.reduce(
func,
|x| x.clone().into(),
|x| x.into(),
|vm| vm.ctx.new_list(Vec::new()).into(),
vm,
))
Expand All @@ -1178,10 +1164,11 @@ impl Py<PyListReverseIterator> {
impl SelfIter for PyListReverseIterator {}
impl IterNext for PyListReverseIterator {
fn next(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
locked_rev_next(&zelf.internal, |list, pos| {
let vec = list.borrow_vec();
Ok(PyIterReturn::from_result(vec.get(pos).cloned().ok_or(None)))
})
Ok(PyIterReturn::from_result(
zelf.internal
.next(|list, pos| list.borrow_vec().get(pos).cloned())
.ok_or(None),
))
}
}

Expand Down
1 change: 1 addition & 0 deletions crates/vm/src/builtins/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ pub use memory::PyMemoryView;
pub(crate) mod module;
pub use module::{PyModule, PyModuleDef, PyModuleSlots};
pub(crate) mod namespace;
pub(crate) mod native_iter;
pub use namespace::PyNamespace;
pub(crate) mod object;
pub use object::PyBaseObject;
Expand Down
Loading
Loading