rustpython_vm/
sliceable.rs

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

rustpython_vm/
sliceable.rs
