rustpython_vm/
sequence.rs

1use crate::{
2    AsObject, PyObject, PyObjectRef, PyResult,
3    function::OptionalArg,
4    sliceable::SequenceIndexOp,
5    types::PyComparisonOp,
6    vm::{MAX_MEMORY_SIZE, VirtualMachine},
7};
8use core::ops::{Deref, Range};
9use optional::Optioned;
10
11pub trait MutObjectSequenceOp {
12    type Inner: ?Sized;
13
14    fn do_get(index: usize, inner: &Self::Inner) -> Option<&PyObject>;
15    fn do_lock(&self) -> impl Deref<Target = Self::Inner>;
16
17    fn mut_count(&self, vm: &VirtualMachine, needle: &PyObject) -> PyResult<usize> {
18        let mut count = 0;
19        self._mut_iter_equal_skeleton::<_, false>(vm, needle, 0..isize::MAX as usize, || {
20            count += 1
21        })?;
22        Ok(count)
23    }
24
25    fn mut_index_range(
26        &self,
27        vm: &VirtualMachine,
28        needle: &PyObject,
29        range: Range<usize>,
30    ) -> PyResult<Optioned<usize>> {
31        self._mut_iter_equal_skeleton::<_, true>(vm, needle, range, || {})
32    }
33
34    fn mut_index(&self, vm: &VirtualMachine, needle: &PyObject) -> PyResult<Optioned<usize>> {
35        self.mut_index_range(vm, needle, 0..isize::MAX as usize)
36    }
37
38    fn mut_contains(&self, vm: &VirtualMachine, needle: &PyObject) -> PyResult<bool> {
39        self.mut_index(vm, needle).map(|x| x.is_some())
40    }
41
42    fn _mut_iter_equal_skeleton<F, const SHORT: bool>(
43        &self,
44        vm: &VirtualMachine,
45        needle: &PyObject,
46        range: Range<usize>,
47        mut f: F,
48    ) -> PyResult<Optioned<usize>>
49    where
50        F: FnMut(),
51    {
52        let mut borrower = None;
53        let mut i = range.start;
54
55        let index = loop {
56            if i >= range.end {
57                break Optioned::<usize>::none();
58            }
59            let guard = if let Some(x) = borrower.take() {
60                x
61            } else {
62                self.do_lock()
63            };
64
65            let elem = if let Some(x) = Self::do_get(i, &guard) {
66                x
67            } else {
68                break Optioned::<usize>::none();
69            };
70
71            if elem.is(needle) {
72                f();
73                if SHORT {
74                    break Optioned::<usize>::some(i);
75                }
76                borrower = Some(guard);
77            } else {
78                let elem = elem.to_owned();
79                drop(guard);
80
81                if elem.rich_compare_bool(needle, PyComparisonOp::Eq, vm)? {
82                    f();
83                    if SHORT {
84                        break Optioned::<usize>::some(i);
85                    }
86                }
87            }
88            i += 1;
89        };
90
91        Ok(index)
92    }
93}
94
95pub trait SequenceExt<T: Clone>
96where
97    Self: AsRef<[T]>,
98{
99    fn mul(&self, vm: &VirtualMachine, n: isize) -> PyResult<Vec<T>> {
100        let n = vm.check_repeat_or_overflow_error(self.as_ref().len(), n)?;
101        if self.as_ref().is_empty() {
102            return Ok(Vec::new());
103        }
104
105        if n > 1 && core::mem::size_of_val(self.as_ref()) >= MAX_MEMORY_SIZE / n {
106            return Err(vm.no_memory_error());
107        }
108
109        let total = n
110            .checked_mul(self.as_ref().len())
111            .ok_or_else(|| vm.no_memory_error())?;
112        let mut v = Vec::new();
113        v.try_reserve_exact(total)
114            .map_err(|_| vm.no_memory_error())?;
115        for _ in 0..n {
116            v.extend_from_slice(self.as_ref());
117        }
118        Ok(v)
119    }
120}
121
122impl<T: Clone> SequenceExt<T> for [T] {}
123
124pub trait SequenceMutExt<T: Clone>
125where
126    Self: AsRef<[T]>,
127{
128    fn as_vec_mut(&mut self) -> &mut Vec<T>;
129
130    fn imul(&mut self, vm: &VirtualMachine, n: isize) -> PyResult<()> {
131        let n = vm.check_repeat_or_overflow_error(self.as_ref().len(), n)?;
132
133        if n > 1 && core::mem::size_of_val(self.as_ref()) >= MAX_MEMORY_SIZE / n {
134            // TODO: make a global static NoMemory shared exc object and return its reference.
135            return Err(vm.no_memory_error());
136        }
137
138        if n == 0 {
139            self.as_vec_mut().clear();
140        } else if n != 1 && !self.as_ref().is_empty() {
141            let len = self.as_ref().len();
142            let v = self.as_vec_mut();
143            v.try_reserve_exact(len * (n - 1))
144                .map_err(|_| vm.no_memory_error())?;
145            for _ in 1..n {
146                v.extend_from_within(..len);
147            }
148        }
149        Ok(())
150    }
151}
152
153impl<T: Clone> SequenceMutExt<T> for Vec<T> {
154    fn as_vec_mut(&mut self) -> &mut Self {
155        self
156    }
157}
158
159#[derive(FromArgs)]
160pub struct OptionalRangeArgs {
161    #[pyarg(positional, default = 0)]
162    start: PyObjectRef,
163    // Platform ssize maximum. Missing is clamped to the sequence length.
164    #[pyarg(positional, optional, py_default = "9223372036854775807")]
165    stop: OptionalArg<PyObjectRef>,
166}
167
168impl OptionalRangeArgs {
169    pub fn saturate(self, len: usize, vm: &VirtualMachine) -> PyResult<(usize, usize)> {
170        let saturate = |obj: PyObjectRef| -> PyResult<_> {
171            Ok(obj.try_index(vm)?.as_bigint().saturated_at(len))
172        };
173        let start = saturate(self.start)?;
174        let stop = self.stop.map_or(Ok(len), saturate)?;
175        Ok((start, stop))
176    }
177}
rustpython_vm/sequence.rs

rustpython_vm/
sequence.rs
