rustpython_vm/
sequence.rs

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

rustpython_vm/
sequence.rs
