RustPython/vm/src/obj/objrange.rs

use std::cell::Cell;
use std::ops::Mul;

use num_bigint::{BigInt, Sign};
use num_integer::Integer;
use num_traits::{One, Signed, Zero};

use crate::function::{OptionalArg, PyFuncArgs};
use crate::pyobject::{Either, PyContext, PyIteratorValue, PyObjectRef, PyRef, PyResult, PyValue};
use crate::vm::VirtualMachine;

use super::objint::{PyInt, PyIntRef};
use super::objslice::PySliceRef;
use super::objtype::PyClassRef;

#[derive(Debug, Clone)]
pub struct PyRange {
    pub start: PyIntRef,
    pub stop: PyIntRef,
    pub step: PyIntRef,
}

impl PyValue for PyRange {
    fn class(vm: &VirtualMachine) -> PyClassRef {
        vm.ctx.range_type()
    }
}

impl PyRange {
    #[inline]
    fn offset(&self, value: &BigInt) -> Option<BigInt> {
        let start = self.start.as_bigint();
        let stop = self.stop.as_bigint();
        let step = self.step.as_bigint();
        match step.sign() {
            Sign::Plus if value >= start && value < stop => Some(value - start),
            Sign::Minus if value <= self.start.as_bigint() && value > stop => Some(start - value),
            _ => None,
        }
    }

    #[inline]
    pub fn index_of(&self, value: &BigInt) -> Option<BigInt> {
        let step = self.step.as_bigint();
        match self.offset(value) {
            Some(ref offset) if offset.is_multiple_of(step) => Some((offset / step).abs()),
            Some(_) | None => None,
        }
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        let start = self.start.as_bigint();
        let stop = self.stop.as_bigint();
        let step = self.step.as_bigint();
        (start <= stop && step.is_negative()) || (start >= stop && step.is_positive())
    }

    #[inline]
    pub fn forward(&self) -> bool {
        self.start.as_bigint() < self.stop.as_bigint()
    }

    #[inline]
    pub fn get<'a, T>(&'a self, index: T) -> Option<BigInt>
    where
        &'a BigInt: Mul<T, Output = BigInt>,
    {
        let start = self.start.as_bigint();
        let stop = self.stop.as_bigint();
        let step = self.step.as_bigint();

        let result = start + step * index;

        if (self.forward() && !self.is_empty() && &result < stop)
            || (!self.forward() && !self.is_empty() && &result > stop)
        {
            Some(result)
        } else {
            None
        }
    }
}

pub fn get_value(obj: &PyObjectRef) -> PyRange {
    obj.payload::<PyRange>().unwrap().clone()
}

pub fn init(context: &PyContext) {
    let range_type = &context.range_type;

    let range_doc = "range(stop) -> range object\n\
                     range(start, stop[, step]) -> range object\n\n\
                     Return an object that produces a sequence of integers from start (inclusive)\n\
                     to stop (exclusive) by step.  range(i, j) produces i, i+1, i+2, ..., j-1.\n\
                     start defaults to 0, and stop is omitted!  range(4) produces 0, 1, 2, 3.\n\
                     These are exactly the valid indices for a list of 4 elements.\n\
                     When step is given, it specifies the increment (or decrement).";

    extend_class!(context, range_type, {
        "__bool__" => context.new_rustfunc(PyRangeRef::bool),
        "__contains__" => context.new_rustfunc(PyRangeRef::contains),
        "__doc__" => context.new_str(range_doc.to_string()),
        "__getitem__" => context.new_rustfunc(PyRangeRef::getitem),
        "__iter__" => context.new_rustfunc(PyRangeRef::iter),
        "__len__" => context.new_rustfunc(PyRangeRef::len),
        "__new__" => context.new_rustfunc(range_new),
        "__repr__" => context.new_rustfunc(PyRangeRef::repr),
        "__reversed__" => context.new_rustfunc(PyRangeRef::reversed),
        "count" => context.new_rustfunc(PyRangeRef::count),
        "index" => context.new_rustfunc(PyRangeRef::index),
        "start" => context.new_property(PyRangeRef::start),
        "stop" => context.new_property(PyRangeRef::stop),
        "step" => context.new_property(PyRangeRef::step),
    });
}

type PyRangeRef = PyRef<PyRange>;

impl PyRangeRef {
    fn new(cls: PyClassRef, stop: PyIntRef, vm: &VirtualMachine) -> PyResult<PyRangeRef> {
        PyRange {
            start: PyInt::new(BigInt::zero()).into_ref(vm),
            stop: stop.clone(),
            step: PyInt::new(BigInt::one()).into_ref(vm),
        }
        .into_ref_with_type(vm, cls)
    }

    fn new_from(
        cls: PyClassRef,
        start: PyIntRef,
        stop: PyIntRef,
        step: OptionalArg<PyIntRef>,
        vm: &VirtualMachine,
    ) -> PyResult<PyRangeRef> {
        PyRange {
            start: start,
            stop: stop,
            step: step
                .into_option()
                .unwrap_or_else(|| PyInt::new(BigInt::one()).into_ref(vm)),
        }
        .into_ref_with_type(vm, cls)
    }

    fn start(self, _vm: &VirtualMachine) -> PyIntRef {
        self.start.clone()
    }

    fn stop(self, _vm: &VirtualMachine) -> PyIntRef {
        self.stop.clone()
    }

    fn step(self, _vm: &VirtualMachine) -> PyIntRef {
        self.step.clone()
    }

    fn iter(self: PyRangeRef, _vm: &VirtualMachine) -> PyIteratorValue {
        PyIteratorValue {
            position: Cell::new(0),
            iterated_obj: self.into_object(),
        }
    }

    fn reversed(self: PyRangeRef, vm: &VirtualMachine) -> PyIteratorValue {
        let start = self.start.as_bigint();
        let stop = self.stop.as_bigint();
        let step = self.step.as_bigint();

        // compute the last element that is actually contained within the range
        // this is the new start
        let remainder = ((stop - start) % step).abs();
        let new_start = if remainder.is_zero() {
            stop - step
        } else {
            stop - &remainder
        };

        let new_stop: BigInt = match step.sign() {
            Sign::Plus => start - 1,
            Sign::Minus => start + 1,
            Sign::NoSign => unreachable!(),
        };

        let reversed = PyRange {
            start: PyInt::new(new_start).into_ref(vm),
            stop: PyInt::new(new_stop).into_ref(vm),
            step: PyInt::new(-step).into_ref(vm),
        };

        PyIteratorValue {
            position: Cell::new(0),
            iterated_obj: reversed.into_ref(vm).into_object(),
        }
    }

    fn len(self, _vm: &VirtualMachine) -> PyInt {
        let start = self.start.as_bigint();
        let stop = self.stop.as_bigint();
        let step = self.step.as_bigint();

        match step.sign() {
            Sign::Plus if start < stop => PyInt::new((stop - start - 1usize) / step + 1),
            Sign::Minus if start > stop => PyInt::new((start - stop - 1usize) / (-step) + 1),
            Sign::Plus | Sign::Minus => PyInt::new(0),
            Sign::NoSign => unreachable!(),
        }
    }

    fn repr(self, _vm: &VirtualMachine) -> String {
        if self.step.as_bigint().is_one() {
            format!("range({}, {})", self.start, self.stop)
        } else {
            format!("range({}, {}, {})", self.start, self.stop, self.step)
        }
    }

    fn bool(self, _vm: &VirtualMachine) -> bool {
        !self.is_empty()
    }

    fn contains(self, needle: PyObjectRef, _vm: &VirtualMachine) -> bool {
        if let Ok(int) = needle.downcast::<PyInt>() {
            match self.offset(int.as_bigint()) {
                Some(ref offset) => offset.is_multiple_of(self.step.as_bigint()),
                None => false,
            }
        } else {
            false
        }
    }

    fn index(self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyInt> {
        if let Ok(int) = needle.downcast::<PyInt>() {
            match self.index_of(int.as_bigint()) {
                Some(idx) => Ok(PyInt::new(idx)),
                None => Err(vm.new_value_error(format!("{} is not in range", int))),
            }
        } else {
            Err(vm.new_value_error("sequence.index(x): x not in sequence".to_string()))
        }
    }

    fn count(self, item: PyObjectRef, _vm: &VirtualMachine) -> PyInt {
        if let Ok(int) = item.downcast::<PyInt>() {
            if self.index_of(int.as_bigint()).is_some() {
                PyInt::new(1)
            } else {
                PyInt::new(0)
            }
        } else {
            PyInt::new(0)
        }
    }

    fn getitem(self, subscript: Either<PyIntRef, PySliceRef>, vm: &VirtualMachine) -> PyResult {
        match subscript {
            Either::A(index) => {
                if let Some(value) = self.get(index.as_bigint()) {
                    Ok(PyInt::new(value).into_ref(vm).into_object())
                } else {
                    Err(vm.new_index_error("range object index out of range".to_string()))
                }
            }
            Either::B(slice) => {
                let new_start = if let Some(int) = slice.start.as_ref() {
                    if let Some(i) = self.get(int) {
                        PyInt::new(i).into_ref(vm)
                    } else {
                        self.start.clone()
                    }
                } else {
                    self.start.clone()
                };

                let new_end = if let Some(int) = slice.stop.as_ref() {
                    if let Some(i) = self.get(int) {
                        PyInt::new(i).into_ref(vm)
                    } else {
                        self.stop.clone()
                    }
                } else {
                    self.stop.clone()
                };

                let new_step = if let Some(int) = slice.step.as_ref() {
                    PyInt::new(int * self.step.as_bigint()).into_ref(vm)
                } else {
                    self.step.clone()
                };

                Ok(PyRange {
                    start: new_start,
                    stop: new_end,
                    step: new_step,
                }
                .into_ref(vm)
                .into_object())
            }
        }
    }
}

fn range_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
    let range = if args.args.len() <= 2 {
        let (cls, stop) = args.bind(vm)?;
        PyRangeRef::new(cls, stop, vm)
    } else {
        let (cls, start, stop, step) = args.bind(vm)?;
        PyRangeRef::new_from(cls, start, stop, step, vm)
    }?;

    Ok(range.into_object())
}