rustpython_vm/
sliceable.rs

1// export through sliceable module, not slice.
2use crate::{
3    Py, 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#[expect(clippy::len_without_is_empty, reason = "Doesn't match CPython code")]
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    /// The index or slice `obj` stands for, or `None` when it is neither.
262    fn try_from_object_opt(vm: &VirtualMachine, obj: &PyObject) -> Option<PyResult<Self>> {
263        const OVERFLOW: &str = "cannot fit 'int' into an index-sized integer";
264        if let Some(i) = obj.downcast_ref::<PyInt>() {
265            // TODO: number protocol
266            Some(
267                i.try_to_primitive(vm)
268                    .map_err(|_| vm.new_index_error(OVERFLOW))
269                    .map(Self::Int),
270            )
271        } else if let Some(slice) = obj.downcast_ref::<PySlice>() {
272            Some(slice.to_saturated(vm).map(Self::Slice))
273        } else {
274            // TODO: __index__ for indices is no more supported?
275            obj.try_index_opt(vm).map(|i| {
276                i?.try_to_primitive(vm)
277                    .map_err(|_| vm.new_index_error(OVERFLOW))
278                    .map(Self::Int)
279            })
280        }
281    }
282
283    pub fn try_from_borrowed_object(
284        vm: &VirtualMachine,
285        obj: &PyObject,
286        type_name: &str,
287    ) -> PyResult<Self> {
288        Self::try_from_object_opt(vm, obj).unwrap_or_else(|| {
289            Err(vm.new_type_error(format!(
290                "{type_name} indices must be integers or slices, not {}",
291                obj.class().slot_name()
292            )))
293        })
294    }
295
296    /// `unicode_subscript`, which turns down what it cannot use in its own words.
297    pub fn try_from_str_subscript(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Self> {
298        Self::try_from_object_opt(vm, obj).unwrap_or_else(|| {
299            Err(vm.new_type_error(format!(
300                "string indices must be integers, not '{}'",
301                obj.class().slot_name()
302            )))
303        })
304    }
305}
306
307pub trait SequenceIndexOp {
308    // Saturate p in range [0, len] inclusive
309    fn saturated_at(&self, len: usize) -> usize;
310    // Use PySliceableSequence::wrap_index for implementors
311    fn wrapped_at(&self, len: usize) -> Option<usize>;
312}
313
314impl SequenceIndexOp for isize {
315    fn saturated_at(&self, len: usize) -> usize {
316        let len = len.to_isize().unwrap_or(Self::MAX);
317        let mut p = *self;
318        if p < 0 {
319            p += len;
320        }
321        p.clamp(0, len) as usize
322    }
323
324    fn wrapped_at(&self, len: usize) -> Option<usize> {
325        let mut p = *self;
326        if p < 0 {
327            // casting to isize is ok because it is used by wrapping_add
328            p = p.wrapping_add(len as Self);
329        }
330        if p < 0 || (p as usize) >= len {
331            None
332        } else {
333            Some(p as usize)
334        }
335    }
336}
337
338impl SequenceIndexOp for BigInt {
339    fn saturated_at(&self, len: usize) -> usize {
340        if self.is_negative() {
341            self.abs()
342                .try_into()
343                .map_or(0, |abs| len.saturating_sub(abs))
344        } else {
345            self.try_into().unwrap_or(len)
346        }
347    }
348
349    fn wrapped_at(&self, _len: usize) -> Option<usize> {
350        unimplemented!("please add one once we need it")
351    }
352}
353
354/// A saturated slice with values ranging in [isize::MIN, isize::MAX]. Used for
355/// sliceable sequences that require indices in the aforementioned range.
356///
357/// Invokes `__index__` on the PySliceRef during construction so as to separate the
358/// transformation from PyObject into isize and the adjusting of the slice to a given
359/// sequence length. The reason this is important is due to the fact that an objects
360/// `__index__` might get a lock on the sequence and cause a deadlock.
361#[derive(Copy, Clone, Debug)]
362pub struct SaturatedSlice {
363    start: isize,
364    stop: isize,
365    step: isize,
366}
367
368impl SaturatedSlice {
369    #[must_use]
370    pub const fn from_parts(start: isize, stop: isize, step: isize) -> Self {
371        Self { start, stop, step }
372    }
373
374    #[must_use]
375    pub const fn start(&self) -> isize {
376        self.start
377    }
378
379    #[must_use]
380    pub const fn stop(&self) -> isize {
381        self.stop
382    }
383
384    #[must_use]
385    pub const fn step(&self) -> isize {
386        self.step
387    }
388
389    // Equivalent to PySlice_Unpack.
390    pub fn with_slice(slice: &Py<PySlice>, vm: &VirtualMachine) -> PyResult<Self> {
391        let step = to_isize_index(vm, slice.step_ref(vm))?.unwrap_or(1);
392        if step == 0 {
393            return Err(vm.new_value_error("slice step cannot be zero"));
394        }
395        let start = to_isize_index(vm, slice.start_ref(vm))?
396            .unwrap_or_else(|| if step.is_negative() { isize::MAX } else { 0 });
397
398        let stop = to_isize_index(vm, &slice.stop)?.unwrap_or_else(|| {
399            if step.is_negative() {
400                isize::MIN
401            } else {
402                isize::MAX
403            }
404        });
405        Ok(Self { start, stop, step })
406    }
407
408    // Equivalent to PySlice_AdjustIndices
409    /// Convert for usage in indexing the underlying rust collections. Called *after*
410    /// __index__ has been called on the Slice which might mutate the collection.
411    #[must_use]
412    pub fn adjust_indices(&self, len: usize) -> (Range<usize>, isize, usize) {
413        if len == 0 {
414            return (0..0, self.step, 0);
415        }
416        let range = if self.step.is_negative() {
417            let stop = if self.stop == -1 {
418                len
419            } else {
420                self.stop.saturating_add(1).saturated_at(len)
421            };
422            let start = if self.start == -1 {
423                len
424            } else {
425                self.start.saturating_add(1).saturated_at(len)
426            };
427            stop..start
428        } else {
429            self.start.saturated_at(len)..self.stop.saturated_at(len)
430        };
431
432        let (range, slice_len) = if range.start >= range.end {
433            (range.start..range.start, 0)
434        } else {
435            let slice_len = (range.end - range.start - 1) / self.step.unsigned_abs() + 1;
436            (range, slice_len)
437        };
438        (range, self.step, slice_len)
439    }
440
441    // PySlice_AdjustIndices, keeping the adjusted start rather than a range.
442    /// The index the slice begins at, clamped into `0..=len` for a positive step
443    /// and into `-1..=len-1` for a negative one, together with its length.
444    ///
445    /// Unlike [`Self::adjust_indices`] this stays meaningful for an empty slice,
446    /// where it is still the position a strided view moves to.
447    #[must_use]
448    pub fn adjust_indices_start(&self, len: usize) -> (isize, usize) {
449        let len = len as isize;
450        let clamp = |i: isize| {
451            if i < 0 {
452                let i = i.saturating_add(len);
453                if i < 0 {
454                    if self.step.is_negative() { -1 } else { 0 }
455                } else {
456                    i
457                }
458            } else if i >= len {
459                if self.step.is_negative() {
460                    len - 1
461                } else {
462                    len
463                }
464            } else {
465                i
466            }
467        };
468        let start = clamp(self.start);
469        let stop = clamp(self.stop);
470        let step = self.step.unsigned_abs();
471        let slice_len = if self.step.is_negative() {
472            if stop < start {
473                (start - stop - 1) as usize / step + 1
474            } else {
475                0
476            }
477        } else if start < stop {
478            (stop - start - 1) as usize / step + 1
479        } else {
480            0
481        };
482        (start, slice_len)
483    }
484
485    #[must_use]
486    pub fn iter(&self, len: usize) -> SaturatedSliceIter {
487        SaturatedSliceIter::new(self, len)
488    }
489}
490
491pub struct SaturatedSliceIter {
492    index: isize,
493    step: isize,
494    len: usize,
495}
496
497impl SaturatedSliceIter {
498    #[must_use]
499    pub fn new(slice: &SaturatedSlice, seq_len: usize) -> Self {
500        let (range, step, len) = slice.adjust_indices(seq_len);
501        Self::from_adjust_indices(range, step, len)
502    }
503
504    #[must_use]
505    pub const fn from_adjust_indices(range: Range<usize>, step: isize, len: usize) -> Self {
506        let index = if step.is_negative() {
507            range.end as isize - 1
508        } else {
509            range.start as isize
510        };
511        Self { index, step, len }
512    }
513
514    #[must_use]
515    pub const fn positive_order(mut self) -> Self {
516        if self.step.is_negative() {
517            self.index += self.step * self.len.saturating_sub(1) as isize;
518            self.step = self.step.saturating_abs()
519        }
520        self
521    }
522}
523
524impl Iterator for SaturatedSliceIter {
525    type Item = usize;
526
527    fn next(&mut self) -> Option<Self::Item> {
528        if self.len == 0 {
529            return None;
530        }
531        self.len -= 1;
532        let ret = self.index as usize;
533        // SAFETY: if index is overflowed, len should be zero
534        self.index = self.index.wrapping_add(self.step);
535        Some(ret)
536    }
537}
538
539// Go from PyObjectRef to isize w/o overflow error, out of range values are substituted by
540// isize::MIN or isize::MAX depending on type and value of step.
541// Equivalent to PyEval_SliceIndex.
542fn to_isize_index(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Option<isize>> {
543    if vm.is_none(obj) {
544        return Ok(None);
545    }
546    let result = obj.try_index_opt(vm).unwrap_or_else(|| {
547        Err(vm.new_type_error("slice indices must be integers or None or have an __index__ method"))
548    })?;
549    let value = result.as_bigint();
550    let is_negative = value.is_negative();
551    Ok(Some(value.to_isize().unwrap_or(if is_negative {
552        isize::MIN
553    } else {
554        isize::MAX
555    })))
556}
rustpython_vm/sliceable.rs

rustpython_vm/
sliceable.rs
