malachite_base/num/iterators/
mod.rs

1// Copyright © 2025 Mikhail Hogrefe
2//
3// This file is part of Malachite.
4//
5// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
6// Lesser General Public License (LGPL) as published by the Free Software Foundation; either version
7// 3 of the License, or (at your option) any later version. See <https://www.gnu.org/licenses/>.
8
9use crate::iterators::bit_distributor::{BitDistributor, BitDistributorOutputType};
10use crate::num::arithmetic::traits::{DivMod, DivisibleBy};
11use crate::num::basic::unsigneds::PrimitiveUnsigned;
12use crate::num::conversion::traits::{ExactFrom, WrappingFrom};
13use core::cmp::Ordering::*;
14use core::marker::PhantomData;
15
16#[doc(hidden)]
17#[derive(Clone, Debug)]
18pub struct SameWidthIteratorToBitChunks<
19    I: Iterator<Item = T>,
20    T: PrimitiveUnsigned,
21    U: PrimitiveUnsigned,
22> {
23    xs: I,
24    width: u64,
25    phantom: PhantomData<*const U>,
26}
27
28impl<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned>
29    SameWidthIteratorToBitChunks<I, T, U>
30{
31    fn next_with_wrapping<F: Fn(T) -> U>(&mut self, wrap: F) -> Option<Option<U>> {
32        self.xs.next().map(|x| {
33            if x.significant_bits() <= self.width {
34                Some(wrap(x))
35            } else {
36                None
37            }
38        })
39    }
40}
41
42const fn same_width_iterator_to_bit_chunks<
43    I: Iterator<Item = T>,
44    T: PrimitiveUnsigned,
45    U: PrimitiveUnsigned,
46>(
47    xs: I,
48    width: u64,
49) -> SameWidthIteratorToBitChunks<I, T, U> {
50    SameWidthIteratorToBitChunks {
51        xs,
52        width,
53        phantom: PhantomData,
54    }
55}
56
57#[doc(hidden)]
58#[derive(Clone, Debug)]
59pub struct EvenFractionIteratorToBitChunks<
60    I: Iterator<Item = T>,
61    T: PrimitiveUnsigned,
62    U: PrimitiveUnsigned,
63> {
64    xs: I,
65    x: T,
66    multiple: u64,
67    x_width: u64,
68    y_width: u64,
69    counter: u64,
70    phantom: PhantomData<*const U>,
71}
72
73impl<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned>
74    EvenFractionIteratorToBitChunks<I, T, U>
75{
76    fn next_with_wrapping<F: Fn(T) -> U>(&mut self, wrap: F) -> Option<Option<U>> {
77        if self.counter == 0 {
78            if let Some(x) = self.xs.next() {
79                if x.significant_bits() > self.x_width {
80                    return Some(None);
81                }
82                self.x = x;
83                self.counter = self.multiple;
84            } else {
85                return None;
86            }
87        } else {
88            self.x >>= self.y_width;
89        }
90        let y = wrap(self.x.mod_power_of_2(self.y_width));
91        self.counter -= 1;
92        Some(Some(y))
93    }
94}
95
96const fn even_fraction_iterator_to_bit_chunks<
97    I: Iterator<Item = T>,
98    T: PrimitiveUnsigned,
99    U: PrimitiveUnsigned,
100>(
101    xs: I,
102    multiple: u64,
103    out_chunk_size: u64,
104) -> EvenFractionIteratorToBitChunks<I, T, U> {
105    EvenFractionIteratorToBitChunks {
106        xs,
107        x: T::ZERO,
108        multiple,
109        x_width: multiple * out_chunk_size,
110        y_width: out_chunk_size,
111        counter: 0,
112        phantom: PhantomData,
113    }
114}
115
116#[doc(hidden)]
117#[derive(Clone, Debug)]
118pub struct EvenMultipleIteratorToBitChunks<
119    I: Iterator<Item = T>,
120    T: PrimitiveUnsigned,
121    U: PrimitiveUnsigned,
122> {
123    xs: I,
124    x_width: u64,
125    y_width: u64,
126    done: bool,
127    phantom: PhantomData<*const U>,
128}
129
130impl<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned>
131    EvenMultipleIteratorToBitChunks<I, T, U>
132{
133    fn next_with_wrapping<F: Fn(T) -> U>(&mut self, wrap: F) -> Option<Option<U>> {
134        if self.done {
135            return None;
136        }
137        let mut y = U::ZERO;
138        let mut shift = 0;
139        while shift < self.y_width {
140            if let Some(x) = self.xs.next() {
141                if x.significant_bits() > self.x_width {
142                    return Some(None);
143                }
144                y |= wrap(x) << shift;
145                shift += self.x_width;
146            } else {
147                self.done = true;
148                break;
149            }
150        }
151        if shift == 0 { None } else { Some(Some(y)) }
152    }
153}
154
155const fn even_multiple_iterator_to_bit_chunks<
156    I: Iterator<Item = T>,
157    T: PrimitiveUnsigned,
158    U: PrimitiveUnsigned,
159>(
160    xs: I,
161    in_chunk_size: u64,
162    out_chunk_size: u64,
163) -> EvenMultipleIteratorToBitChunks<I, T, U> {
164    EvenMultipleIteratorToBitChunks {
165        xs,
166        x_width: in_chunk_size,
167        y_width: out_chunk_size,
168        done: false,
169        phantom: PhantomData,
170    }
171}
172
173#[doc(hidden)]
174#[derive(Clone, Debug)]
175pub struct IrregularIteratorToBitChunks<
176    I: Iterator<Item = T>,
177    T: PrimitiveUnsigned,
178    U: PrimitiveUnsigned,
179> {
180    xs: I,
181    x: T,
182    x_width: u64,
183    y_width: u64,
184    remaining_x_bits: u64,
185    in_inner_loop: bool,
186    phantom: PhantomData<*const U>,
187}
188
189impl<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned>
190    IrregularIteratorToBitChunks<I, T, U>
191{
192    fn next_with_wrapping<F: Fn(T) -> U>(&mut self, wrap: F) -> Option<Option<U>> {
193        let mut y = U::ZERO;
194        let mut remaining_y_bits = self.y_width;
195        loop {
196            if !self.in_inner_loop {
197                if let Some(x) = self.xs.next() {
198                    if x.significant_bits() > self.x_width {
199                        return Some(None);
200                    }
201                    self.x = x;
202                } else {
203                    break;
204                }
205                self.remaining_x_bits = self.x_width;
206                self.in_inner_loop = true;
207            }
208            while self.remaining_x_bits != 0 {
209                let y_index = self.y_width - remaining_y_bits;
210                if self.remaining_x_bits <= remaining_y_bits {
211                    y |= wrap(self.x) << y_index;
212                    remaining_y_bits -= self.remaining_x_bits;
213                    self.remaining_x_bits = 0;
214                } else {
215                    y |= wrap(self.x).mod_power_of_2(remaining_y_bits) << y_index;
216                    self.x >>= remaining_y_bits;
217                    self.remaining_x_bits -= remaining_y_bits;
218                    remaining_y_bits = 0;
219                }
220                if remaining_y_bits == 0 {
221                    return Some(Some(y));
222                }
223            }
224            self.in_inner_loop = false;
225        }
226        if y == U::ZERO { None } else { Some(Some(y)) }
227    }
228}
229
230const fn irregular_iterator_to_bit_chunks<
231    I: Iterator<Item = T>,
232    T: PrimitiveUnsigned,
233    U: PrimitiveUnsigned,
234>(
235    xs: I,
236    in_chunk_size: u64,
237    out_chunk_size: u64,
238) -> IrregularIteratorToBitChunks<I, T, U> {
239    IrregularIteratorToBitChunks {
240        xs,
241        x: T::ZERO,
242        x_width: in_chunk_size,
243        y_width: out_chunk_size,
244        remaining_x_bits: 0,
245        in_inner_loop: false,
246        phantom: PhantomData,
247    }
248}
249
250/// Regroups an iterator of bit chunks into another iterator of bit chunks, possibly with a
251/// different chunk size.
252///
253/// This `enum` is created by [`iterator_to_bit_chunks`]; see its documentation for more.
254#[derive(Clone, Debug)]
255pub enum IteratorToBitChunks<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned> {
256    SameWidth(SameWidthIteratorToBitChunks<I, T, U>),
257    EvenFraction(EvenFractionIteratorToBitChunks<I, T, U>),
258    EvenMultiple(EvenMultipleIteratorToBitChunks<I, T, U>),
259    Irregular(IrregularIteratorToBitChunks<I, T, U>),
260}
261
262impl<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned>
263    IteratorToBitChunks<I, T, U>
264{
265    pub(crate) fn next_with_wrapping<F: Fn(T) -> U>(&mut self, wrap: F) -> Option<Option<U>> {
266        match self {
267            IteratorToBitChunks::SameWidth(xs) => xs.next_with_wrapping(wrap),
268            IteratorToBitChunks::EvenFraction(xs) => xs.next_with_wrapping(wrap),
269            IteratorToBitChunks::EvenMultiple(xs) => xs.next_with_wrapping(wrap),
270            IteratorToBitChunks::Irregular(xs) => xs.next_with_wrapping(wrap),
271        }
272    }
273}
274
275impl<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned + WrappingFrom<T>> Iterator
276    for IteratorToBitChunks<I, T, U>
277{
278    type Item = Option<U>;
279
280    #[inline]
281    fn next(&mut self) -> Option<Option<U>> {
282        self.next_with_wrapping(U::wrapping_from)
283    }
284}
285
286/// Regroups an iterator of bit chunks into another iterator of bit chunks, possibly with a
287/// different chunk size.
288///
289/// In other words, let $A$ be the input chunk size and $B$ the output chunk size. Let's consider
290/// the lowest $A$ bits of each unsigned value produced by the input iterator, and concatenate them
291/// (least-significant bits first) into a single bit sequence. Then we chop the sequence up into
292/// chunks of $B$ bits, and have the output iterator return each chunk.
293///
294/// Let $(x\_i)\_{i=0}^{n-1}$ be the input iterator, where $n$ may be $\infty$. If $n$ is finite, we
295/// assume that $x\_{n-1} \neq 0$. Then we define the bit sequence $b\_{k=0}^\infty$ such that $b
296/// \in \\{0, 1\\}$, $b\_k=0$ for $k \geq An$, and
297/// $$
298/// x_i = \sum_{k=0}^{A-1} b_{Ai+k}2^k.
299/// $$
300/// Then, let $(y\_j)\_{j=0}^{m-1}$ be a sequence such that
301/// $$
302/// y_j = \sum_{k=0}^{B-1} b_{Bi+k}2^k.
303/// $$
304/// Then we have $f((x\_i)\_{i=0}^{n-1}) = (y\_j)\_{j=0}^{m-1}$. Note that the sequence $y$ is not
305/// uniquely specified, since it may contain arbitrarily many trailing zeros. However, if $x$ is
306/// finite, $y$ is guaranteed to also be finite.
307///
308/// The output length is $An/B + O(1)$, where $n$ is `xs.count()`, $A$ is `in_chunk_size`, and $B$
309/// is `out_chunk_size`.
310///
311/// # Complexity per iteration
312/// Constant time and additional memory.
313///
314/// # Examples
315/// ```
316/// use itertools::Itertools;
317/// use malachite_base::num::iterators::iterator_to_bit_chunks;
318///
319/// assert_eq!(
320///     iterator_to_bit_chunks::<_, u16, u32>([123, 456].iter().cloned(), 10, 10)
321///         .map(Option::unwrap)
322///         .collect_vec(),
323///     &[123, 456]
324/// );
325/// assert_eq!(
326///     iterator_to_bit_chunks::<_, u16, u16>([0b000111111, 0b110010010].iter().cloned(), 9, 3)
327///         .map(Option::unwrap)
328///         .collect_vec(),
329///     &[0b111, 0b111, 0b000, 0b010, 0b010, 0b110]
330/// );
331/// assert_eq!(
332///     iterator_to_bit_chunks::<_, u16, u32>(
333///         [0b111, 0b111, 0b000, 0b010, 0b010, 0b110].iter().cloned(),
334///         3,
335///         9
336///     )
337///     .map(Option::unwrap)
338///     .collect_vec(),
339///     &[0b000111111, 0b110010010]
340/// );
341/// assert_eq!(
342///     iterator_to_bit_chunks::<_, u32, u16>(
343///         [0b1010101, 0b1111101, 0b0100001, 0b110010].iter().cloned(),
344///         7,
345///         6
346///     )
347///     .map(Option::unwrap)
348///     .collect_vec(),
349///     &[0b010101, 0b111011, 0b000111, 0b010010, 0b110]
350/// );
351/// ```
352pub fn iterator_to_bit_chunks<I: Iterator<Item = T>, T: PrimitiveUnsigned, U: PrimitiveUnsigned>(
353    xs: I,
354    in_chunk_size: u64,
355    out_chunk_size: u64,
356) -> IteratorToBitChunks<I, T, U> {
357    assert_ne!(in_chunk_size, 0);
358    assert_ne!(out_chunk_size, 0);
359    assert!(in_chunk_size <= T::WIDTH);
360    assert!(out_chunk_size <= U::WIDTH);
361    match in_chunk_size.cmp(&out_chunk_size) {
362        Equal => {
363            return IteratorToBitChunks::SameWidth(same_width_iterator_to_bit_chunks(
364                xs,
365                in_chunk_size,
366            ));
367        }
368        Less => {
369            if out_chunk_size.divisible_by(in_chunk_size) {
370                return IteratorToBitChunks::EvenMultiple(even_multiple_iterator_to_bit_chunks(
371                    xs,
372                    in_chunk_size,
373                    out_chunk_size,
374                ));
375            }
376        }
377        Greater => {
378            let (multiple, remainder) = in_chunk_size.div_mod(out_chunk_size);
379            if remainder == 0 {
380                return IteratorToBitChunks::EvenFraction(even_fraction_iterator_to_bit_chunks(
381                    xs,
382                    multiple,
383                    out_chunk_size,
384                ));
385            }
386        }
387    }
388    IteratorToBitChunks::Irregular(irregular_iterator_to_bit_chunks(
389        xs,
390        in_chunk_size,
391        out_chunk_size,
392    ))
393}
394
395/// A `struct` that holds the state of the ruler sequence.
396///
397/// This `struct` is created by [`ruler_sequence`]; see its documentation for more.
398#[derive(Clone, Debug, Eq, Hash, PartialEq)]
399pub struct RulerSequence<T: ExactFrom<u32>> {
400    i: u64,
401    phantom: PhantomData<*const T>,
402}
403
404impl<T: ExactFrom<u32>> Iterator for RulerSequence<T> {
405    type Item = T;
406
407    fn next(&mut self) -> Option<T> {
408        self.i += 1;
409        Some(T::exact_from(self.i.trailing_zeros()))
410    }
411}
412
413/// Returns the ruler sequence.
414///
415/// The ruler sequence (<https://oeis.org/A007814>) is the number of times that 2 divides the
416/// numbers $1, 2, 3, \ldots$.
417///
418/// $(x_i)_{i=1}^\infty = t_i$, where for each $i$, $i = (2k_i+1)2^{t_i}$ for some $k_i\in
419/// \mathbb{Z}$.
420///
421/// The $n$th term of this sequence is no greater than $\log_2(n + 1)$. Every number occurs
422/// infinitely many times, and any number's first occurrence is after all smaller numbers have
423/// occured.
424///
425/// The output length is infinite.
426///
427/// # Complexity per iteration
428/// Constant time and additional memory.
429///
430/// # Examples
431/// ```
432/// use malachite_base::iterators::prefix_to_string;
433/// use malachite_base::num::iterators::ruler_sequence;
434///
435/// assert_eq!(
436///     prefix_to_string(ruler_sequence::<u32>(), 20),
437///     "[0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 4, 0, 1, 0, 2, ...]"
438/// );
439/// ```
440pub const fn ruler_sequence<T: ExactFrom<u32>>() -> RulerSequence<T> {
441    RulerSequence {
442        i: 0,
443        phantom: PhantomData,
444    }
445}
446
447/// A `struct` that holds the state of a bit distributor sequence.
448///
449/// This `struct` is created by [`bit_distributor_sequence`]; see its documentation for more.
450#[derive(Clone, Debug)]
451pub struct BitDistributorSequence {
452    bit_distributor: BitDistributor,
453}
454
455impl Iterator for BitDistributorSequence {
456    type Item = usize;
457
458    fn next(&mut self) -> Option<usize> {
459        let i = self.bit_distributor.get_output(1);
460        self.bit_distributor.increment_counter();
461        Some(i)
462    }
463}
464
465/// Returns a sequence based on a [`BitDistributor`].
466///
467/// The sequence is obtained by taking the second output of a two-output [`BitDistributor`]. If both
468/// output types are normal with weight 1, the sequence is <https://oeis.org/A059905>.
469///
470/// The smaller the first output type is relative to the second (where tiny outputs are smaller than
471/// normal outputs), the slower the sequence will grow.
472///
473/// - If the first output type is normal and the second is tiny, the sequence grows as $O(n)$.
474/// - If the first output type is tiny and the second is normal, the sequence grows as $O(\log n)$.
475/// - If both output types are normal with weights $p$ and $q$, the sequence grows as
476///   $O(n^\frac{p}{p+q})$.
477/// - The output types cannot both be tiny.
478///
479/// Every number occurs infinitely many times, and any number's first occurrence is after all
480/// smaller numbers have occured. The sequence increases by no more than 1 at each step, but may
481/// decrease by an arbitrarily large amount.
482///
483/// The output length is infinite.
484///
485/// # Complexity per iteration
486/// Constant time and additional memory.
487///
488/// # Panics
489/// Panics if both output types are tiny.
490///
491/// # Examples
492/// ```
493/// use malachite_base::iterators::bit_distributor::BitDistributorOutputType;
494/// use malachite_base::iterators::prefix_to_string;
495/// use malachite_base::num::iterators::bit_distributor_sequence;
496///
497/// assert_eq!(
498///     prefix_to_string(
499///         bit_distributor_sequence(
500///             BitDistributorOutputType::normal(1),
501///             BitDistributorOutputType::normal(2)
502///         ),
503///         50
504///     ),
505///     "[0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, \
506///     3, 2, 3, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 0, 1, ...]"
507/// );
508/// assert_eq!(
509///     prefix_to_string(
510///         bit_distributor_sequence(
511///             BitDistributorOutputType::normal(2),
512///             BitDistributorOutputType::normal(1)
513///         ),
514///         50
515///     ),
516///     "[0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7, 4, 5, 6, 7, 8, 9, 10, 11, 8, 9, 10, 11, 12, 13, 14, \
517///     15, 12, 13, 14, 15, 0, 1, 2, 3, 0, 1, 2, 3, 4, 5, 6, 7, 4, 5, 6, 7, 8, 9, ...]"
518/// );
519/// ```
520pub fn bit_distributor_sequence(
521    x_output_type: BitDistributorOutputType,
522    y_output_type: BitDistributorOutputType,
523) -> BitDistributorSequence {
524    BitDistributorSequence {
525        bit_distributor: BitDistributor::new(&[y_output_type, x_output_type]),
526    }
527}
malachite_base/num/iterators/mod.rs

malachite_base/num/iterators/
mod.rs
