rustpython_vm/builtins/
slice.rs

1// sliceobject.{h,c} in CPython
2// spell-checker:ignore sliceobject
3use core::cell::Cell;
4use core::ptr::NonNull;
5use rustpython_common::wtf8::{Wtf8Buf, wtf8_concat};
6
7use super::{PyGenericAlias, PyStrRef, PyTupleRef, PyType, PyTypeRef};
8use crate::{
9    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
10    class::{PyClassDef, PyClassImpl},
11    common::hash::{PyHash, PyUHash},
12    function::{ArgIndex, FuncArgs, OptionalArg, PyComparisonValue},
13    sliceable::SaturatedSlice,
14    types::{Comparable, Constructor, Hashable, PyComparisonOp, Representable},
15};
16use malachite_bigint::{BigInt, ToBigInt};
17use num_traits::{One, Signed, Zero};
18
19#[pyclass(module = false, name = "slice", unhashable = true, traverse = "manual")]
20#[derive(Debug)]
21pub struct PySlice {
22    #[pymember]
23    pub start: Option<PyObjectRef>,
24    #[pymember]
25    pub stop: PyObjectRef,
26    #[pymember]
27    pub step: Option<PyObjectRef>,
28}
29
30// SAFETY: Traverse properly visits all owned PyObjectRefs
31unsafe impl crate::object::Traverse for PySlice {
32    fn traverse(&self, traverse_fn: &mut crate::object::TraverseFn<'_>) {
33        self.start.traverse(traverse_fn);
34        self.stop.traverse(traverse_fn);
35        self.step.traverse(traverse_fn);
36    }
37
38    fn clear(&mut self, out: &mut Vec<PyObjectRef>) {
39        if let Some(start) = self.start.take() {
40            out.push(start);
41        }
42        // stop is not Option, so it will be freed when payload is dropped
43        // (via drop_in_place on freelist pop, or Box::from_raw on dealloc)
44        if let Some(step) = self.step.take() {
45            out.push(step);
46        }
47    }
48}
49
50thread_local! {
51    static SLICE_FREELIST: Cell<crate::object::FreeList<PySlice>> = const { Cell::new(crate::object::FreeList::new()) };
52}
53
54impl PyPayload for PySlice {
55    const MAX_FREELIST: usize = 1;
56    const HAS_FREELIST: bool = true;
57
58    #[inline]
59    fn class(ctx: &Context) -> &'static Py<PyType> {
60        ctx.types.slice_type
61    }
62
63    #[inline]
64    unsafe fn freelist_push(obj: *mut PyObject) -> bool {
65        SLICE_FREELIST
66            .try_with(|fl| {
67                let mut list = fl.take();
68                let stored = if list.len() < Self::MAX_FREELIST {
69                    list.push(obj);
70                    true
71                } else {
72                    false
73                };
74                fl.set(list);
75                stored
76            })
77            .unwrap_or(false)
78    }
79
80    #[inline]
81    unsafe fn freelist_pop(_payload: &Self) -> Option<NonNull<PyObject>> {
82        SLICE_FREELIST
83            .try_with(|fl| {
84                let mut list = fl.take();
85                let result = list.pop().map(|p| unsafe { NonNull::new_unchecked(p) });
86                fl.set(list);
87                result
88            })
89            .ok()
90            .flatten()
91    }
92}
93
94impl PySlice {
95    pub(crate) fn start_ref<'a>(&'a self, vm: &'a VirtualMachine) -> &'a PyObject {
96        match &self.start {
97            Some(v) => v,
98            None => vm.ctx.none.as_object(),
99        }
100    }
101
102    pub(crate) fn step_ref<'a>(&'a self, vm: &'a VirtualMachine) -> &'a PyObject {
103        match &self.step {
104            Some(v) => v,
105            None => vm.ctx.none.as_object(),
106        }
107    }
108
109    pub(crate) fn inner_indices(
110        &self,
111        length: &BigInt,
112        vm: &VirtualMachine,
113    ) -> PyResult<(BigInt, BigInt, BigInt)> {
114        // Calculate step
115        let step: BigInt;
116        if vm.is_none(self.step_ref(vm)) {
117            step = One::one();
118        } else {
119            // Clone the value, not the reference.
120            let this_step = self.step_ref(vm).try_index(vm)?;
121            step = this_step.as_bigint().clone();
122
123            if step.is_zero() {
124                return Err(vm.new_value_error("slice step cannot be zero."));
125            }
126        }
127
128        // For convenience
129        let backwards = step.is_negative();
130
131        // Each end of the array
132        let lower = if backwards {
133            (-1_i8).to_bigint().unwrap()
134        } else {
135            Zero::zero()
136        };
137
138        let upper = if backwards {
139            lower.clone() + length
140        } else {
141            length.clone()
142        };
143
144        // Calculate start
145        let mut start: BigInt;
146        if vm.is_none(self.start_ref(vm)) {
147            // Default
148            start = if backwards {
149                upper.clone()
150            } else {
151                lower.clone()
152            };
153        } else {
154            let this_start = self.start_ref(vm).try_index(vm)?;
155            start = this_start.as_bigint().clone();
156
157            if start < Zero::zero() {
158                // From end of array
159                start += length;
160
161                if start < lower {
162                    start = lower.clone();
163                }
164            } else if start > upper {
165                start = upper.clone();
166            }
167        }
168
169        // Calculate Stop
170        let mut stop: BigInt;
171        if vm.is_none(&self.stop) {
172            stop = if backwards { lower } else { upper };
173        } else {
174            let this_stop = self.stop.try_index(vm)?;
175            stop = this_stop.as_bigint().clone();
176
177            if stop < Zero::zero() {
178                // From end of array
179                stop += length;
180                if stop < lower {
181                    stop = lower;
182                }
183            } else if stop > upper {
184                stop = upper;
185            }
186        }
187
188        Ok((start, stop, step))
189    }
190
191    // TODO: Uncomment when Python adds __class_getitem__ to slice
192    // #[pyclassmethod]
193    fn __class_getitem__(
194        cls: PyTypeRef,
195        args: PyObjectRef,
196        vm: &VirtualMachine,
197    ) -> PyResult<PyGenericAlias> {
198        PyGenericAlias::from_args(cls, args, vm)
199    }
200}
201
202#[pyclass(with(Comparable, Representable, Hashable))]
203impl Py<PySlice> {
204    #[pyslot]
205    fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
206        let slice: PySlice = match args.args.len() {
207            0 => {
208                return Err(vm.new_arity_type_error(PySlice::NAME, 1..=3, 0));
209            }
210            1 => {
211                let stop = args.bind_for(vm, PySlice::NAME)?;
212                PySlice {
213                    start: None,
214                    stop,
215                    step: None,
216                }
217            }
218            _ => {
219                let (start, stop, step): (PyObjectRef, PyObjectRef, OptionalArg<PyObjectRef>) =
220                    args.bind_for(vm, PySlice::NAME)?;
221                PySlice {
222                    start: Some(start),
223                    stop,
224                    step: step.into_option(),
225                }
226            }
227        };
228        slice.into_ref_with_type(vm, cls).map(Into::into)
229    }
230
231    #[pymethod]
232    fn indices(&self, object: ArgIndex, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
233        let length = object.into_int_ref();
234        let length = length.as_bigint();
235        if length.is_negative() {
236            return Err(vm.new_value_error("length should not be negative."));
237        }
238        let (start, stop, step) = self.inner_indices(length, vm)?;
239        Ok(vm.new_tuple((start, stop, step)))
240    }
241
242    #[allow(clippy::type_complexity)]
243    #[pymethod]
244    fn __reduce__(
245        zelf: PyRef<PySlice>,
246    ) -> (
247        PyTypeRef,
248        (Option<PyObjectRef>, PyObjectRef, Option<PyObjectRef>),
249    ) {
250        (
251            zelf.class().to_owned(),
252            (zelf.start.clone(), zelf.stop.clone(), zelf.step.clone()),
253        )
254    }
255}
256
257impl Py<PySlice> {
258    pub fn to_saturated(&self, vm: &VirtualMachine) -> PyResult<SaturatedSlice> {
259        SaturatedSlice::with_slice(self, vm)
260    }
261}
262
263impl Hashable for PySlice {
264    #[inline]
265    fn hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyHash> {
266        const XXPRIME_1: PyUHash = if cfg!(target_pointer_width = "64") {
267            11400714785074694791
268        } else {
269            2654435761
270        };
271        const XXPRIME_2: PyUHash = if cfg!(target_pointer_width = "64") {
272            14029467366897019727
273        } else {
274            2246822519
275        };
276        const XXPRIME_5: PyUHash = if cfg!(target_pointer_width = "64") {
277            2870177450012600261
278        } else {
279            374761393
280        };
281        const ROTATE: u32 = if cfg!(target_pointer_width = "64") {
282            31
283        } else {
284            13
285        };
286
287        let mut acc = XXPRIME_5;
288        for part in &[zelf.start_ref(vm), &zelf.stop, zelf.step_ref(vm)] {
289            let lane = part.hash(vm)? as PyUHash;
290            if lane == u64::MAX as PyUHash {
291                return Ok(-1 as PyHash);
292            }
293            acc = acc.wrapping_add(lane.wrapping_mul(XXPRIME_2));
294            acc = acc.rotate_left(ROTATE);
295            acc = acc.wrapping_mul(XXPRIME_1);
296        }
297        if acc == u64::MAX as PyUHash {
298            return Ok(1546275796 as PyHash);
299        }
300        Ok(acc as PyHash)
301    }
302}
303
304impl Comparable for PySlice {
305    fn cmp(
306        zelf: &Py<Self>,
307        other: &PyObject,
308        op: PyComparisonOp,
309        vm: &VirtualMachine,
310    ) -> PyResult<PyComparisonValue> {
311        let other = class_or_notimplemented!(Self, other);
312
313        let ret = match op {
314            PyComparisonOp::Lt | PyComparisonOp::Le => None
315                .or_else(|| {
316                    vm.bool_seq_lt(zelf.start_ref(vm), other.start_ref(vm))
317                        .transpose()
318                })
319                .or_else(|| vm.bool_seq_lt(&zelf.stop, &other.stop).transpose())
320                .or_else(|| {
321                    vm.bool_seq_lt(zelf.step_ref(vm), other.step_ref(vm))
322                        .transpose()
323                })
324                .unwrap_or_else(|| Ok(op == PyComparisonOp::Le))?,
325            PyComparisonOp::Eq | PyComparisonOp::Ne => {
326                let eq = vm.identical_or_equal(zelf.start_ref(vm), other.start_ref(vm))?
327                    && vm.identical_or_equal(&zelf.stop, &other.stop)?
328                    && vm.identical_or_equal(zelf.step_ref(vm), other.step_ref(vm))?;
329                if op == PyComparisonOp::Ne { !eq } else { eq }
330            }
331            PyComparisonOp::Gt | PyComparisonOp::Ge => None
332                .or_else(|| {
333                    vm.bool_seq_gt(zelf.start_ref(vm), other.start_ref(vm))
334                        .transpose()
335                })
336                .or_else(|| vm.bool_seq_gt(&zelf.stop, &other.stop).transpose())
337                .or_else(|| {
338                    vm.bool_seq_gt(zelf.step_ref(vm), other.step_ref(vm))
339                        .transpose()
340                })
341                .unwrap_or_else(|| Ok(op == PyComparisonOp::Ge))?,
342        };
343
344        Ok(PyComparisonValue::Implemented(ret))
345    }
346}
347
348impl Representable for PySlice {
349    #[inline]
350    fn repr_wtf8(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
351        let start_repr = zelf.start_ref(vm).repr(vm)?;
352        let stop_repr = zelf.stop.repr(vm)?;
353        let step_repr = zelf.step_ref(vm).repr(vm)?;
354
355        Ok(wtf8_concat!(
356            "slice(",
357            start_repr.as_wtf8(),
358            ", ",
359            stop_repr.as_wtf8(),
360            ", ",
361            step_repr.as_wtf8(),
362            ")"
363        ))
364    }
365}
366
367#[pyclass(module = false, name = "ellipsis")]
368#[derive(Debug)]
369pub struct PyEllipsis;
370
371impl PyPayload for PyEllipsis {
372    #[inline]
373    fn class(ctx: &Context) -> &'static Py<PyType> {
374        ctx.types.ellipsis_type
375    }
376}
377
378impl Constructor for PyEllipsis {
379    type Args = ();
380
381    fn slot_new(_cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
382        let _: () = args.bind_for(vm, Self::NAME)?;
383        Ok(vm.ctx.ellipsis.clone().into())
384    }
385
386    fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
387        unreachable!("Ellipsis is a singleton")
388    }
389}
390
391#[pyclass(with(Constructor, Representable), flags(IMMUTABLETYPE))]
392impl Py<PyEllipsis> {
393    #[pymethod]
394    fn __reduce__(&self, vm: &VirtualMachine) -> PyStrRef {
395        vm.ctx.names.Ellipsis.to_owned()
396    }
397}
398
399impl Representable for PyEllipsis {
400    #[inline]
401    fn repr(_zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyStrRef> {
402        Ok(vm.ctx.names.Ellipsis.to_owned())
403    }
404
405    #[cold]
406    fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
407        unreachable!("use repr instead")
408    }
409}
410
411pub(crate) fn init(ctx: &'static Context) {
412    PySlice::extend_class(ctx, ctx.types.slice_type);
413    PyEllipsis::extend_class(ctx, ctx.types.ellipsis_type);
414}
rustpython_vm/builtins/slice.rs

rustpython_vm/builtins/
slice.rs
