Rollup merge of #123564 - scottmcm:step-by-div-zero, r=joboet

Don't emit divide-by-zero panic paths in `StepBy::len`

I happened to notice today that there's actually two such calls emitted in the assembly: <https://rust.godbolt.org/z/1Wbbd3Ts6>

Since they're impossible, hopefully telling LLVM that will also help optimizations elsewhere.

Author: matthiaskrgr

library/core/src/iter/adapters/step_by.rs

@@ -1,6 +1,7 @@
 use crate::{
     intrinsics,
     iter::{from_fn, TrustedLen, TrustedRandomAccess},
+    num::NonZeroUsize,
     ops::{Range, Try},
 };
 
@@ -22,7 +23,11 @@ pub struct StepBy<I> {
     /// Additionally this type-dependent preprocessing means specialized implementations
     /// cannot be used interchangeably.
     iter: I,
-    step: usize,
+    /// This field is `step - 1`, aka the correct amount to pass to `nth` when iterating.
+    /// It MUST NOT be `usize::MAX`, as `unsafe` code depends on being able to add one
+    /// without the risk of overflow.  (This is important so that length calculations
+    /// don't need to check for division-by-zero, for example.)
+    step_minus_one: usize,
     first_take: bool,
 }
 
@@ -31,7 +36,16 @@ impl<I> StepBy<I> {
     pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> {
         assert!(step != 0);
         let iter = <I as SpecRangeSetup<I>>::setup(iter, step);
-        StepBy { iter, step: step - 1, first_take: true }
+        StepBy { iter, step_minus_one: step - 1, first_take: true }
+    }
+
+    /// The `step` that was originally passed to `Iterator::step_by(step)`,
+    /// aka `self.step_minus_one + 1`.
+    #[inline]
+    fn original_step(&self) -> NonZeroUsize {
+        // SAFETY: By type invariant, `step_minus_one` cannot be `MAX`, which
+        // means the addition cannot overflow and the result cannot be zero.
+        unsafe { NonZeroUsize::new_unchecked(intrinsics::unchecked_add(self.step_minus_one, 1)) }
     }
 }
 
@@ -81,8 +95,8 @@ where
     // The zero-based index starting from the end of the iterator of the
     // last element. Used in the `DoubleEndedIterator` implementation.
     fn next_back_index(&self) -> usize {
-        let rem = self.iter.len() % (self.step + 1);
-        if self.first_take { if rem == 0 { self.step } else { rem - 1 } } else { rem }
+        let rem = self.iter.len() % self.original_step();
+        if self.first_take { if rem == 0 { self.step_minus_one } else { rem - 1 } } else { rem }
     }
 }
 
@@ -209,30 +223,30 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
 
     #[inline]
     default fn spec_next(&mut self) -> Option<I::Item> {
-        let step_size = if self.first_take { 0 } else { self.step };
+        let step_size = if self.first_take { 0 } else { self.step_minus_one };
         self.first_take = false;
         self.iter.nth(step_size)
     }
 
     #[inline]
     default fn spec_size_hint(&self) -> (usize, Option<usize>) {
         #[inline]
-        fn first_size(step: usize) -> impl Fn(usize) -> usize {
-            move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) }
+        fn first_size(step: NonZeroUsize) -> impl Fn(usize) -> usize {
+            move |n| if n == 0 { 0 } else { 1 + (n - 1) / step }
         }
 
         #[inline]
-        fn other_size(step: usize) -> impl Fn(usize) -> usize {
-            move |n| n / (step + 1)
+        fn other_size(step: NonZeroUsize) -> impl Fn(usize) -> usize {
+            move |n| n / step
         }
 
         let (low, high) = self.iter.size_hint();
 
         if self.first_take {
-            let f = first_size(self.step);
+            let f = first_size(self.original_step());
             (f(low), high.map(f))
         } else {
-            let f = other_size(self.step);
+            let f = other_size(self.original_step());
             (f(low), high.map(f))
         }
     }
@@ -247,10 +261,9 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
             }
             n -= 1;
         }
-        // n and self.step are indices, we need to add 1 to get the amount of elements
+        // n and self.step_minus_one 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;
+        let mut step = self.original_step().get();
         // n + 1 could overflow
         // thus, if n is usize::MAX, instead of adding one, we call .nth(step)
         if n == usize::MAX {
@@ -288,8 +301,11 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
         R: Try<Output = Acc>,
     {
         #[inline]
-        fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
-            move || iter.nth(step)
+        fn nth<I: Iterator>(
+            iter: &mut I,
+            step_minus_one: usize,
+        ) -> impl FnMut() -> Option<I::Item> + '_ {
+            move || iter.nth(step_minus_one)
         }
 
         if self.first_take {
@@ -299,16 +315,19 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
                 Some(x) => acc = f(acc, x)?,
             }
         }
-        from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f)
+        from_fn(nth(&mut self.iter, self.step_minus_one)).try_fold(acc, f)
     }
 
     default fn spec_fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
     where
         F: FnMut(Acc, Self::Item) -> Acc,
     {
         #[inline]
-        fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ {
-            move || iter.nth(step)
+        fn nth<I: Iterator>(
+            iter: &mut I,
+            step_minus_one: usize,
+        ) -> impl FnMut() -> Option<I::Item> + '_ {
+            move || iter.nth(step_minus_one)
         }
 
         if self.first_take {
@@ -318,7 +337,7 @@ unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
                 Some(x) => acc = f(acc, x),
             }
         }
-        from_fn(nth(&mut self.iter, self.step)).fold(acc, f)
+        from_fn(nth(&mut self.iter, self.step_minus_one)).fold(acc, f)
     }
 }
 
@@ -336,7 +355,7 @@ unsafe impl<I: DoubleEndedIterator + ExactSizeIterator> StepByBackImpl<I> for St
         // is out of bounds because the length of `self.iter` does not exceed
         // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is
         // zero-indexed
-        let n = n.saturating_mul(self.step + 1).saturating_add(self.next_back_index());
+        let n = n.saturating_mul(self.original_step().get()).saturating_add(self.next_back_index());
         self.iter.nth_back(n)
     }
 
@@ -348,16 +367,16 @@ unsafe impl<I: DoubleEndedIterator + ExactSizeIterator> StepByBackImpl<I> for St
         #[inline]
         fn nth_back<I: DoubleEndedIterator>(
             iter: &mut I,
-            step: usize,
+            step_minus_one: usize,
         ) -> impl FnMut() -> Option<I::Item> + '_ {
-            move || iter.nth_back(step)
+            move || iter.nth_back(step_minus_one)
         }
 
         match self.next_back() {
             None => try { init },
             Some(x) => {
                 let acc = f(init, x)?;
-                from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f)
+                from_fn(nth_back(&mut self.iter, self.step_minus_one)).try_fold(acc, f)
             }
         }
     }
@@ -371,16 +390,16 @@ unsafe impl<I: DoubleEndedIterator + ExactSizeIterator> StepByBackImpl<I> for St
         #[inline]
         fn nth_back<I: DoubleEndedIterator>(
             iter: &mut I,
-            step: usize,
+            step_minus_one: usize,
         ) -> impl FnMut() -> Option<I::Item> + '_ {
-            move || iter.nth_back(step)
+            move || iter.nth_back(step_minus_one)
         }
 
         match self.next_back() {
             None => init,
             Some(x) => {
                 let acc = f(init, x);
-                from_fn(nth_back(&mut self.iter, self.step)).fold(acc, f)
+                from_fn(nth_back(&mut self.iter, self.step_minus_one)).fold(acc, f)
             }
         }
     }
@@ -424,8 +443,7 @@ macro_rules! spec_int_ranges {
             fn spec_next(&mut self) -> Option<$t> {
                 // if a step size larger than the type has been specified fall back to
                 // t::MAX, in which case remaining will be at most 1.
-                // The `+ 1` can't overflow since the constructor substracted 1 from the original value.
-                let step = <$t>::try_from(self.step + 1).unwrap_or(<$t>::MAX);
+                let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX);
                 let remaining = self.iter.end;
                 if remaining > 0 {
                     let val = self.iter.start;
@@ -474,7 +492,7 @@ macro_rules! spec_int_ranges {
             {
                 // if a step size larger than the type has been specified fall back to
                 // t::MAX, in which case remaining will be at most 1.
-                let step = <$t>::try_from(self.step + 1).unwrap_or(<$t>::MAX);
+                let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX);
                 let remaining = self.iter.end;
                 let mut acc = init;
                 let mut val = self.iter.start;
@@ -500,7 +518,7 @@ macro_rules! spec_int_ranges_r {
             fn spec_next_back(&mut self) -> Option<Self::Item>
                 where Range<$t>: DoubleEndedIterator + ExactSizeIterator,
             {
-                let step = (self.step + 1) as $t;
+                let step = self.original_step().get() as $t;
                 let remaining = self.iter.end;
                 if remaining > 0 {
                     let start = self.iter.start;

tests/codegen/step_by-overflow-checks.rs

@@ -0,0 +1,26 @@
+//@ compile-flags: -O
+
+#![crate_type = "lib"]
+
+use std::iter::StepBy;
+use std::slice::Iter;
+
+// The constructor for `StepBy` ensures we can never end up needing to do zero
+// checks on denominators, so check that the code isn't emitting panic paths.
+
+// CHECK-LABEL: @step_by_len_std
+#[no_mangle]
+pub fn step_by_len_std(x: &StepBy<Iter<i32>>) -> usize {
+    // CHECK-NOT: div_by_zero
+    // CHECK: udiv
+    // CHECK-NOT: div_by_zero
+    x.len()
+}
+
+// CHECK-LABEL: @step_by_len_naive
+#[no_mangle]
+pub fn step_by_len_naive(x: Iter<i32>, step_minus_one: usize) -> usize {
+    // CHECK: udiv
+    // CHECK: call{{.+}}div_by_zero
+    x.len() / (step_minus_one + 1)
+}