Skip to content

Commit 52ce150

Browse files
youknowoneyouknowone
committed
optimize range.__contains__ and iter_search
1 parent fe06583 commit 52ce150

1 file changed

Lines changed: 66 additions & 64 deletions

File tree

‎vm/src/builtins/range.rs‎

Lines changed: 66 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -32,15 +32,15 @@ enum SearchType {
3232
// and place in vm.rs for all sequences to be able to use it.
3333
#[inline]
3434
fn iter_search(
35-
obj: PyObjectRef,
36-
item: PyObjectRef,
35+
obj: &PyObject,
36+
item: &PyObject,
3737
flag: SearchType,
3838
vm: &VirtualMachine,
3939
) -> PyResult<usize> {
4040
let mut count = 0;
4141
let iter = obj.get_iter(vm)?;
4242
for element in iter.iter_without_hint::<PyObjectRef>(vm)? {
43-
if vm.bool_eq(&item, &*element?)? {
43+
if vm.bool_eq(item, &*element?)? {
4444
match flag {
4545
SearchType::Index => return Ok(count),
4646
SearchType::Contains => return Ok(1),
@@ -53,8 +53,7 @@ fn iter_search(
5353
SearchType::Contains => Ok(0),
5454
SearchType::Index => Err(vm.new_value_error(format!(
5555
"{} not in range",
56-
&item
57-
.repr(vm)
56+
item.repr(vm)
5857
.map(|v| v.as_str().to_owned())
5958
.unwrap_or_else(|_| "value".to_owned())
6059
))),
@@ -174,7 +173,15 @@ pub fn init(context: &Context) {
174173
PyRangeIterator::extend_class(context, context.types.range_iterator_type);
175174
}
176175

177-
#[pyclass(with(AsMapping, AsSequence, Hashable, Comparable, Iterable, Representable))]
176+
#[pyclass(with(
177+
Py,
178+
AsMapping,
179+
AsSequence,
180+
Hashable,
181+
Comparable,
182+
Iterable,
183+
Representable
184+
))]
178185
impl PyRange {
179186
fn new(cls: PyTypeRef, stop: ArgIndex, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
180187
PyRange {
@@ -265,26 +272,6 @@ impl PyRange {
265272
!self.is_empty()
266273
}
267274

268-
#[pymethod(magic)]
269-
fn contains(&self, needle: PyObjectRef, vm: &VirtualMachine) -> bool {
270-
// Only accept ints, not subclasses.
271-
if let Some(int) = needle.payload_if_exact::<PyInt>(vm) {
272-
match self.offset(int.as_bigint()) {
273-
Some(ref offset) => offset.is_multiple_of(self.step.as_bigint()),
274-
None => false,
275-
}
276-
} else {
277-
iter_search(
278-
self.clone().into_pyobject(vm),
279-
needle,
280-
SearchType::Contains,
281-
vm,
282-
)
283-
.unwrap_or(0)
284-
!= 0
285-
}
286-
}
287-
288275
#[pymethod(magic)]
289276
fn reduce(&self, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef) {
290277
let range_parameters: Vec<PyObjectRef> = [&self.start, &self.stop, &self.step]
@@ -295,43 +282,6 @@ impl PyRange {
295282
(vm.ctx.types.range_type.to_owned(), range_parameters_tuple)
296283
}
297284

298-
#[pymethod]
299-
fn index(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<BigInt> {
300-
if let Ok(int) = needle.clone().downcast::<PyInt>() {
301-
match self.index_of(int.as_bigint()) {
302-
Some(idx) => Ok(idx),
303-
None => Err(vm.new_value_error(format!("{int} is not in range"))),
304-
}
305-
} else {
306-
// Fallback to iteration.
307-
Ok(BigInt::from_bytes_be(
308-
Sign::Plus,
309-
&iter_search(
310-
self.clone().into_pyobject(vm),
311-
needle,
312-
SearchType::Index,
313-
vm,
314-
)?
315-
.to_be_bytes(),
316-
))
317-
}
318-
}
319-
320-
#[pymethod]
321-
fn count(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
322-
if let Ok(int) = item.clone().downcast::<PyInt>() {
323-
if self.index_of(int.as_bigint()).is_some() {
324-
Ok(1)
325-
} else {
326-
Ok(0)
327-
}
328-
} else {
329-
// Dealing with classes who might compare equal with ints in their
330-
// __eq__, slow search.
331-
iter_search(self.clone().into_pyobject(vm), item, SearchType::Count, vm)
332-
}
333-
}
334-
335285
#[pymethod(magic)]
336286
fn getitem(&self, subscript: PyObjectRef, vm: &VirtualMachine) -> PyResult {
337287
match RangeIndex::try_from_object(vm, subscript)? {
@@ -374,6 +324,58 @@ impl PyRange {
374324
}
375325
}
376326

327+
#[pyclass]
328+
impl Py<PyRange> {
329+
fn contains_inner(&self, needle: &PyObject, vm: &VirtualMachine) -> bool {
330+
// Only accept ints, not subclasses.
331+
if let Some(int) = needle.downcast_ref_if_exact::<PyInt>(vm) {
332+
match self.offset(int.as_bigint()) {
333+
Some(ref offset) => offset.is_multiple_of(self.step.as_bigint()),
334+
None => false,
335+
}
336+
} else {
337+
iter_search(self.as_object(), needle, SearchType::Contains, vm).unwrap_or(0) != 0
338+
}
339+
}
340+
341+
#[pymethod(magic)]
342+
fn contains(&self, needle: PyObjectRef, vm: &VirtualMachine) -> bool {
343+
self.contains_inner(&needle, vm)
344+
}
345+
346+
#[pymethod]
347+
fn index(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<BigInt> {
348+
if let Ok(int) = needle.clone().downcast::<PyInt>() {
349+
match self.index_of(int.as_bigint()) {
350+
Some(idx) => Ok(idx),
351+
None => Err(vm.new_value_error(format!("{int} is not in range"))),
352+
}
353+
} else {
354+
// Fallback to iteration.
355+
Ok(BigInt::from_bytes_be(
356+
Sign::Plus,
357+
&iter_search(self.as_object(), &needle, SearchType::Index, vm)?.to_be_bytes(),
358+
))
359+
}
360+
}
361+
362+
#[pymethod]
363+
fn count(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
364+
if let Ok(int) = item.clone().downcast::<PyInt>() {
365+
let count = if self.index_of(int.as_bigint()).is_some() {
366+
1
367+
} else {
368+
0
369+
};
370+
Ok(count)
371+
} else {
372+
// Dealing with classes who might compare equal with ints in their
373+
// __eq__, slow search.
374+
iter_search(self.as_object(), &item, SearchType::Count, vm)
375+
}
376+
}
377+
}
378+
377379
impl PyRange {
378380
fn protocol_length(&self, vm: &VirtualMachine) -> PyResult<usize> {
379381
PyInt::from(self.len())
@@ -408,7 +410,7 @@ impl AsSequence for PyRange {
408410
.ok_or_else(|| vm.new_index_error("index out of range".to_owned()))
409411
}),
410412
contains: atomic_func!(|seq, needle, vm| {
411-
Ok(PyRange::sequence_downcast(seq).contains(needle.to_owned(), vm))
413+
Ok(PyRange::sequence_downcast(seq).contains_inner(needle, vm))
412414
}),
413415
..PySequenceMethods::NOT_IMPLEMENTED
414416
});

0 commit comments

Comments
 (0)