Skip to content

Commit 18ed00a

Browse files
OddCoincidenceOddCoincidence
committed
Range cleanups
1 parent cecc2e1 commit 18ed00a

File tree

2 files changed

+62
-57
lines changed

2 files changed

+62
-57
lines changed

‎vm/src/function.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ impl PyFuncArgs {
143143
///
144144
/// If the given `FromArgs` includes any conversions, exceptions raised
145145
/// during the conversion will halt the binding and return the error.
146-
fn bind<T: FromArgs>(mut self, vm: &VirtualMachine) -> PyResult<T> {
146+
pub fn bind<T: FromArgs>(mut self, vm: &VirtualMachine) -> PyResult<T> {
147147
let given_args = self.args.len();
148148
let bound = match T::from_args(vm, &mut self) {
149149
Ok(args) => args,

‎vm/src/obj/objrange.rs

Lines changed: 61 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -5,25 +5,23 @@ use num_bigint::{BigInt, Sign};
55
use num_integer::Integer;
66
use num_traits::{One, Signed, ToPrimitive, Zero};
77

8-
use crate::function::PyFuncArgs;
8+
use crate::function::{OptionalArg, PyFuncArgs};
99
use crate::pyobject::{
10-
PyContext, PyIteratorValue, PyObject, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
10+
PyContext, PyIteratorValue, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
1111
};
1212
use crate::vm::VirtualMachine;
1313

14-
use super::objint::{self, PyInt};
14+
use super::objint::{self, PyInt, PyIntRef};
1515
use super::objslice::PySlice;
1616
use super::objtype;
17-
use crate::obj::objtype::PyClassRef;
18-
19-
pub type PyRangeRef = PyRef<PyRange>;
17+
use super::objtype::PyClassRef;
2018

2119
#[derive(Debug, Clone)]
2220
pub struct PyRange {
2321
// Unfortunately Rust's built in range type doesn't support things like indexing
2422
// or ranges where start > end so we need to roll our own.
2523
pub start: BigInt,
26-
pub end: BigInt,
24+
pub stop: BigInt,
2725
pub step: BigInt,
2826
}
2927

@@ -37,11 +35,11 @@ impl PyRange {
3735
#[inline]
3836
pub fn try_len(&self) -> Option<usize> {
3937
match self.step.sign() {
40-
Sign::Plus if self.start < self.end => ((&self.end - &self.start - 1usize)
38+
Sign::Plus if self.start < self.stop => ((&self.stop - &self.start - 1usize)
4139
/ &self.step)
4240
.to_usize()
4341
.map(|sz| sz + 1),
44-
Sign::Minus if self.start > self.end => ((&self.start - &self.end - 1usize)
42+
Sign::Minus if self.start > self.stop => ((&self.start - &self.stop - 1usize)
4543
/ (-&self.step))
4644
.to_usize()
4745
.map(|sz| sz + 1),
@@ -57,8 +55,8 @@ impl PyRange {
5755
#[inline]
5856
fn offset(&self, value: &BigInt) -> Option<BigInt> {
5957
match self.step.sign() {
60-
Sign::Plus if *value >= self.start && *value < self.end => Some(value - &self.start),
61-
Sign::Minus if *value <= self.start && *value > self.end => Some(&self.start - value),
58+
Sign::Plus if *value >= self.start && *value < self.stop => Some(value - &self.start),
59+
Sign::Minus if *value <= self.start && *value > self.stop => Some(&self.start - value),
6260
_ => None,
6361
}
6462
}
@@ -92,13 +90,13 @@ impl PyRange {
9290

9391
#[inline]
9492
pub fn is_empty(&self) -> bool {
95-
(self.start <= self.end && self.step.is_negative())
96-
|| (self.start >= self.end && self.step.is_positive())
93+
(self.start <= self.stop && self.step.is_negative())
94+
|| (self.start >= self.stop && self.step.is_positive())
9795
}
9896

9997
#[inline]
10098
pub fn forward(&self) -> bool {
101-
self.start < self.end
99+
self.start < self.stop
102100
}
103101

104102
#[inline]
@@ -108,8 +106,8 @@ impl PyRange {
108106
{
109107
let result = &self.start + &self.step * index;
110108

111-
if (self.forward() && !self.is_empty() && result < self.end)
112-
|| (!self.forward() && !self.is_empty() && result > self.end)
109+
if (self.forward() && !self.is_empty() && result < self.stop)
110+
|| (!self.forward() && !self.is_empty() && result > self.stop)
113111
{
114112
Some(result)
115113
} else {
@@ -121,22 +119,22 @@ impl PyRange {
121119
pub fn reversed(&self) -> Self {
122120
// compute the last element that is actually contained within the range
123121
// this is the new start
124-
let remainder = ((&self.end - &self.start) % &self.step).abs();
122+
let remainder = ((&self.stop - &self.start) % &self.step).abs();
125123
let start = if remainder.is_zero() {
126-
&self.end - &self.step
124+
&self.stop - &self.step
127125
} else {
128-
&self.end - &remainder
126+
&self.stop - &remainder
129127
};
130128

131129
match self.step.sign() {
132130
Sign::Plus => PyRange {
133131
start,
134-
end: &self.start - 1,
132+
stop: &self.start - 1,
135133
step: -&self.step,
136134
},
137135
Sign::Minus => PyRange {
138136
start,
139-
end: &self.start + 1,
137+
stop: &self.start + 1,
140138
step: -&self.step,
141139
},
142140
Sign::NoSign => unreachable!(),
@@ -145,9 +143,9 @@ impl PyRange {
145143

146144
pub fn repr(&self) -> String {
147145
if self.step == BigInt::one() {
148-
format!("range({}, {})", self.start, self.end)
146+
format!("range({}, {})", self.start, self.stop)
149147
} else {
150-
format!("range({}, {}, {})", self.start, self.end, self.step)
148+
format!("range({}, {}, {})", self.start, self.stop, self.step)
151149
}
152150
}
153151
}
@@ -185,40 +183,47 @@ pub fn init(context: &PyContext) {
185183
});
186184
}
187185

188-
fn range_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
189-
arg_check!(
190-
vm,
191-
args,
192-
required = [(cls, None), (first, Some(vm.ctx.int_type()))],
193-
optional = [
194-
(second, Some(vm.ctx.int_type())),
195-
(step, Some(vm.ctx.int_type()))
196-
]
197-
);
186+
type PyRangeRef = PyRef<PyRange>;
198187

199-
let start = if second.is_some() {
200-
objint::get_value(first).clone()
201-
} else {
202-
BigInt::zero()
203-
};
188+
impl PyRangeRef {
189+
fn new(cls: PyClassRef, stop: PyIntRef, vm: &VirtualMachine) -> PyResult<PyRangeRef> {
190+
PyRange {
191+
start: Zero::zero(),
192+
stop: stop.value.clone(),
193+
step: One::one(),
194+
}
195+
.into_ref_with_type(vm, cls)
196+
}
204197

205-
let end = if let Some(pyint) = second {
206-
objint::get_value(pyint).clone()
207-
} else {
208-
objint::get_value(first).clone()
209-
};
198+
fn new_from(
199+
cls: PyClassRef,
200+
start: PyIntRef,
201+
stop: PyIntRef,
202+
step: OptionalArg<PyIntRef>,
203+
vm: &VirtualMachine,
204+
) -> PyResult<PyRangeRef> {
205+
PyRange {
206+
start: start.value.clone(),
207+
stop: stop.value.clone(),
208+
step: step
209+
.into_option()
210+
.map(|i| i.value.clone())
211+
.unwrap_or_else(One::one),
212+
}
213+
.into_ref_with_type(vm, cls)
214+
}
215+
}
210216

211-
let step = if let Some(pyint) = step {
212-
objint::get_value(pyint).clone()
217+
fn range_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
218+
let range = if args.args.len() <= 2 {
219+
let (cls, stop) = args.bind(vm)?;
220+
PyRangeRef::new(cls, stop, vm)
213221
} else {
214-
BigInt::one()
215-
};
222+
let (cls, start, stop, step) = args.bind(vm)?;
223+
PyRangeRef::new_from(cls, start, stop, step, vm)
224+
}?;
216225

217-
if step.is_zero() {
218-
Err(vm.new_value_error("range with 0 step size".to_string()))
219-
} else {
220-
Ok(PyObject::new(PyRange { start, end, step }, cls.clone()))
221-
}
226+
Ok(range.into_object())
222227
}
223228

224229
fn range_iter(range: PyRangeRef, _vm: &VirtualMachine) -> PyIteratorValue {
@@ -282,10 +287,10 @@ fn range_getitem(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
282287
if let Some(i) = range.get(int) {
283288
i
284289
} else {
285-
range.end
290+
range.stop
286291
}
287292
} else {
288-
range.end
293+
range.stop
289294
};
290295

291296
let new_step = if let Some(int) = step {
@@ -296,7 +301,7 @@ fn range_getitem(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
296301

297302
Ok(PyRange {
298303
start: new_start,
299-
end: new_end,
304+
stop: new_end,
300305
step: new_step,
301306
}
302307
.into_ref(vm)
@@ -384,7 +389,7 @@ fn range_start(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
384389

385390
fn range_stop(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
386391
arg_check!(vm, args, required = [(zelf, Some(vm.ctx.range_type()))]);
387-
Ok(vm.ctx.new_int(get_value(zelf).end))
392+
Ok(vm.ctx.new_int(get_value(zelf).stop))
388393
}
389394

390395
fn range_step(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {

0 commit comments

Comments
 (0)