snarkvm_utilities/
bititerator.rs

1// Copyright 2024 Aleo Network Foundation
2// This file is part of the snarkVM library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use std::iter::ExactSizeIterator;
17
18/// Iterates over a slice of `u64` in *big-endian* order.
19#[derive(Debug)]
20pub struct BitIteratorBE<Slice> {
21    s: Slice,
22    n: usize,
23}
24
25#[allow(clippy::len_without_is_empty)]
26impl<Slice: AsRef<[u64]>> BitIteratorBE<Slice> {
27    pub fn new(s: Slice) -> Self {
28        let n = s.as_ref().len() * 64;
29        BitIteratorBE { s, n }
30    }
31
32    /// Construct an iterator that automatically skips any leading zeros.
33    /// That is, it skips all zeros before the most-significant one.
34    pub fn new_without_leading_zeros(s: Slice) -> impl Iterator<Item = bool> {
35        Self::new(s).skip_while(|b| !b)
36    }
37}
38
39impl<Slice: AsRef<[u64]>> Iterator for BitIteratorBE<Slice> {
40    type Item = bool;
41
42    fn next(&mut self) -> Option<bool> {
43        if self.n == 0 {
44            None
45        } else {
46            self.n -= 1;
47            let part = self.n / 64;
48            let bit = self.n - (64 * part);
49
50            Some(self.s.as_ref()[part] & (1 << bit) > 0)
51        }
52    }
53}
54
55impl<Slice: AsRef<[u64]>> ExactSizeIterator for BitIteratorBE<Slice> {
56    fn len(&self) -> usize {
57        self.n
58    }
59}
60
61/// Iterates over a slice of `u64` in *little-endian* order.
62#[derive(Debug)]
63pub struct BitIteratorLE<Slice: AsRef<[u64]>> {
64    s: Slice,
65    n: usize,
66    max_len: usize,
67}
68
69impl<Slice: AsRef<[u64]>> BitIteratorLE<Slice> {
70    pub fn new(s: Slice) -> Self {
71        let n = 0;
72        let max_len = s.as_ref().len() * 64;
73        BitIteratorLE { s, n, max_len }
74    }
75}
76
77impl<Slice: AsRef<[u64]>> Iterator for BitIteratorLE<Slice> {
78    type Item = bool;
79
80    fn next(&mut self) -> Option<bool> {
81        if self.n == self.max_len {
82            None
83        } else {
84            let part = self.n / 64;
85            let bit = self.n - (64 * part);
86            self.n += 1;
87
88            Some(self.s.as_ref()[part] & (1 << bit) > 0)
89        }
90    }
91}
92
93impl<Slice: AsRef<[u64]>> ExactSizeIterator for BitIteratorLE<Slice> {
94    fn len(&self) -> usize {
95        self.max_len
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn test_bititerator_be() {
105        let mut five = BitIteratorBE::new(&[5]);
106
107        for _ in 0..61 {
108            assert_eq!(Some(false), five.next());
109        }
110        assert_eq!(Some(true), five.next());
111        assert_eq!(Some(false), five.next());
112        assert_eq!(Some(true), five.next());
113        assert_eq!(None, five.next());
114    }
115
116    #[test]
117    fn test_bititerator_be_without_leading_zeros() {
118        let mut five = BitIteratorBE::new_without_leading_zeros(&[5]);
119
120        assert_eq!(Some(true), five.next());
121        assert_eq!(Some(false), five.next());
122        assert_eq!(Some(true), five.next());
123        assert_eq!(None, five.next());
124    }
125
126    #[test]
127    fn test_bititerator_le() {
128        let mut five = BitIteratorLE::new(&[5]);
129
130        assert_eq!(Some(true), five.next());
131        assert_eq!(Some(false), five.next());
132        assert_eq!(Some(true), five.next());
133        for _ in 0..61 {
134            assert_eq!(Some(false), five.next());
135        }
136        assert_eq!(None, five.next());
137    }
138}