Skip to main content

kaspa_utils/
iter.rs

1use itertools::Itertools;
2
3/// Format all iterator elements, separated by `sep`.
4///
5/// Unlike the underlying `itertools::format`, **does not panic** if `fmt` is called more than once.
6/// Should be used for logging purposes since `itertools::format` will panic when used by multiple loggers.
7pub trait IterExtensions: Iterator {
8    fn reusable_format(self, sep: &str) -> ReusableIterFormat<'_, Self>
9    where
10        Self: Sized,
11    {
12        ReusableIterFormat::new(self.format(sep))
13    }
14
15    /// Provides a run-length-encoding iterator that yields the cumulative count
16    /// of elements seen so far, along with the value of the element. Useful for creating
17    /// compressed representations of sequences with repeating elements
18    fn rle_cumulative(self) -> impl Iterator<Item = (usize, Self::Item)>
19    where
20        Self: Sized,
21        Self::Item: PartialEq,
22    {
23        let mut cumulative: usize = 0;
24        self.dedup_with_count().map(move |(count, value)| {
25            cumulative += count;
26            (cumulative, value)
27        })
28    }
29}
30
31pub trait IterExtensionsRle<T>: Iterator<Item = (usize, T)>
32where
33    T: Clone,
34{
35    /// Expands a run-length encoded iterator back into its original sequence of elements.
36    /// It takes an iterator of (cumulative_count, item) tuples and yields the repeated items
37    fn expand_rle(self) -> impl Iterator<Item = T>
38    where
39        Self: Sized,
40    {
41        self.scan(0usize, |prev, (cum, item)| {
42            let count = cum.checked_sub(*prev).filter(|&c| c > 0).expect("cumulative counts must be strictly increasing");
43            *prev = cum;
44            Some((count, item))
45        })
46        .flat_map(|(count, item)| core::iter::repeat_n(item, count))
47    }
48}
49
50impl<I, T> IterExtensionsRle<T> for I
51where
52    I: Iterator<Item = (usize, T)>,
53    T: Clone,
54{
55}
56
57impl<T: ?Sized> IterExtensions for T where T: Iterator {}
58
59pub struct ReusableIterFormat<'a, I> {
60    inner: itertools::Format<'a, I>,
61}
62
63impl<'a, I> ReusableIterFormat<'a, I> {
64    pub fn new(inner: itertools::Format<'a, I>) -> Self {
65        Self { inner }
66    }
67}
68
69impl<I> core::fmt::Display for ReusableIterFormat<'_, I>
70where
71    I: core::clone::Clone,
72    I: Iterator,
73    I::Item: core::fmt::Display,
74{
75    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
76        // Clone the inner format to workaround the `Format: was already formatted once` internal error
77        self.inner.clone().fmt(f)
78    }
79}
80
81impl<I> core::fmt::Debug for ReusableIterFormat<'_, I>
82where
83    I: core::clone::Clone,
84    I: Iterator,
85    I::Item: core::fmt::Debug,
86{
87    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
88        // Clone the inner format to workaround the `Format: was already formatted once` internal error
89        self.inner.clone().fmt(f)
90    }
91}
92
93#[cfg(feature = "std")]
94/// Returns an iterator over powers of two up to (the rounded up) available parallelism: `2, 4, 8, ..., 2^(available_parallelism.log2().ceil())`,
95/// i.e., for `std::thread::available_parallelism = 15` the function will return `2, 4, 8, 16`
96pub fn parallelism_in_power_steps() -> impl Iterator<Item = usize> {
97    (1..=(std::thread::available_parallelism().unwrap().get() as f64).log2().ceil() as u32).map(|x| 2usize.pow(x))
98}