Skip to content

Commit 4f97c0f

Browse files
committed
Extract saturate function.
1 parent 83e2490 commit 4f97c0f

File tree

2 files changed

+19
-43
lines changed

2 files changed

+19
-43
lines changed

vm/src/builtins/list.rs

Lines changed: 5 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use crate::{
1010
protocol::{PyIterReturn, PyMappingMethods, PySequence, PySequenceMethods},
1111
recursion::ReprGuard,
1212
sequence::{MutObjectSequenceOp, SequenceExt, SequenceMutExt},
13-
sliceable::{SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
13+
sliceable::{pyint_saturate_index, SequenceIndex, SliceableSequenceMutOp, SliceableSequenceOp},
1414
types::{
1515
AsMapping, AsSequence, Comparable, Constructor, Hashable, Initializer, IterNext,
1616
IterNextIterable, Iterable, PyComparisonOp, Unconstructible, Unhashable,
@@ -19,7 +19,6 @@ use crate::{
1919
vm::VirtualMachine,
2020
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
2121
};
22-
use num_traits::Signed;
2322
use std::{fmt, ops::DerefMut};
2423

2524
/// Built-in mutable sequence.
@@ -277,54 +276,18 @@ impl PyList {
277276
vm: &VirtualMachine,
278277
) -> PyResult<usize> {
279278
let len = self.len();
280-
let start: OptionalArg<PyIntRef> = match start {
281-
OptionalArg::Missing => OptionalArg::Missing,
282-
OptionalArg::Present(obj) => OptionalArg::Present(obj.try_into_value(vm)?),
283-
};
284-
let stop: OptionalArg<PyIntRef> = match stop {
285-
OptionalArg::Missing => OptionalArg::Missing,
286-
OptionalArg::Present(obj) => OptionalArg::Present(obj.try_into_value(vm)?),
287-
};
288279
let start: usize = match start {
289280
OptionalArg::Missing => 0,
290281
OptionalArg::Present(obj) => {
291-
let bigint = obj.as_bigint();
292-
if bigint.is_negative() {
293-
let idx: Result<usize, _> = bigint.abs().try_into();
294-
match idx {
295-
Ok(abs) => {
296-
if abs >= len {
297-
0
298-
} else {
299-
len - abs
300-
}
301-
}
302-
Err(_) => 0,
303-
}
304-
} else {
305-
bigint.try_into().unwrap_or(len)
306-
}
282+
let int: PyIntRef = obj.try_into_value(vm)?;
283+
pyint_saturate_index(int, len)
307284
}
308285
};
309286
let stop: usize = match stop {
310287
OptionalArg::Missing => len,
311288
OptionalArg::Present(obj) => {
312-
let bigint = obj.as_bigint();
313-
if bigint.is_negative() {
314-
let idx: Result<usize, _> = bigint.abs().try_into();
315-
match idx {
316-
Ok(abs) => {
317-
if abs >= len {
318-
0
319-
} else {
320-
len - abs
321-
}
322-
}
323-
Err(_) => 0,
324-
}
325-
} else {
326-
bigint.try_into().unwrap_or(len)
327-
}
289+
let int: PyIntRef = obj.try_into_value(vm)?;
290+
pyint_saturate_index(int, len)
328291
}
329292
};
330293
let index = self.mut_index_range(vm, &needle, start..stop)?;

vm/src/sliceable.rs

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
// export through slicable module, not slice.
22
use crate::{
3-
builtins::{int::PyInt, slice::PySlice},
3+
builtins::{int::PyInt, slice::PySlice, PyIntRef},
44
AsObject, PyObject, PyResult, VirtualMachine,
55
};
66
use num_traits::{Signed, ToPrimitive};
@@ -320,6 +320,19 @@ pub fn saturate_index(p: isize, len: usize) -> usize {
320320
p as usize
321321
}
322322

323+
// Saturate p in range [0, len] inclusive
324+
pub fn pyint_saturate_index(p: PyIntRef, len: usize) -> usize {
325+
let bigint = p.as_bigint();
326+
if bigint.is_negative() {
327+
bigint
328+
.abs()
329+
.try_into()
330+
.map_or(0, |abs| len.saturating_sub(abs))
331+
} else {
332+
bigint.try_into().unwrap_or(len)
333+
}
334+
}
335+
323336
/// A saturated slice with values ranging in [isize::MIN, isize::MAX]. Used for
324337
/// slicable sequences that require indices in the aforementioned range.
325338
///

0 commit comments

Comments
 (0)