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

rustpython_vm/builtins/
slice.rs
