rustpython_vm/builtins/
slice.rs

1// sliceobject.{h,c} in CPython
2// spell-checker:ignore sliceobject
3use super::{PyStrRef, PyTupleRef, PyType, PyTypeRef};
4use crate::{
5    class::PyClassImpl,
6    common::hash::{PyHash, PyUHash},
7    convert::ToPyObject,
8    function::{ArgIndex, FuncArgs, OptionalArg, PyComparisonValue},
9    sliceable::SaturatedSlice,
10    types::{Comparable, Constructor, Hashable, PyComparisonOp, Representable},
11    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
12};
13use malachite_bigint::{BigInt, ToBigInt};
14use num_traits::{One, Signed, Zero};
15
16#[pyclass(module = false, name = "slice", unhashable = true, traverse)]
17#[derive(Debug)]
18pub struct PySlice {
19    pub start: Option<PyObjectRef>,
20    pub stop: PyObjectRef,
21    pub step: Option<PyObjectRef>,
22}
23
24impl PyPayload for PySlice {
25    fn class(ctx: &Context) -> &'static Py<PyType> {
26        ctx.types.slice_type
27    }
28}
29
30#[pyclass(with(Comparable, Representable, Hashable))]
31impl PySlice {
32    #[pygetset]
33    fn start(&self, vm: &VirtualMachine) -> PyObjectRef {
34        self.start.clone().to_pyobject(vm)
35    }
36
37    pub(crate) fn start_ref<'a>(&'a self, vm: &'a VirtualMachine) -> &'a PyObject {
38        match &self.start {
39            Some(v) => v,
40            None => vm.ctx.none.as_object(),
41        }
42    }
43
44    #[pygetset]
45    pub(crate) fn stop(&self, _vm: &VirtualMachine) -> PyObjectRef {
46        self.stop.clone()
47    }
48
49    #[pygetset]
50    fn step(&self, vm: &VirtualMachine) -> PyObjectRef {
51        self.step.clone().to_pyobject(vm)
52    }
53
54    pub(crate) fn step_ref<'a>(&'a self, vm: &'a VirtualMachine) -> &'a PyObject {
55        match &self.step {
56            Some(v) => v,
57            None => vm.ctx.none.as_object(),
58        }
59    }
60
61    pub fn to_saturated(&self, vm: &VirtualMachine) -> PyResult<SaturatedSlice> {
62        SaturatedSlice::with_slice(self, vm)
63    }
64
65    #[pyslot]
66    fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
67        let slice: PySlice = match args.args.len() {
68            0 => {
69                return Err(
70                    vm.new_type_error("slice() must have at least one arguments.".to_owned())
71                );
72            }
73            1 => {
74                let stop = args.bind(vm)?;
75                PySlice {
76                    start: None,
77                    stop,
78                    step: None,
79                }
80            }
81            _ => {
82                let (start, stop, step): (PyObjectRef, PyObjectRef, OptionalArg<PyObjectRef>) =
83                    args.bind(vm)?;
84                PySlice {
85                    start: Some(start),
86                    stop,
87                    step: step.into_option(),
88                }
89            }
90        };
91        slice.into_ref_with_type(vm, cls).map(Into::into)
92    }
93
94    pub(crate) fn inner_indices(
95        &self,
96        length: &BigInt,
97        vm: &VirtualMachine,
98    ) -> PyResult<(BigInt, BigInt, BigInt)> {
99        // Calculate step
100        let step: BigInt;
101        if vm.is_none(self.step_ref(vm)) {
102            step = One::one();
103        } else {
104            // Clone the value, not the reference.
105            let this_step = self.step(vm).try_index(vm)?;
106            step = this_step.as_bigint().clone();
107
108            if step.is_zero() {
109                return Err(vm.new_value_error("slice step cannot be zero.".to_owned()));
110            }
111        }
112
113        // For convenience
114        let backwards = step.is_negative();
115
116        // Each end of the array
117        let lower = if backwards {
118            (-1_i8).to_bigint().unwrap()
119        } else {
120            Zero::zero()
121        };
122
123        let upper = if backwards {
124            lower.clone() + length
125        } else {
126            length.clone()
127        };
128
129        // Calculate start
130        let mut start: BigInt;
131        if vm.is_none(self.start_ref(vm)) {
132            // Default
133            start = if backwards {
134                upper.clone()
135            } else {
136                lower.clone()
137            };
138        } else {
139            let this_start = self.start(vm).try_index(vm)?;
140            start = this_start.as_bigint().clone();
141
142            if start < Zero::zero() {
143                // From end of array
144                start += length;
145
146                if start < lower {
147                    start = lower.clone();
148                }
149            } else if start > upper {
150                start = upper.clone();
151            }
152        }
153
154        // Calculate Stop
155        let mut stop: BigInt;
156        if vm.is_none(&self.stop) {
157            stop = if backwards { lower } else { upper };
158        } else {
159            let this_stop = self.stop(vm).try_index(vm)?;
160            stop = this_stop.as_bigint().clone();
161
162            if stop < Zero::zero() {
163                // From end of array
164                stop += length;
165                if stop < lower {
166                    stop = lower;
167                }
168            } else if stop > upper {
169                stop = upper;
170            }
171        }
172
173        Ok((start, stop, step))
174    }
175
176    #[pymethod]
177    fn indices(&self, length: ArgIndex, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
178        let length = length.as_bigint();
179        if length.is_negative() {
180            return Err(vm.new_value_error("length should not be negative.".to_owned()));
181        }
182        let (start, stop, step) = self.inner_indices(length, vm)?;
183        Ok(vm.new_tuple((start, stop, step)))
184    }
185
186    #[allow(clippy::type_complexity)]
187    #[pymethod(magic)]
188    fn reduce(
189        zelf: PyRef<Self>,
190    ) -> PyResult<(
191        PyTypeRef,
192        (Option<PyObjectRef>, PyObjectRef, Option<PyObjectRef>),
193    )> {
194        Ok((
195            zelf.class().to_owned(),
196            (zelf.start.clone(), zelf.stop.clone(), zelf.step.clone()),
197        ))
198    }
199}
200
201impl Hashable for PySlice {
202    #[inline]
203    fn hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyHash> {
204        const XXPRIME_1: PyUHash = if cfg!(target_pointer_width = "64") {
205            11400714785074694791
206        } else {
207            2654435761
208        };
209        const XXPRIME_2: PyUHash = if cfg!(target_pointer_width = "64") {
210            14029467366897019727
211        } else {
212            2246822519
213        };
214        const XXPRIME_5: PyUHash = if cfg!(target_pointer_width = "64") {
215            2870177450012600261
216        } else {
217            374761393
218        };
219        const ROTATE: u32 = if cfg!(target_pointer_width = "64") {
220            31
221        } else {
222            13
223        };
224
225        let mut acc = XXPRIME_5;
226        for part in [zelf.start_ref(vm), &zelf.stop, zelf.step_ref(vm)].iter() {
227            let lane = part.hash(vm)? as PyUHash;
228            if lane == u64::MAX as PyUHash {
229                return Ok(-1 as PyHash);
230            }
231            acc = acc.wrapping_add(lane.wrapping_mul(XXPRIME_2));
232            acc = acc.rotate_left(ROTATE);
233            acc = acc.wrapping_mul(XXPRIME_1);
234        }
235        if acc == u64::MAX as PyUHash {
236            return Ok(1546275796 as PyHash);
237        }
238        Ok(acc as PyHash)
239    }
240}
241
242impl Comparable for PySlice {
243    fn cmp(
244        zelf: &Py<Self>,
245        other: &PyObject,
246        op: PyComparisonOp,
247        vm: &VirtualMachine,
248    ) -> PyResult<PyComparisonValue> {
249        let other = class_or_notimplemented!(Self, other);
250
251        let ret = match op {
252            PyComparisonOp::Lt | PyComparisonOp::Le => None
253                .or_else(|| {
254                    vm.bool_seq_lt(zelf.start_ref(vm), other.start_ref(vm))
255                        .transpose()
256                })
257                .or_else(|| vm.bool_seq_lt(&zelf.stop, &other.stop).transpose())
258                .or_else(|| {
259                    vm.bool_seq_lt(zelf.step_ref(vm), other.step_ref(vm))
260                        .transpose()
261                })
262                .unwrap_or_else(|| Ok(op == PyComparisonOp::Le))?,
263            PyComparisonOp::Eq | PyComparisonOp::Ne => {
264                let eq = vm.identical_or_equal(zelf.start_ref(vm), other.start_ref(vm))?
265                    && vm.identical_or_equal(&zelf.stop, &other.stop)?
266                    && vm.identical_or_equal(zelf.step_ref(vm), other.step_ref(vm))?;
267                if op == PyComparisonOp::Ne {
268                    !eq
269                } else {
270                    eq
271                }
272            }
273            PyComparisonOp::Gt | PyComparisonOp::Ge => None
274                .or_else(|| {
275                    vm.bool_seq_gt(zelf.start_ref(vm), other.start_ref(vm))
276                        .transpose()
277                })
278                .or_else(|| vm.bool_seq_gt(&zelf.stop, &other.stop).transpose())
279                .or_else(|| {
280                    vm.bool_seq_gt(zelf.step_ref(vm), other.step_ref(vm))
281                        .transpose()
282                })
283                .unwrap_or_else(|| Ok(op == PyComparisonOp::Ge))?,
284        };
285
286        Ok(PyComparisonValue::Implemented(ret))
287    }
288}
289
290impl Representable for PySlice {
291    #[inline]
292    fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
293        let start_repr = zelf.start_ref(vm).repr(vm)?;
294        let stop_repr = &zelf.stop.repr(vm)?;
295        let step_repr = zelf.step_ref(vm).repr(vm)?;
296
297        Ok(format!(
298            "slice({}, {}, {})",
299            start_repr.as_str(),
300            stop_repr.as_str(),
301            step_repr.as_str()
302        ))
303    }
304}
305
306#[pyclass(module = false, name = "EllipsisType")]
307#[derive(Debug)]
308pub struct PyEllipsis;
309
310impl PyPayload for PyEllipsis {
311    fn class(ctx: &Context) -> &'static Py<PyType> {
312        ctx.types.ellipsis_type
313    }
314}
315
316impl Constructor for PyEllipsis {
317    type Args = ();
318
319    fn py_new(_cls: PyTypeRef, _args: Self::Args, vm: &VirtualMachine) -> PyResult {
320        Ok(vm.ctx.ellipsis.clone().into())
321    }
322}
323
324#[pyclass(with(Constructor, Representable))]
325impl PyEllipsis {
326    #[pymethod(magic)]
327    fn reduce(&self, vm: &VirtualMachine) -> PyStrRef {
328        vm.ctx.names.Ellipsis.to_owned()
329    }
330}
331
332impl Representable for PyEllipsis {
333    #[inline]
334    fn repr(_zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyStrRef> {
335        Ok(vm.ctx.names.Ellipsis.to_owned())
336    }
337
338    #[cold]
339    fn repr_str(_zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<String> {
340        unreachable!("use repr instead")
341    }
342}
343
344pub fn init(ctx: &Context) {
345    PySlice::extend_class(ctx, ctx.types.slice_type);
346    PyEllipsis::extend_class(ctx, ctx.types.ellipsis_type);
347}
Morty Proxy This is a proxified and sanitized view of the page, visit original site.