From 37771d42d72d80a3025cf1bbf2aa58df79669a4e Mon Sep 17 00:00:00 2001 From: oberien Date: Thu, 18 Jan 2018 20:49:32 +0100 Subject: [PATCH 1/6] Specialize StepBy::nth --- src/libcore/iter/mod.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 06c29b47bf921..443b1567e1b4e 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -694,6 +694,18 @@ impl Iterator for StepBy where I: Iterator { (f(inner_hint.0), inner_hint.1.map(f)) } } + + #[inline] + fn nth(&mut self, mut n: usize) -> Option { + if self.first_take { + if n == 0 { + self.first_take = false; + return self.iter.next() + } + n -= 1; + } + self.iter.nth(n * self.step) + } } // StepBy can only make the iterator shorter, so the len will still fit. From 5850f0b742430a1b1b5ae7d2491dce6ca10e20d3 Mon Sep 17 00:00:00 2001 From: oberien Date: Fri, 19 Jan 2018 14:55:20 +0100 Subject: [PATCH 2/6] Fix off-by-ones --- src/libcore/iter/mod.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 443b1567e1b4e..33c00989fd286 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -698,13 +698,16 @@ impl Iterator for StepBy where I: Iterator { #[inline] fn nth(&mut self, mut n: usize) -> Option { if self.first_take { + self.first_take = false; + let first = self.iter.next(); if n == 0 { - self.first_take = false; - return self.iter.next() + return first; } n -= 1; } - self.iter.nth(n * self.step) + // n and self.step are indices, thus we need to add 1 before multiplying. + // After that we need to subtract 1 from the result to convert it back to an index. + self.iter.nth((n + 1) * (self.step + 1) - 1) } } From d33cc12eed3df459db3c9ae2dd89df9cc6e45dd6 Mon Sep 17 00:00:00 2001 From: oberien Date: Fri, 19 Jan 2018 14:55:34 +0100 Subject: [PATCH 3/6] Unit Tests --- src/libcore/tests/iter.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index 8997cf9c6bff9..e887499165966 100644 --- a/src/libcore/tests/iter.rs +++ b/src/libcore/tests/iter.rs @@ -161,6 +161,24 @@ fn test_iterator_step_by() { assert_eq!(it.next(), None); } +#[test] +fn test_iterator_step_by_nth() { + let mut it = (0..16).step_by(5); + assert_eq!(it.nth(0), Some(0)); + assert_eq!(it.nth(0), Some(5)); + assert_eq!(it.nth(0), Some(10)); + assert_eq!(it.nth(0), Some(15)); + assert_eq!(it.nth(0), None); + + let it = (0..18).step_by(5); + assert_eq!(it.clone().nth(0), Some(0)); + assert_eq!(it.clone().nth(1), Some(5)); + assert_eq!(it.clone().nth(2), Some(10)); + assert_eq!(it.clone().nth(3), Some(15)); + assert_eq!(it.clone().nth(4), None); + assert_eq!(it.clone().nth(42), None); +} + #[test] #[should_panic] fn test_iterator_step_by_zero() { From f08dec114f6008cb7a906ffd55c221fd30d70987 Mon Sep 17 00:00:00 2001 From: oberien Date: Fri, 19 Jan 2018 21:07:01 +0100 Subject: [PATCH 4/6] Handle Overflow --- src/libcore/iter/mod.rs | 31 ++++++++++++++++++++++++--- src/libcore/tests/iter.rs | 44 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 3 deletions(-) diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 33c00989fd286..57e7e03a6cebc 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -705,9 +705,34 @@ impl Iterator for StepBy where I: Iterator { } n -= 1; } - // n and self.step are indices, thus we need to add 1 before multiplying. - // After that we need to subtract 1 from the result to convert it back to an index. - self.iter.nth((n + 1) * (self.step + 1) - 1) + // n and self.step are indices, we need to add 1 to get the amount of elements + // When calling `.nth`, we need to subtract 1 again to convert back to an index + // step + 1 can't overflow because `.step_by` sets `self.step` to `step - 1` + let mut step = self.step + 1; + // n + 1 could overflow + // thus, if n is usize::MAX, instead of adding one, we call .nth(step) + if n == usize::MAX { + self.iter.nth(step - 1); + } else { + n += 1; + } + + // overflow handling + while n.checked_mul(step).is_none() { + let div_n = usize::MAX / n; + let div_step = usize::MAX / step; + let nth_n = div_n * n; + let nth_step = div_step * step; + let nth = if nth_n > nth_step { + step -= div_n; + nth_n + } else { + n -= div_step; + nth_step + }; + self.iter.nth(nth - 1); + } + self.iter.nth(n * step - 1) } } diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index e887499165966..e52e119ff59b9 100644 --- a/src/libcore/tests/iter.rs +++ b/src/libcore/tests/iter.rs @@ -179,6 +179,50 @@ fn test_iterator_step_by_nth() { assert_eq!(it.clone().nth(42), None); } +#[test] +fn test_iterator_step_by_nth_overflow() { + #[cfg(target_pointer_width = "8")] + type Bigger = u16; + #[cfg(target_pointer_width = "16")] + type Bigger = u32; + #[cfg(target_pointer_width = "32")] + type Bigger = u64; + #[cfg(target_pointer_width = "64")] + type Bigger = u128; + + #[derive(Clone)] + struct Test(Bigger); + impl<'a> Iterator for &'a mut Test { + type Item = i32; + fn next(&mut self) -> Option { Some(21) } + fn nth(&mut self, n: usize) -> Option { + self.0 += n as Bigger + 1; + Some(42) + } + } + + let mut it = Test(0); + let root = usize::MAX >> (::std::mem::size_of::() * 8 / 2); + let n = root + 20; + (&mut it).step_by(n).nth(n); + assert_eq!(it.0, n as Bigger * n as Bigger); + + // large step + let mut it = Test(0); + (&mut it).step_by(usize::MAX).nth(5); + assert_eq!(it.0, (usize::MAX as Bigger) * 5); + + // n + 1 overflows + let mut it = Test(0); + (&mut it).step_by(2).nth(usize::MAX); + assert_eq!(it.0, (usize::MAX as Bigger) * 2); + + // n + 1 overflows + let mut it = Test(0); + (&mut it).step_by(1).nth(usize::MAX); + assert_eq!(it.0, (usize::MAX as Bigger) * 1); +} + #[test] #[should_panic] fn test_iterator_step_by_zero() { From f72b7f7c86b16fe8b76c4eb02052943dd6ef1ab2 Mon Sep 17 00:00:00 2001 From: oberien Date: Fri, 19 Jan 2018 22:34:22 +0100 Subject: [PATCH 5/6] Optimize StepBy::nth overflow handling --- src/libcore/iter/mod.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 57e7e03a6cebc..5923e6dc3ee28 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -307,6 +307,7 @@ use fmt; use iter_private::TrustedRandomAccess; use ops::Try; use usize; +use intrinsics; #[stable(feature = "rust1", since = "1.0.0")] pub use self::iterator::Iterator; @@ -718,7 +719,11 @@ impl Iterator for StepBy where I: Iterator { } // overflow handling - while n.checked_mul(step).is_none() { + loop { + let mul = n.checked_mul(step); + if unsafe { intrinsics::likely(mul.is_some())} { + return self.iter.nth(mul.unwrap() - 1); + } let div_n = usize::MAX / n; let div_step = usize::MAX / step; let nth_n = div_n * n; @@ -732,7 +737,6 @@ impl Iterator for StepBy where I: Iterator { }; self.iter.nth(nth - 1); } - self.iter.nth(n * step - 1) } } From 4a0da4cf2c7a2b5903fd1b8bc124f8963ce1b535 Mon Sep 17 00:00:00 2001 From: oberien Date: Sat, 20 Jan 2018 00:41:21 +0100 Subject: [PATCH 6/6] Spacing --- src/libcore/iter/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 5923e6dc3ee28..7314fac282b66 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -721,7 +721,7 @@ impl Iterator for StepBy where I: Iterator { // overflow handling loop { let mul = n.checked_mul(step); - if unsafe { intrinsics::likely(mul.is_some())} { + if unsafe { intrinsics::likely(mul.is_some()) } { return self.iter.nth(mul.unwrap() - 1); } let div_n = usize::MAX / n;