rustpython_vm/
sliceable.rs

1// export through sliceable module, not slice.
2use crate::{
3    PyObject, PyResult, VirtualMachine,
4    builtins::{int::PyInt, slice::PySlice},
5};
6use core::ops::Range;
7use malachite_bigint::BigInt;
8use num_traits::{Signed, ToPrimitive};
9
10pub trait SliceableSequenceMutOp
11where
12    Self: AsRef<[Self::Item]>,
13{
14    type Item: Clone;
15    fn do_set(&mut self, index: usize, value: Self::Item);
16    fn do_delete(&mut self, index: usize);
17    /// as CPython, length of range and items could be different, function must act like Vec::splice()
18    fn do_set_range(&mut self, range: Range<usize>, items: &[Self::Item]);
19    fn do_delete_range(&mut self, range: Range<usize>);
20    fn do_set_indexes<I>(&mut self, indexes: I, items: &[Self::Item])
21    where
22        I: Iterator<Item = usize>;
23    /// indexes must be positive order
24    fn do_delete_indexes<I>(&mut self, range: Range<usize>, indexes: I)
25    where
26        I: Iterator<Item = usize>;
27
28    fn setitem_by_index(
29        &mut self,
30        vm: &VirtualMachine,
31        index: isize,
32        value: Self::Item,
33    ) -> PyResult<()> {
34        let pos = self
35            .as_ref()
36            .wrap_index(index)
37            .ok_or_else(|| vm.new_index_error("assignment index out of range"))?;
38        self.do_set(pos, value);
39        Ok(())
40    }
41
42    fn setitem_by_slice_no_resize(
43        &mut self,
44        vm: &VirtualMachine,
45        slice: SaturatedSlice,
46        items: &[Self::Item],
47    ) -> PyResult<()> {
48        let (range, step, slice_len) = slice.adjust_indices(self.as_ref().len());
49        if slice_len != items.len() {
50            Err(vm.new_buffer_error("Existing exports of data: object cannot be re-sized"))
51        } else if step == 1 {
52            self.do_set_range(range, items);
53            Ok(())
54        } else {
55            self.do_set_indexes(
56                SaturatedSliceIter::from_adjust_indices(range, step, slice_len),
57                items,
58            );
59            Ok(())
60        }
61    }
62
63    fn setitem_by_slice(
64        &mut self,
65        vm: &VirtualMachine,
66        slice: SaturatedSlice,
67        items: &[Self::Item],
68    ) -> PyResult<()> {
69        let (range, step, slice_len) = slice.adjust_indices(self.as_ref().len());
70        if step == 1 {
71            self.do_set_range(range, items);
72            Ok(())
73        } else if slice_len == items.len() {
74            self.do_set_indexes(
75                SaturatedSliceIter::from_adjust_indices(range, step, slice_len),
76                items,
77            );
78            Ok(())
79        } else {
80            Err(vm.new_value_error(format!(
81                "attempt to assign sequence of size {} to extended slice of size {}",
82                items.len(),
83                slice_len
84            )))
85        }
86    }
87
88    fn delitem_by_index(&mut self, vm: &VirtualMachine, index: isize) -> PyResult<()> {
89        let pos = self
90            .as_ref()
91            .wrap_index(index)
92            .ok_or_else(|| vm.new_index_error("assignment index out of range"))?;
93        self.do_delete(pos);
94        Ok(())
95    }
96
97    fn delitem_by_slice(&mut self, _vm: &VirtualMachine, slice: SaturatedSlice) -> PyResult<()> {
98        let (range, step, slice_len) = slice.adjust_indices(self.as_ref().len());
99        if slice_len == 0 {
100            Ok(())
101        } else if step == 1 || step == -1 {
102            self.do_set_range(range, &[]);
103            Ok(())
104        } else {
105            self.do_delete_indexes(
106                range.clone(),
107                SaturatedSliceIter::from_adjust_indices(range, step, slice_len).positive_order(),
108            );
109            Ok(())
110        }
111    }
112}
113
114impl<T: Clone> SliceableSequenceMutOp for Vec<T> {
115    type Item = T;
116
117    fn do_set(&mut self, index: usize, value: Self::Item) {
118        self[index] = value;
119    }
120
121    fn do_delete(&mut self, index: usize) {
122        self.remove(index);
123    }
124
125    fn do_set_range(&mut self, range: Range<usize>, items: &[Self::Item]) {
126        self.splice(range, items.to_vec());
127    }
128
129    fn do_delete_range(&mut self, range: Range<usize>) {
130        self.drain(range);
131    }
132
133    fn do_set_indexes<I>(&mut self, indexes: I, items: &[Self::Item])
134    where
135        I: Iterator<Item = usize>,
136    {
137        for (i, item) in indexes.zip(items) {
138            self.do_set(i, item.clone());
139        }
140    }
141
142    fn do_delete_indexes<I>(&mut self, range: Range<usize>, indexes: I)
143    where
144        I: Iterator<Item = usize>,
145    {
146        let mut indexes = indexes.peekable();
147        let mut deleted = 0;
148
149        // passing whole range, swap or overlap
150        for i in range.clone() {
151            if indexes.peek() == Some(&i) {
152                indexes.next();
153                deleted += 1;
154            } else {
155                self.swap(i - deleted, i);
156            }
157        }
158        // then drain (the values to delete should now be contiguous at the end of the range)
159        self.drain((range.end - deleted)..range.end);
160    }
161}
162
163#[allow(clippy::len_without_is_empty)]
164pub trait SliceableSequenceOp {
165    type Item;
166    type Sliced;
167
168    fn do_get(&self, index: usize) -> Self::Item;
169    fn do_slice(&self, range: Range<usize>) -> Self::Sliced;
170    fn do_slice_reverse(&self, range: Range<usize>) -> Self::Sliced;
171    fn do_stepped_slice(&self, range: Range<usize>, step: usize) -> Self::Sliced;
172    fn do_stepped_slice_reverse(&self, range: Range<usize>, step: usize) -> Self::Sliced;
173    fn empty() -> Self::Sliced;
174
175    fn len(&self) -> usize;
176
177    fn wrap_index(&self, p: isize) -> Option<usize> {
178        p.wrapped_at(self.len())
179    }
180
181    fn saturate_index(&self, p: isize) -> usize {
182        p.saturated_at(self.len())
183    }
184
185    fn getitem_by_slice(
186        &self,
187        _vm: &VirtualMachine,
188        slice: SaturatedSlice,
189    ) -> PyResult<Self::Sliced> {
190        let (range, step, slice_len) = slice.adjust_indices(self.len());
191        let sliced = if slice_len == 0 {
192            Self::empty()
193        } else if step == 1 {
194            self.do_slice(range)
195        } else if step == -1 {
196            self.do_slice_reverse(range)
197        } else if step.is_positive() {
198            self.do_stepped_slice(range, step.unsigned_abs())
199        } else {
200            self.do_stepped_slice_reverse(range, step.unsigned_abs())
201        };
202        Ok(sliced)
203    }
204
205    fn getitem_by_index(&self, vm: &VirtualMachine, index: isize) -> PyResult<Self::Item> {
206        let pos = self
207            .wrap_index(index)
208            .ok_or_else(|| vm.new_index_error("index out of range"))?;
209        Ok(self.do_get(pos))
210    }
211}
212
213impl<T: Clone> SliceableSequenceOp for [T] {
214    type Item = T;
215    type Sliced = Vec<T>;
216
217    #[inline]
218    fn do_get(&self, index: usize) -> Self::Item {
219        self[index].clone()
220    }
221
222    #[inline]
223    fn do_slice(&self, range: Range<usize>) -> Self::Sliced {
224        self[range].to_vec()
225    }
226
227    #[inline]
228    fn do_slice_reverse(&self, range: Range<usize>) -> Self::Sliced {
229        let mut slice = self[range].to_vec();
230        slice.reverse();
231        slice
232    }
233
234    #[inline]
235    fn do_stepped_slice(&self, range: Range<usize>, step: usize) -> Self::Sliced {
236        self[range].iter().step_by(step).cloned().collect()
237    }
238
239    #[inline]
240    fn do_stepped_slice_reverse(&self, range: Range<usize>, step: usize) -> Self::Sliced {
241        self[range].iter().rev().step_by(step).cloned().collect()
242    }
243
244    #[inline(always)]
245    fn empty() -> Self::Sliced {
246        Vec::new()
247    }
248
249    #[inline(always)]
250    fn len(&self) -> usize {
251        self.len()
252    }
253}
254
255pub enum SequenceIndex {
256    Int(isize),
257    Slice(SaturatedSlice),
258}
259
260impl SequenceIndex {
261    pub fn try_from_borrowed_object(
262        vm: &VirtualMachine,
263        obj: &PyObject,
264        type_name: &str,
265    ) -> PyResult<Self> {
266        if let Some(i) = obj.downcast_ref::<PyInt>() {
267            // TODO: number protocol
268            i.try_to_primitive(vm)
269                .map_err(|_| vm.new_index_error("cannot fit 'int' into an index-sized integer"))
270                .map(Self::Int)
271        } else if let Some(slice) = obj.downcast_ref::<PySlice>() {
272            slice.to_saturated(vm).map(Self::Slice)
273        } else if let Some(i) = obj.try_index_opt(vm) {
274            // TODO: __index__ for indices is no more supported?
275            i?.try_to_primitive(vm)
276                .map_err(|_| vm.new_index_error("cannot fit 'int' into an index-sized integer"))
277                .map(Self::Int)
278        } else {
279            Err(vm.new_type_error(format!(
280                "{} indices must be integers or slices or classes that override __index__ operator, not '{}'",
281                type_name,
282                obj.class()
283            )))
284        }
285    }
286}
287
288pub trait SequenceIndexOp {
289    // Saturate p in range [0, len] inclusive
290    fn saturated_at(&self, len: usize) -> usize;
291    // Use PySliceableSequence::wrap_index for implementors
292    fn wrapped_at(&self, len: usize) -> Option<usize>;
293}
294
295impl SequenceIndexOp for isize {
296    fn saturated_at(&self, len: usize) -> usize {
297        let len = len.to_isize().unwrap_or(Self::MAX);
298        let mut p = *self;
299        if p < 0 {
300            p += len;
301        }
302        p.clamp(0, len) as usize
303    }
304
305    fn wrapped_at(&self, len: usize) -> Option<usize> {
306        let mut p = *self;
307        if p < 0 {
308            // casting to isize is ok because it is used by wrapping_add
309            p = p.wrapping_add(len as Self);
310        }
311        if p < 0 || (p as usize) >= len {
312            None
313        } else {
314            Some(p as usize)
315        }
316    }
317}
318
319impl SequenceIndexOp for BigInt {
320    fn saturated_at(&self, len: usize) -> usize {
321        if self.is_negative() {
322            self.abs()
323                .try_into()
324                .map_or(0, |abs| len.saturating_sub(abs))
325        } else {
326            self.try_into().unwrap_or(len)
327        }
328    }
329
330    fn wrapped_at(&self, _len: usize) -> Option<usize> {
331        unimplemented!("please add one once we need it")
332    }
333}
334
335/// A saturated slice with values ranging in [isize::MIN, isize::MAX]. Used for
336/// sliceable sequences that require indices in the aforementioned range.
337///
338/// Invokes `__index__` on the PySliceRef during construction so as to separate the
339/// transformation from PyObject into isize and the adjusting of the slice to a given
340/// sequence length. The reason this is important is due to the fact that an objects
341/// `__index__` might get a lock on the sequence and cause a deadlock.
342#[derive(Copy, Clone, Debug)]
343pub struct SaturatedSlice {
344    start: isize,
345    stop: isize,
346    step: isize,
347}
348
349impl SaturatedSlice {
350    // Equivalent to PySlice_Unpack.
351    pub fn with_slice(slice: &PySlice, vm: &VirtualMachine) -> PyResult<Self> {
352        let step = to_isize_index(vm, slice.step_ref(vm))?.unwrap_or(1);
353        if step == 0 {
354            return Err(vm.new_value_error("slice step cannot be zero"));
355        }
356        let start = to_isize_index(vm, slice.start_ref(vm))?
357            .unwrap_or_else(|| if step.is_negative() { isize::MAX } else { 0 });
358
359        let stop = to_isize_index(vm, &slice.stop(vm))?.unwrap_or_else(|| {
360            if step.is_negative() {
361                isize::MIN
362            } else {
363                isize::MAX
364            }
365        });
366        Ok(Self { start, stop, step })
367    }
368
369    // Equivalent to PySlice_AdjustIndices
370    /// Convert for usage in indexing the underlying rust collections. Called *after*
371    /// __index__ has been called on the Slice which might mutate the collection.
372    pub fn adjust_indices(&self, len: usize) -> (Range<usize>, isize, usize) {
373        if len == 0 {
374            return (0..0, self.step, 0);
375        }
376        let range = if self.step.is_negative() {
377            let stop = if self.stop == -1 {
378                len
379            } else {
380                self.stop.saturating_add(1).saturated_at(len)
381            };
382            let start = if self.start == -1 {
383                len
384            } else {
385                self.start.saturating_add(1).saturated_at(len)
386            };
387            stop..start
388        } else {
389            self.start.saturated_at(len)..self.stop.saturated_at(len)
390        };
391
392        let (range, slice_len) = if range.start >= range.end {
393            (range.start..range.start, 0)
394        } else {
395            let slice_len = (range.end - range.start - 1) / self.step.unsigned_abs() + 1;
396            (range, slice_len)
397        };
398        (range, self.step, slice_len)
399    }
400
401    pub fn iter(&self, len: usize) -> SaturatedSliceIter {
402        SaturatedSliceIter::new(self, len)
403    }
404}
405
406pub struct SaturatedSliceIter {
407    index: isize,
408    step: isize,
409    len: usize,
410}
411
412impl SaturatedSliceIter {
413    pub fn new(slice: &SaturatedSlice, seq_len: usize) -> Self {
414        let (range, step, len) = slice.adjust_indices(seq_len);
415        Self::from_adjust_indices(range, step, len)
416    }
417
418    pub const fn from_adjust_indices(range: Range<usize>, step: isize, len: usize) -> Self {
419        let index = if step.is_negative() {
420            range.end as isize - 1
421        } else {
422            range.start as isize
423        };
424        Self { index, step, len }
425    }
426
427    pub const fn positive_order(mut self) -> Self {
428        if self.step.is_negative() {
429            self.index += self.step * self.len.saturating_sub(1) as isize;
430            self.step = self.step.saturating_abs()
431        }
432        self
433    }
434}
435
436impl Iterator for SaturatedSliceIter {
437    type Item = usize;
438
439    fn next(&mut self) -> Option<Self::Item> {
440        if self.len == 0 {
441            return None;
442        }
443        self.len -= 1;
444        let ret = self.index as usize;
445        // SAFETY: if index is overflowed, len should be zero
446        self.index = self.index.wrapping_add(self.step);
447        Some(ret)
448    }
449}
450
451// Go from PyObjectRef to isize w/o overflow error, out of range values are substituted by
452// isize::MIN or isize::MAX depending on type and value of step.
453// Equivalent to PyEval_SliceIndex.
454fn to_isize_index(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Option<isize>> {
455    if vm.is_none(obj) {
456        return Ok(None);
457    }
458    let result = obj.try_index_opt(vm).unwrap_or_else(|| {
459        Err(vm.new_type_error("slice indices must be integers or None or have an __index__ method"))
460    })?;
461    let value = result.as_bigint();
462    let is_negative = value.is_negative();
463    Ok(Some(value.to_isize().unwrap_or(if is_negative {
464        isize::MIN
465    } else {
466        isize::MAX
467    })))
468}
rustpython_vm/sliceable.rs

rustpython_vm/
sliceable.rs
