1 use crate::{intrinsics, iter::from_fn, ops::Try}; 2 3 /// An iterator for stepping iterators by a custom amount. 4 /// 5 /// This `struct` is created by the [`step_by`] method on [`Iterator`]. See 6 /// its documentation for more. 7 /// 8 /// [`step_by`]: Iterator::step_by 9 /// [`Iterator`]: trait.Iterator.html 10 #[must_use = "iterators are lazy and do nothing unless consumed"] 11 #[stable(feature = "iterator_step_by", since = "1.28.0")] 12 #[derive(Clone, Debug)] 13 pub struct StepBy<I> { 14 iter: I, 15 step: usize, 16 first_take: bool, 17 } 18 19 impl<I> StepBy<I> { new(iter: I, step: usize) -> StepBy<I>20 pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> { 21 assert!(step != 0); 22 StepBy { iter, step: step - 1, first_take: true } 23 } 24 } 25 26 #[stable(feature = "iterator_step_by", since = "1.28.0")] 27 impl<I> Iterator for StepBy<I> 28 where 29 I: Iterator, 30 { 31 type Item = I::Item; 32 33 #[inline] next(&mut self) -> Option<Self::Item>34 fn next(&mut self) -> Option<Self::Item> { 35 if self.first_take { 36 self.first_take = false; 37 self.iter.next() 38 } else { 39 self.iter.nth(self.step) 40 } 41 } 42 43 #[inline] size_hint(&self) -> (usize, Option<usize>)44 fn size_hint(&self) -> (usize, Option<usize>) { 45 #[inline] 46 fn first_size(step: usize) -> impl Fn(usize) -> usize { 47 move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) } 48 } 49 50 #[inline] 51 fn other_size(step: usize) -> impl Fn(usize) -> usize { 52 move |n| n / (step + 1) 53 } 54 55 let (low, high) = self.iter.size_hint(); 56 57 if self.first_take { 58 let f = first_size(self.step); 59 (f(low), high.map(f)) 60 } else { 61 let f = other_size(self.step); 62 (f(low), high.map(f)) 63 } 64 } 65 66 #[inline] nth(&mut self, mut n: usize) -> Option<Self::Item>67 fn nth(&mut self, mut n: usize) -> Option<Self::Item> { 68 if self.first_take { 69 self.first_take = false; 70 let first = self.iter.next(); 71 if n == 0 { 72 return first; 73 } 74 n -= 1; 75 } 76 // n and self.step are indices, we need to add 1 to get the amount of elements 77 // When calling `.nth`, we need to subtract 1 again to convert back to an index 78 // step + 1 can't overflow because `.step_by` sets `self.step` to `step - 1` 79 let mut step = self.step + 1; 80 // n + 1 could overflow 81 // thus, if n is usize::MAX, instead of adding one, we call .nth(step) 82 if n == usize::MAX { 83 self.iter.nth(step - 1); 84 } else { 85 n += 1; 86 } 87 88 // overflow handling 89 loop { 90 let mul = n.checked_mul(step); 91 { 92 if intrinsics::likely(mul.is_some()) { 93 return self.iter.nth(mul.unwrap() - 1); 94 } 95 } 96 let div_n = usize::MAX / n; 97 let div_step = usize::MAX / step; 98 let nth_n = div_n * n; 99 let nth_step = div_step * step; 100 let nth = if nth_n > nth_step { 101 step -= div_n; 102 nth_n 103 } else { 104 n -= div_step; 105 nth_step 106 }; 107 self.iter.nth(nth - 1); 108 } 109 } 110 try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R where F: FnMut(Acc, Self::Item) -> R, R: Try<Output = Acc>,111 fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R 112 where 113 F: FnMut(Acc, Self::Item) -> R, 114 R: Try<Output = Acc>, 115 { 116 #[inline] 117 fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ { 118 move || iter.nth(step) 119 } 120 121 if self.first_take { 122 self.first_take = false; 123 match self.iter.next() { 124 None => return try { acc }, 125 Some(x) => acc = f(acc, x)?, 126 } 127 } 128 from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f) 129 } 130 fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc where F: FnMut(Acc, Self::Item) -> Acc,131 fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc 132 where 133 F: FnMut(Acc, Self::Item) -> Acc, 134 { 135 #[inline] 136 fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ { 137 move || iter.nth(step) 138 } 139 140 if self.first_take { 141 self.first_take = false; 142 match self.iter.next() { 143 None => return acc, 144 Some(x) => acc = f(acc, x), 145 } 146 } 147 from_fn(nth(&mut self.iter, self.step)).fold(acc, f) 148 } 149 } 150 151 impl<I> StepBy<I> 152 where 153 I: ExactSizeIterator, 154 { 155 // The zero-based index starting from the end of the iterator of the 156 // last element. Used in the `DoubleEndedIterator` implementation. next_back_index(&self) -> usize157 fn next_back_index(&self) -> usize { 158 let rem = self.iter.len() % (self.step + 1); 159 if self.first_take { 160 if rem == 0 { self.step } else { rem - 1 } 161 } else { 162 rem 163 } 164 } 165 } 166 167 #[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")] 168 impl<I> DoubleEndedIterator for StepBy<I> 169 where 170 I: DoubleEndedIterator + ExactSizeIterator, 171 { 172 #[inline] next_back(&mut self) -> Option<Self::Item>173 fn next_back(&mut self) -> Option<Self::Item> { 174 self.iter.nth_back(self.next_back_index()) 175 } 176 177 #[inline] nth_back(&mut self, n: usize) -> Option<Self::Item>178 fn nth_back(&mut self, n: usize) -> Option<Self::Item> { 179 // `self.iter.nth_back(usize::MAX)` does the right thing here when `n` 180 // is out of bounds because the length of `self.iter` does not exceed 181 // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is 182 // zero-indexed 183 let n = n.saturating_mul(self.step + 1).saturating_add(self.next_back_index()); 184 self.iter.nth_back(n) 185 } 186 try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R where F: FnMut(Acc, Self::Item) -> R, R: Try<Output = Acc>,187 fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R 188 where 189 F: FnMut(Acc, Self::Item) -> R, 190 R: Try<Output = Acc>, 191 { 192 #[inline] 193 fn nth_back<I: DoubleEndedIterator>( 194 iter: &mut I, 195 step: usize, 196 ) -> impl FnMut() -> Option<I::Item> + '_ { 197 move || iter.nth_back(step) 198 } 199 200 match self.next_back() { 201 None => try { init }, 202 Some(x) => { 203 let acc = f(init, x)?; 204 from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f) 205 } 206 } 207 } 208 209 #[inline] rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc where Self: Sized, F: FnMut(Acc, Self::Item) -> Acc,210 fn rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc 211 where 212 Self: Sized, 213 F: FnMut(Acc, Self::Item) -> Acc, 214 { 215 #[inline] 216 fn nth_back<I: DoubleEndedIterator>( 217 iter: &mut I, 218 step: usize, 219 ) -> impl FnMut() -> Option<I::Item> + '_ { 220 move || iter.nth_back(step) 221 } 222 223 match self.next_back() { 224 None => init, 225 Some(x) => { 226 let acc = f(init, x); 227 from_fn(nth_back(&mut self.iter, self.step)).fold(acc, f) 228 } 229 } 230 } 231 } 232 233 // StepBy can only make the iterator shorter, so the len will still fit. 234 #[stable(feature = "iterator_step_by", since = "1.28.0")] 235 impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {} 236