step_count/
lib.rs

1#![no_std]
2#![feature(step_trait)]
3//! Iterator utility for counting the number of iterations with an arbitrary
4//! type that implements the [`Step`] trait.
5
6use core::iter::Step;
7
8/// Consumes the iterator, counting the number of iterations.
9/// This uses the [`Step`] trait to keep track of the count of iterations.
10/// The count starts from the default value provided by the [`Default`] trait.
11///
12/// _Note_: This function takes any implementation of [`IntoIterator`],
13/// which includes iterators themselves.
14///
15/// # Panics
16///
17/// Panics if the count exceeds the capacity of the count type (`T`).
18///
19/// # Examples
20///
21/// Basic usage:
22///
23/// ```
24/// # use step_count::step_count;
25/// let arr = [1, 2, 3];
26/// let count: usize = step_count(arr);
27/// assert_eq!(count, 3);
28/// ```
29///
30/// Overflow:
31///
32/// ```should_panic
33/// # use step_count::step_count;
34/// let arr = [(); u8::MAX as usize + 1];
35/// step_count::<u8>(arr);
36/// ```
37#[inline]
38pub fn step_count<T: Step + Default>(iter: impl IntoIterator) -> T {
39    step_count_from(iter, T::default())
40}
41
42/// Consumes the iterator, counting the number of iterations,
43/// starting from a given value.
44/// This uses the [`Step`] trait to keep track of the count of iterations.
45///
46/// _Note_: This function takes any implementation of [`IntoIterator`],
47/// which includes iterators themselves.
48///
49/// # Panics
50///
51/// Panics if the count exceeds the capacity of the count type (`T`).
52///
53/// # Examples
54///
55/// Basic usage:
56///
57/// ```
58/// # use step_count::step_count_from;
59/// let arr = [1, 2, 3];
60/// let count: u16 = step_count_from(arr, 2);
61/// assert_eq!(count, 5);
62/// ```
63///
64/// Overflow:
65///
66/// ```should_panic
67/// # use step_count::step_count_from;
68/// let arr = [(); u8::MAX as usize - 1];
69/// step_count_from::<u8>(arr, 2);
70/// ```
71#[inline]
72pub fn step_count_from<T: Step>(iter: impl IntoIterator, start: T) -> T {
73    let iter = iter.into_iter();
74    iter.fold(start, |count, _| T::forward(count, 1))
75}
76
77/// Consumes the iterator, counting the number of iterations.
78/// This uses the [`Step`] trait to keep track of the count of iterations.
79/// The count starts from the default value provided by the [`Default`] trait.
80/// Returns [`None`] if the count exceeded the capcity of the count type (`T`).
81///
82/// This function always fully consumes the iterator, even if [`None`] is returned.
83///
84/// _Note_: This function takes any implementation of [`IntoIterator`],
85/// which includes iterators themselves.
86///
87/// # Examples
88///
89/// Basic usage:
90///
91/// ```
92/// # use step_count::step_count_checked;
93/// let arr = [1, 2, 3];
94/// let count: Option<u16> = step_count_checked(arr);
95/// assert_eq!(count, Some(3));
96/// ```
97///
98/// Overflow:
99///
100/// ```
101/// # use step_count::step_count_checked;
102/// let arr = [(); u8::MAX as usize + 1];
103/// let count: Option<u8> = step_count_checked(arr);
104/// assert!(count.is_none());
105/// ```
106///
107/// Consumption:
108///
109/// ```
110/// # use step_count::step_count_checked;
111/// let mut range = -1..u8::MAX as isize;
112/// let count: Option<u8> = step_count_checked(&mut range);
113/// assert!(count.is_none());
114/// assert!(range.is_empty());
115/// ```
116#[inline]
117pub fn step_count_checked<T: Step + Default>(iter: impl IntoIterator) -> Option<T> {
118    step_count_from_checked(iter, T::default())
119}
120
121/// Consumes the iterator, counting the number of iterations,
122/// starting from a given value.
123/// This uses the [`Step`] trait to keep track of the count of iterations.
124/// Returns [`None`] if the count exceeded the capcity of the count type (`T`).
125///
126/// This function always fully consumes the iterator, even if [`None`] is returned.
127///
128/// _Note_: This function takes any implementation of [`IntoIterator`],
129/// which includes iterators themselves.
130///
131/// # Examples
132///
133/// Basic usage:
134///
135/// ```
136/// # use step_count::step_count_from_checked;
137/// let arr = [1, 2, 3];
138/// let count: Option<u16> = step_count_from_checked(arr, 2);
139/// assert_eq!(count, Some(5));
140/// ```
141///
142/// Overflow:
143///
144/// ```
145/// # use step_count::step_count_from_checked;
146/// let arr = [(); u8::MAX as usize - 1];
147/// let count: Option<u8> = step_count_from_checked(arr, 2);
148/// assert!(count.is_none());
149/// ```
150///
151/// Consumption:
152///
153/// ```
154/// # use step_count::step_count_from_checked;
155/// let mut range = -2..=i8::MAX as isize;
156/// let count: Option<i8> = step_count_from_checked(&mut range, -1);
157/// assert!(count.is_none());
158/// assert!(range.is_empty());
159/// ```
160#[inline]
161pub fn step_count_from_checked<T: Step>(iter: impl IntoIterator, start: T) -> Option<T> {
162    let mut iter = iter.into_iter();
163    let result = iter.try_fold(start, |count, _| T::forward_checked(count, 1));
164    iter.last();
165    result
166}
167
168/// Convenience trait to allow using `step_count*` functions as methods.
169/// This trait is implemented for every [`Iterator`].
170pub trait StepCount: Iterator {
171    /// Consumes the iterator, counting the number of iterations.
172    /// This uses the [`Step`] trait to keep track of the count of iterations.
173    /// The count starts from the default value provided by the [`Default`] trait.
174    ///
175    /// # Panics
176    ///
177    /// Panics if the count exceeds the capacity of the count type (`T`).
178    ///
179    /// # Examples
180    ///
181    /// Basic usage:
182    ///
183    /// ```
184    /// # use step_count::StepCount;
185    /// let arr = [1, 2, 3];
186    /// let count: usize = arr.into_iter().step_count();
187    /// assert_eq!(count, 3);
188    /// ```
189    ///
190    /// Overflow:
191    ///
192    /// ```should_panic
193    /// # use step_count::StepCount;
194    /// let range = 0..u8::MAX as usize + 1;
195    /// range.step_count::<u8>();
196    /// ```
197    #[inline]
198    fn step_count<T: Step + Default>(self) -> T
199    where
200        Self: Sized,
201    {
202        step_count(self)
203    }
204    /// Consumes the iterator, counting the number of iterations,
205    /// starting from a given value.
206    /// This uses the [`Step`] trait to keep track of the count of iterations.
207    ///
208    /// # Panics
209    ///
210    /// Panics if the count exceeds the capacity of the count type (`T`).
211    ///
212    /// # Examples
213    ///
214    /// Basic usage:
215    ///
216    /// ```
217    /// # use step_count::StepCount;
218    /// let arr = [1, 2, 3];
219    /// let count: u16 = arr.into_iter().step_count_from(2);
220    /// assert_eq!(count, 5);
221    /// ```
222    ///
223    /// Overflow:
224    ///
225    /// ```should_panic
226    /// # use step_count::StepCount;
227    /// let range = 0..u8::MAX as usize - 1;
228    /// range.step_count_from::<u8>(2);
229    /// ```
230    #[inline]
231    fn step_count_from<T: Step>(self, start: T) -> T
232    where
233        Self: Sized,
234    {
235        step_count_from(self, start)
236    }
237    /// Consumes the iterator, counting the number of iterations.
238    /// This uses the [`Step`] trait to keep track of the count of iterations.
239    /// The count starts from the default value provided by the [`Default`] trait.
240    /// Returns [`None`] if the count exceeded the capcity of the count type (`T`).
241    ///
242    /// This function always fully consumes the iterator, even if [`None`] is returned.
243    ///
244    /// # Examples
245    ///
246    /// Basic usage:
247    ///
248    /// ```
249    /// # use step_count::StepCount;
250    /// let arr = [1, 2, 3];
251    /// let count: Option<u16> = arr.into_iter().step_count_checked();
252    /// assert_eq!(count, Some(3));
253    /// ```
254    ///
255    /// Overflow:
256    ///
257    /// ```
258    /// # use step_count::StepCount;
259    /// let range = 0..u8::MAX as usize + 1;
260    /// let count: Option<u8> = range.step_count_checked();
261    /// assert!(count.is_none());
262    /// ```
263    ///
264    /// Consumption:
265    ///
266    /// ```
267    /// # use step_count::StepCount;
268    /// let mut range = -1..u8::MAX as isize;
269    /// let count: Option<u8> = range.by_ref().step_count_checked();
270    /// assert!(count.is_none());
271    /// assert!(range.is_empty());
272    /// ```
273    #[inline]
274    fn step_count_checked<T: Step + Default>(self) -> Option<T>
275    where
276        Self: Sized,
277    {
278        step_count_checked(self)
279    }
280    /// Consumes the iterator, counting the number of iterations,
281    /// starting from a given value.
282    /// This uses the [`Step`] trait to keep track of the count of iterations.
283    /// Returns [`None`] if the count exceeded the capcity of the count type (`T`).
284    ///
285    /// This function always fully consumes the iterator, even if [`None`] is returned.
286    ///
287    /// # Examples
288    ///
289    /// Basic usage:
290    ///
291    /// ```
292    /// # use step_count::StepCount;
293    /// let arr = [1, 2, 3];
294    /// let count: Option<u16> = arr.into_iter().step_count_from_checked(2);
295    /// assert_eq!(count, Some(5));
296    /// ```
297    ///
298    /// Overflow:
299    ///
300    /// ```
301    /// # use step_count::StepCount;
302    /// let range = 0..u8::MAX as usize - 1;
303    /// let count: Option<u8> = range.step_count_from_checked(2);
304    /// assert!(count.is_none());
305    /// ```
306    ///
307    /// Consumption:
308    ///
309    /// ```
310    /// # use step_count::StepCount;
311    /// let mut range = -2..i8::MAX as isize;
312    /// let count: Option<i8> = range.by_ref().step_count_from_checked(-1);
313    /// assert!(count.is_none());
314    /// assert!(range.is_empty());
315    /// ```
316    #[inline]
317    fn step_count_from_checked<T: Step>(self, start: T) -> Option<T>
318    where
319        Self: Sized,
320    {
321        step_count_from_checked(self, start)
322    }
323}
324
325impl<T: Iterator> StepCount for T {}