fp/prime/
iter.rs

1use itertools::Itertools;
2
3pub struct BitflagIterator {
4    remaining: u8,
5    flag: u64,
6}
7
8impl BitflagIterator {
9    pub fn new(flag: u64) -> Self {
10        Self {
11            remaining: u8::MAX,
12            flag,
13        }
14    }
15
16    pub fn new_fixed_length(flag: u64, remaining: usize) -> Self {
17        assert!(remaining <= 64);
18        let remaining = remaining as u8;
19        Self { remaining, flag }
20    }
21
22    pub fn set_bit_iterator(flag: u64) -> impl Iterator<Item = usize> {
23        Self::new(flag)
24            .enumerate()
25            .filter_map(|(idx, v)| if v { Some(idx) } else { None })
26    }
27}
28
29impl Iterator for BitflagIterator {
30    type Item = bool;
31
32    fn next(&mut self) -> Option<Self::Item> {
33        if self.remaining > 64 && self.flag == 0 || self.remaining == 0 {
34            None
35        } else {
36            self.remaining -= 1;
37            let result = self.flag & 1 == 1;
38            self.flag >>= 1;
39            Some(result)
40        }
41    }
42}
43
44/// Iterates through all combinations of numbers from 0 to `p - 1` of length `len`.
45///
46/// # Example
47/// ```
48/// # use fp::prime::{iter::combinations, ValidPrime};
49/// let mut iter = combinations(ValidPrime::new(3), 2);
50///
51/// assert_eq!(iter.next(), Some(vec![0, 0]));
52/// assert_eq!(iter.next(), Some(vec![0, 1]));
53/// assert_eq!(iter.next(), Some(vec![0, 2]));
54/// assert_eq!(iter.next(), Some(vec![1, 0]));
55/// assert_eq!(iter.next(), Some(vec![1, 1]));
56/// assert_eq!(iter.next(), Some(vec![1, 2]));
57/// assert_eq!(iter.next(), Some(vec![2, 0]));
58/// assert_eq!(iter.next(), Some(vec![2, 1]));
59/// assert_eq!(iter.next(), Some(vec![2, 2]));
60/// assert_eq!(iter.next(), None);
61/// ```
62pub fn combinations(p: impl Into<u32>, len: usize) -> impl Iterator<Item = Vec<u32>> {
63    let p = p.into();
64    (0..len).map(|_| 0..p).multi_cartesian_product()
65}
66
67/// Iterates through all numbers with the same number of bits. It may panic or return nonsense
68/// after all valid values are returned.
69pub struct BinomialIterator {
70    value: u32,
71}
72
73impl BinomialIterator {
74    pub fn new(n: usize) -> Self {
75        Self {
76            value: (1 << n) - 1,
77        }
78    }
79}
80
81impl Iterator for BinomialIterator {
82    type Item = u32;
83
84    fn next(&mut self) -> Option<Self::Item> {
85        let v = self.value;
86        let c = v & v.wrapping_neg();
87        let r = v + c;
88        let n = (r ^ v).wrapping_shr(2 + v.trailing_zeros());
89        self.value = r | n;
90        Some(v)
91    }
92}
fp/prime/iter.rs

fp/prime/
iter.rs
