1use core::ops::ControlFlow;
2
3use crate::{
4 try_trait_v2::Try, DoubleEndedLender, ExactSizeLender, FallibleLend, FallibleLender, FallibleLending, Lend, Lender,
5 Lending,
6};
7#[derive(Clone, Debug)]
8#[must_use = "lenders are lazy and do nothing unless consumed"]
9pub struct StepBy<L> {
10 lender: L,
11 step: usize,
12 first_take: bool,
13}
14impl<L> StepBy<L> {
15 pub(crate) fn new(lender: L, step: usize) -> Self {
16 assert!(step != 0);
17 StepBy { lender, step: step - 1, first_take: true }
18 }
19 pub fn into_inner(self) -> L {
20 self.lender
21 }
22 pub fn into_parts(self) -> (L, usize) {
23 (self.lender, self.step)
24 }
25}
26impl<'lend, L> Lending<'lend> for StepBy<L>
27where
28 L: Lender,
29{
30 type Lend = Lend<'lend, L>;
31}
32impl<L> Lender for StepBy<L>
33where
34 L: Lender,
35{
36 #[inline]
37 fn next(&mut self) -> Option<Lend<'_, Self>> {
38 if self.first_take {
39 self.first_take = false;
40 self.lender.next()
41 } else {
42 self.lender.nth(self.step)
43 }
44 }
45 #[inline]
46 fn size_hint(&self) -> (usize, Option<usize>) {
47 let (low, high) = self.lender.size_hint();
48 let step = self.step;
49 if self.first_take {
50 let f = move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) };
51 (f(low), high.map(f))
52 } else {
53 let f = move |n| n / (step + 1);
54 (f(low), high.map(f))
55 }
56 }
57 #[inline]
58 fn nth(&mut self, mut n: usize) -> Option<Lend<'_, Self>> {
59 if self.first_take {
60 self.first_take = false;
61 if n == 0 {
62 return self.lender.next();
63 }
64 n -= 1;
65 }
66 let mut step = self.step + 1;
67 if n == usize::MAX {
68 self.lender.nth(step - 1);
69 } else {
70 n += 1;
71 }
72 loop {
73 let mul = n.checked_mul(step);
74 if let Some(mul) = mul {
75 return self.lender.nth(mul - 1);
76 }
77 let div_n = usize::MAX / n;
78 let div_step = usize::MAX / step;
79 let nth_n = div_n * n;
80 let nth_step = div_step * step;
81 let nth = if nth_n > nth_step {
82 step -= div_n;
83 nth_n
84 } else {
85 n -= div_step;
86 nth_step
87 };
88 self.lender.nth(nth - 1);
89 }
90 }
91 fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
92 where
93 Self: Sized,
94 F: FnMut(B, Lend<'_, Self>) -> R,
95 R: Try<Output = B>,
96 {
97 let mut acc = init;
98 if self.first_take {
99 self.first_take = false;
100 match self.lender.next() {
101 None => return R::from_output(acc),
102 Some(x) => {
103 acc = match f(acc, x).branch() {
104 ControlFlow::Break(b) => return R::from_residual(b),
105 ControlFlow::Continue(c) => c,
106 }
107 }
108 }
109 }
110 while let Some(x) = self.lender.nth(self.step) {
111 acc = match f(acc, x).branch() {
112 ControlFlow::Break(b) => return R::from_residual(b),
113 ControlFlow::Continue(c) => c,
114 };
115 }
116 R::from_output(acc)
117 }
118 fn fold<B, F>(mut self, init: B, mut f: F) -> B
119 where
120 Self: Sized,
121 F: FnMut(B, Lend<'_, Self>) -> B,
122 {
123 let mut acc = init;
124 if self.first_take {
125 self.first_take = false;
126 match self.lender.next() {
127 None => return acc,
128 Some(x) => acc = f(acc, x),
129 }
130 }
131 while let Some(x) = self.lender.nth(self.step) {
132 acc = f(acc, x);
133 }
134 acc
135 }
136}
137impl<L> StepBy<L>
138where
139 L: ExactSizeLender,
140{
141 fn next_back_index(&self) -> usize {
142 let rem = self.lender.len() % (self.step + 1);
143 if self.first_take {
144 if rem == 0 {
145 self.step
146 } else {
147 rem - 1
148 }
149 } else {
150 rem
151 }
152 }
153}
154impl<L> DoubleEndedLender for StepBy<L>
155where
156 L: DoubleEndedLender + ExactSizeLender,
157{
158 #[inline]
159 fn next_back(&mut self) -> Option<Lend<'_, Self>> {
160 self.lender.nth_back(self.next_back_index())
161 }
162
163 #[inline]
164 fn nth_back(&mut self, n: usize) -> Option<Lend<'_, Self>> {
165 let n = n.saturating_mul(self.step + 1).saturating_add(self.next_back_index());
166 self.lender.nth_back(n)
167 }
168
169 fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
170 where
171 Self: Sized,
172 F: FnMut(B, Lend<'_, Self>) -> R,
173 R: Try<Output = B>,
174 {
175 let mut acc = init;
176 match self.next_back() {
177 None => return R::from_output(acc),
178 Some(x) => {
179 acc = match f(acc, x).branch() {
180 ControlFlow::Break(b) => return R::from_residual(b),
181 ControlFlow::Continue(c) => c,
182 };
183 }
184 }
185 while let Some(x) = self.lender.nth_back(self.step) {
186 acc = match f(acc, x).branch() {
187 ControlFlow::Break(b) => return R::from_residual(b),
188 ControlFlow::Continue(c) => c,
189 };
190 }
191 R::from_output(acc)
192 }
193
194 #[inline]
195 fn rfold<B, F>(mut self, init: B, mut f: F) -> B
196 where
197 Self: Sized,
198 F: FnMut(B, Lend<'_, Self>) -> B,
199 {
200 let mut acc = init;
201 match self.next_back() {
202 None => return acc,
203 Some(x) => acc = f(acc, x),
204 }
205 while let Some(x) = self.lender.nth_back(self.step) {
206 acc = f(acc, x);
207 }
208 acc
209 }
210}
211impl<L> ExactSizeLender for StepBy<L> where L: ExactSizeLender {}
212
213impl<'lend, L> FallibleLending<'lend> for StepBy<L>
214where
215 L: FallibleLender,
216{
217 type Lend = FallibleLend<'lend, L>;
218}
219impl<L> FallibleLender for StepBy<L>
220where
221 L: FallibleLender,
222{
223 type Error = L::Error;
224
225 #[inline]
226 fn next(&mut self) -> Result<Option<FallibleLend<'_, Self>>, Self::Error> {
227 if self.first_take {
228 self.first_take = false;
229 self.lender.next()
230 } else {
231 self.lender.nth(self.step)
232 }
233 }
234 #[inline]
235 fn size_hint(&self) -> (usize, Option<usize>) {
236 let (low, high) = self.lender.size_hint();
237 let step = self.step;
238 if self.first_take {
239 let f = move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) };
240 (f(low), high.map(f))
241 } else {
242 let f = move |n| n / (step + 1);
243 (f(low), high.map(f))
244 }
245 }
246 #[inline]
247 fn nth(&mut self, mut n: usize) -> Result<Option<FallibleLend<'_, Self>>, Self::Error> {
248 if self.first_take {
249 self.first_take = false;
250 if n == 0 {
251 return self.lender.next();
252 }
253 n -= 1;
254 }
255 let mut step = self.step + 1;
256 if n == usize::MAX {
257 self.lender.nth(step - 1)?;
258 } else {
259 n += 1;
260 }
261 loop {
262 let mul = n.checked_mul(step);
263 if let Some(mul) = mul {
264 return self.lender.nth(mul - 1);
265 }
266 let div_n = usize::MAX / n;
267 let div_step = usize::MAX / step;
268 let nth_n = div_n * n;
269 let nth_step = div_step * step;
270 let nth = if nth_n > nth_step {
271 step -= div_n;
272 nth_n
273 } else {
274 n -= div_step;
275 nth_step
276 };
277 self.lender.nth(nth - 1)?;
278 }
279 }
280 fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> Result<R, Self::Error>
281 where
282 Self: Sized,
283 F: FnMut(B, FallibleLend<'_, Self>) -> Result<R, Self::Error>,
284 R: Try<Output = B>,
285 {
286 let mut acc = init;
287 if self.first_take {
288 self.first_take = false;
289 match self.lender.next()? {
290 None => return Ok(R::from_output(acc)),
291 Some(x) => {
292 acc = match f(acc, x)?.branch() {
293 ControlFlow::Break(b) => return Ok(R::from_residual(b)),
294 ControlFlow::Continue(c) => c,
295 }
296 }
297 }
298 }
299 while let Some(x) = self.lender.nth(self.step)? {
300 acc = match f(acc, x)?.branch() {
301 ControlFlow::Break(b) => return Ok(R::from_residual(b)),
302 ControlFlow::Continue(c) => c,
303 };
304 }
305 Ok(R::from_output(acc))
306 }
307 fn fold<B, F>(mut self, init: B, mut f: F) -> Result<B, Self::Error>
308 where
309 Self: Sized,
310 F: FnMut(B, FallibleLend<'_, Self>) -> Result<B, Self::Error>,
311 {
312 let mut acc = init;
313 if self.first_take {
314 self.first_take = false;
315 match self.lender.next()? {
316 None => return Ok(acc),
317 Some(x) => acc = f(acc, x)?,
318 }
319 }
320 while let Some(x) = self.lender.nth(self.step)? {
321 acc = f(acc, x)?;
322 }
323 Ok(acc)
324 }
325}