Skip to content

Commit 613109d

Browse files
committed
Auto merge of #24720 - critiqjo:stepby-sizehint, r=alexcrichton
`Iterator::size_hint` can be easily implemented for `StepBy`. #23708
2 parents 1320c29 + 2a8fc9b commit 613109d

File tree

2 files changed

+54
-16
lines changed

2 files changed

+54
-16
lines changed

src/libcore/iter.rs

+41-15
Original file line numberDiff line numberDiff line change
@@ -2402,14 +2402,10 @@ pub trait Step: PartialOrd {
24022402
/// Steps `self` if possible.
24032403
fn step(&self, by: &Self) -> Option<Self>;
24042404

2405-
/// Returns the number of steps between two step objects.
2405+
/// Returns the number of steps between two step objects. The count is
2406+
/// inclusive of `start` and exclusive of `end`.
24062407
///
2407-
/// `start` should always be less than `end`, so the result should never
2408-
/// be negative.
2409-
///
2410-
/// `by` must be > 0.
2411-
///
2412-
/// Returns `None` if it is not possible to calculate steps_between
2408+
/// Returns `None` if it is not possible to calculate `steps_between`
24132409
/// without overflow.
24142410
fn steps_between(start: &Self, end: &Self, by: &Self) -> Option<usize>;
24152411
}
@@ -2424,9 +2420,16 @@ macro_rules! step_impl_unsigned {
24242420
#[inline]
24252421
#[allow(trivial_numeric_casts)]
24262422
fn steps_between(start: &$t, end: &$t, by: &$t) -> Option<usize> {
2427-
if *start <= *end {
2423+
if *by == 0 { return None; }
2424+
if *start < *end {
24282425
// Note: We assume $t <= usize here
2429-
Some((*end - *start) as usize / (*by as usize))
2426+
let diff = (*end - *start) as usize;
2427+
let by = *by as usize;
2428+
if diff % by > 0 {
2429+
Some(diff / by + 1)
2430+
} else {
2431+
Some(diff / by)
2432+
}
24302433
} else {
24312434
Some(0)
24322435
}
@@ -2444,16 +2447,29 @@ macro_rules! step_impl_signed {
24442447
#[inline]
24452448
#[allow(trivial_numeric_casts)]
24462449
fn steps_between(start: &$t, end: &$t, by: &$t) -> Option<usize> {
2447-
if *start <= *end {
2450+
if *by == 0 { return None; }
2451+
let mut diff: usize;
2452+
let mut by_u: usize;
2453+
if *by > 0 {
2454+
if *start >= *end {
2455+
return Some(0);
2456+
}
24482457
// Note: We assume $t <= isize here
24492458
// Use .wrapping_sub and cast to usize to compute the
24502459
// difference that may not fit inside the range of isize.
2451-
Some(
2452-
((*end as isize).wrapping_sub(*start as isize) as usize
2453-
/ (*by as usize))
2454-
)
2460+
diff = (*end as isize).wrapping_sub(*start as isize) as usize;
2461+
by_u = *by as usize;
24552462
} else {
2456-
Some(0)
2463+
if *start <= *end {
2464+
return Some(0);
2465+
}
2466+
diff = (*start as isize).wrapping_sub(*end as isize) as usize;
2467+
by_u = (*by as isize).wrapping_mul(-1) as usize;
2468+
}
2469+
if diff % by_u > 0 {
2470+
Some(diff / by_u + 1)
2471+
} else {
2472+
Some(diff / by_u)
24572473
}
24582474
}
24592475
}
@@ -2675,6 +2691,16 @@ impl<A: Step + Zero + Clone> Iterator for StepBy<A, ops::Range<A>> {
26752691
None
26762692
}
26772693
}
2694+
2695+
#[inline]
2696+
fn size_hint(&self) -> (usize, Option<usize>) {
2697+
match Step::steps_between(&self.range.start,
2698+
&self.range.end,
2699+
&self.step_by) {
2700+
Some(hint) => (hint, Some(hint)),
2701+
None => (0, None)
2702+
}
2703+
}
26782704
}
26792705

26802706
macro_rules! range_exact_iter_impl {

src/libcoretest/iter.rs

+13-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
use core::iter::*;
1212
use core::iter::order::*;
1313
use core::iter::MinMaxResult::*;
14-
use core::isize;
14+
use core::{i8, i16, isize};
1515
use core::usize;
1616
use core::cmp;
1717

@@ -786,6 +786,18 @@ fn test_range_step() {
786786
assert_eq!((200..255).step_by(50).collect::<Vec<u8>>(), [200, 250]);
787787
assert_eq!((200..-5).step_by(1).collect::<Vec<isize>>(), []);
788788
assert_eq!((200..200).step_by(1).collect::<Vec<isize>>(), []);
789+
790+
assert_eq!((0..20).step_by(1).size_hint(), (20, Some(20)));
791+
assert_eq!((0..20).step_by(21).size_hint(), (1, Some(1)));
792+
assert_eq!((0..20).step_by(5).size_hint(), (4, Some(4)));
793+
assert_eq!((20..0).step_by(-5).size_hint(), (4, Some(4)));
794+
assert_eq!((20..0).step_by(-6).size_hint(), (4, Some(4)));
795+
assert_eq!((20..-5).step_by(1).size_hint(), (0, Some(0)));
796+
assert_eq!((20..20).step_by(1).size_hint(), (0, Some(0)));
797+
assert_eq!((0..1).step_by(0).size_hint(), (0, None));
798+
assert_eq!((i8::MAX..i8::MIN).step_by(i8::MIN).size_hint(), (2, Some(2)));
799+
assert_eq!((i16::MIN..i16::MAX).step_by(i16::MAX).size_hint(), (3, Some(3)));
800+
assert_eq!((isize::MIN..isize::MAX).step_by(1).size_hint(), (usize::MAX, Some(usize::MAX)));
789801
}
790802

791803
#[test]

0 commit comments

Comments
 (0)