rustpython_vm/
sequence.rs

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

rustpython_vm/
sequence.rs
