toolbox_rs/
bit_weight_iterator.rs

1use crate::{enumerative_source_coding::decode_u64, math::choose};
2
3/// Implement an iterator of all 64 bit integers with fixed weight
4pub struct U64BitWeightIterator {
5    weight: u64,
6    ordinal: u64,
7    max: u64,
8}
9
10impl U64BitWeightIterator {
11    pub fn with_weight(weight: u64) -> Self {
12        U64BitWeightIterator {
13            weight,
14            ordinal: 0,
15            max: choose(64u64, weight),
16        }
17    }
18}
19
20impl Iterator for U64BitWeightIterator {
21    type Item = u64;
22
23    fn next(&mut self) -> Option<Self::Item> {
24        if self.ordinal < self.max {
25            self.ordinal += 1;
26            return Some(decode_u64(self.weight, self.ordinal - 1));
27        }
28        None
29    }
30}
31
32#[cfg(test)]
33mod tests {
34    use super::U64BitWeightIterator;
35
36    #[test]
37    fn trivial_iterator_of_weight_one() {
38        let result: Vec<u64> = U64BitWeightIterator::with_weight(1).collect();
39        assert_eq!(result.len(), 64);
40        let expected: [u64; 64] = core::array::from_fn(|i| 1 << i);
41        assert_eq!(result, expected);
42    }
43
44    #[test]
45    fn trivial_iterator_of_weight_63() {
46        let result: Vec<u64> = U64BitWeightIterator::with_weight(63).collect();
47        assert_eq!(result.len(), 64);
48        let expected: [u64; 64] =
49            core::array::from_fn(|i| 0xFFFF_FFFF_FFFF_FFFF ^ (1u64 << (63 - i)));
50        assert_eq!(result, expected);
51    }
52}
toolbox_rs/bit_weight_iterator.rs

toolbox_rs/
bit_weight_iterator.rs
