malachite_base/num/iterators/
mod.rs

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

malachite_base/num/iterators/
mod.rs
