From 57bdc26db3e6e3619ea0a5b1e452b9cac7f5527d Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Wed, 13 Nov 2024 22:11:02 +0100
Subject: [PATCH] ff: add unit tests for BitIterator

---
 ff/src/bits.rs | 117 ++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 115 insertions(+), 2 deletions(-)

diff --git a/ff/src/bits.rs b/ff/src/bits.rs
index 049aaa37d..3bf4f3bf9 100644
--- a/ff/src/bits.rs
+++ b/ff/src/bits.rs
@@ -8,7 +8,7 @@ pub struct BitIteratorBE<Slice: AsRef<[u64]>> {
 impl<Slice: AsRef<[u64]>> BitIteratorBE<Slice> {
     pub fn new(s: Slice) -> Self {
         let n = s.as_ref().len() * 64;
-        BitIteratorBE { s, n }
+        Self { s, n }
     }
 
     /// Construct an iterator that automatically skips any leading zeros.
@@ -46,7 +46,7 @@ impl<Slice: AsRef<[u64]>> BitIteratorLE<Slice> {
     pub fn new(s: Slice) -> Self {
         let n = 0;
         let max_len = s.as_ref().len() * 64;
-        BitIteratorLE { s, n, max_len }
+        Self { s, n, max_len }
     }
 
     /// Construct an iterator that automatically skips any trailing zeros.
@@ -80,3 +80,116 @@ impl<Slice: AsRef<[u64]>> Iterator for BitIteratorLE<Slice> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ark_std::vec::Vec;
+
+    #[test]
+    fn test_bit_iterator_be() {
+        // Test with a simple case of a single 64-bit integer: 0b1010
+        let data = [0b1010u64];
+        let mut iter = BitIteratorBE::new(&data);
+
+        // The iterator should return the bits in big-endian order
+        // The first 60 bits are zeros
+        for _ in 0..60 {
+            assert_eq!(iter.next(), Some(false));
+        }
+        assert_eq!(iter.next(), Some(true)); // 3rd bit
+        assert_eq!(iter.next(), Some(false)); // 2nd bit
+        assert_eq!(iter.next(), Some(true)); // 1st bit
+        assert_eq!(iter.next(), Some(false)); // 0th bit
+        assert_eq!(iter.next(), None); // End of iteration
+
+        // Test with the without_leading_zeros method
+        let data = [0b0000_0000_0000_0000_0000_0000_0000_1010u64];
+        let iter: Vec<bool> = BitIteratorBE::without_leading_zeros(&data).collect();
+        assert_eq!(iter, vec![true, false, true, false]); // Only the significant bits
+
+        // Test with all zeros
+        let data = [0u64];
+        let iter: Vec<bool> = BitIteratorBE::without_leading_zeros(&data).collect();
+        assert!(iter.is_empty()); // Should be empty because all bits are zeros
+    }
+
+    #[test]
+    fn test_bit_iterator_le() {
+        // Test with a simple case of a single 64-bit integer: 0b1010
+        let data = [0b1010u64];
+        let mut iter = BitIteratorLE::new(&data);
+
+        // The iterator should return the bits in little-endian order
+        assert_eq!(iter.next(), Some(false)); // 0th bit
+        assert_eq!(iter.next(), Some(true)); // 1st bit
+        assert_eq!(iter.next(), Some(false)); // 2nd bit
+        assert_eq!(iter.next(), Some(true)); // 3rd bit
+        for _ in 4..64 {
+            assert_eq!(iter.next(), Some(false)); // The remaining bits are zeros
+        }
+        assert_eq!(iter.next(), None); // End of iteration
+
+        // Test with the without_trailing_zeros method
+        let data = [0b0000_0000_0000_0000_0000_0000_0000_1010u64];
+        let iter: Vec<bool> = BitIteratorLE::without_trailing_zeros(&data).collect();
+        assert_eq!(iter, vec![false, true, false, true]); // Only the significant bits
+
+        // Test with all zeros
+        let data = [0u64];
+        let iter: Vec<bool> = BitIteratorLE::without_trailing_zeros(&data).collect();
+        assert!(iter.is_empty()); // Should be empty because all bits are zeros
+    }
+
+    #[test]
+    fn test_bit_iterator_be_multiple_integers() {
+        // Test with multiple 64-bit integers: [0b1010, 0b1111]
+        let data = [0b1010u64, 0b1111u64];
+        let mut iter = BitIteratorBE::new(&data);
+
+        // First integer (0b1111) in big-endian order
+        for _ in 0..60 {
+            assert_eq!(iter.next(), Some(false));
+        }
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(true));
+
+        // Second integer (0b1010) in big-endian order
+        for _ in 0..60 {
+            assert_eq!(iter.next(), Some(false));
+        }
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(false));
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(false));
+        assert_eq!(iter.next(), None); // End of iteration
+    }
+
+    #[test]
+    fn test_bit_iterator_le_multiple_integers() {
+        // Test with multiple 64-bit integers: [0b1010, 0b1111]
+        let data = [0b1010u64, 0b1111u64];
+        let mut iter = BitIteratorLE::new(&data);
+
+        // First integer (0b1010) in little-endian order
+        assert_eq!(iter.next(), Some(false));
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(false));
+        assert_eq!(iter.next(), Some(true));
+        for _ in 4..64 {
+            assert_eq!(iter.next(), Some(false));
+        }
+
+        // Second integer (0b1111) in little-endian order
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(true));
+        assert_eq!(iter.next(), Some(true));
+        for _ in 4..64 {
+            assert_eq!(iter.next(), Some(false));
+        }
+        assert_eq!(iter.next(), None); // End of iteration
+    }
+}