Skip to main content

commonware_utils/bitmap/
mod.rs

1//! Bitmap implementation
2//!
3//! The bitmap is a compact representation of a sequence of bits, using chunks of bytes for a
4//! more-efficient memory layout than doing [`Vec<bool>`].
5
6#[cfg(not(feature = "std"))]
7use alloc::{collections::VecDeque, vec::Vec};
8use bytes::{Buf, BufMut};
9use commonware_codec::{util::at_least, EncodeSize, Error as CodecError, Read, ReadExt, Write};
10use core::{
11    fmt::{self, Formatter, Write as _},
12    iter,
13    ops::{BitAnd, BitOr, BitXor, Index, Range},
14};
15#[cfg(feature = "std")]
16use std::collections::VecDeque;
17
18mod prunable;
19pub use prunable::Prunable;
20
21pub mod historical;
22
23/// The default [BitMap] chunk size in bytes.
24pub const DEFAULT_CHUNK_SIZE: usize = 8;
25
26/// A bitmap that stores data in chunks of N bytes.
27///
28/// # Panics
29///
30/// Operations panic if `bit / CHUNK_SIZE_BITS > usize::MAX`. On 32-bit systems
31/// with N=32, this occurs at bit >= 1,099,511,627,776.
32#[derive(Clone, PartialEq, Eq, Hash)]
33pub struct BitMap<const N: usize = DEFAULT_CHUNK_SIZE> {
34    /// The bitmap itself, in chunks of size N bytes. Within each byte, lowest order bits are
35    /// treated as coming before higher order bits in the bit ordering.
36    ///
37    /// Invariant: `chunks.len() == len.div_ceil(CHUNK_SIZE_BITS)`
38    /// Invariant: All bits at index `i` where `i >= len` must be 0.
39    chunks: VecDeque<[u8; N]>,
40
41    /// The total number of bits stored in the bitmap.
42    len: u64,
43}
44
45impl<const N: usize> BitMap<N> {
46    const _CHUNK_SIZE_NON_ZERO_ASSERT: () = assert!(N > 0, "chunk size must be > 0");
47
48    /// The size of a chunk in bits.
49    pub const CHUNK_SIZE_BITS: u64 = (N * 8) as u64;
50
51    /// A chunk of all 0s.
52    const EMPTY_CHUNK: [u8; N] = [0u8; N];
53
54    /// A chunk of all 1s.
55    const FULL_CHUNK: [u8; N] = [u8::MAX; N];
56
57    /* Constructors */
58
59    /// Create a new empty bitmap.
60    pub const fn new() -> Self {
61        #[allow(path_statements)]
62        Self::_CHUNK_SIZE_NON_ZERO_ASSERT; // Prevent compilation for N == 0
63
64        Self {
65            chunks: VecDeque::new(),
66            len: 0,
67        }
68    }
69
70    // Create a new empty bitmap with the capacity to hold `size` bits without reallocating.
71    pub fn with_capacity(size: u64) -> Self {
72        #[allow(path_statements)]
73        Self::_CHUNK_SIZE_NON_ZERO_ASSERT; // Prevent compilation for N == 0
74
75        Self {
76            chunks: VecDeque::with_capacity(size.div_ceil(Self::CHUNK_SIZE_BITS) as usize),
77            len: 0,
78        }
79    }
80
81    /// Create a new bitmap with `size` bits, with all bits set to 0.
82    pub fn zeroes(size: u64) -> Self {
83        #[allow(path_statements)]
84        Self::_CHUNK_SIZE_NON_ZERO_ASSERT; // Prevent compilation for N == 0
85
86        let num_chunks = size.div_ceil(Self::CHUNK_SIZE_BITS) as usize;
87        let mut chunks = VecDeque::with_capacity(num_chunks);
88        for _ in 0..num_chunks {
89            chunks.push_back(Self::EMPTY_CHUNK);
90        }
91        Self { chunks, len: size }
92    }
93
94    /// Create a new bitmap with `size` bits, with all bits set to 1.
95    pub fn ones(size: u64) -> Self {
96        #[allow(path_statements)]
97        Self::_CHUNK_SIZE_NON_ZERO_ASSERT; // Prevent compilation for N == 0
98
99        let num_chunks = size.div_ceil(Self::CHUNK_SIZE_BITS) as usize;
100        let mut chunks = VecDeque::with_capacity(num_chunks);
101        for _ in 0..num_chunks {
102            chunks.push_back(Self::FULL_CHUNK);
103        }
104        let mut result = Self { chunks, len: size };
105        // Clear trailing bits to maintain invariant
106        result.clear_trailing_bits();
107        result
108    }
109
110    /* Length */
111
112    /// Return the number of bits currently stored in the bitmap.
113    #[inline]
114    pub const fn len(&self) -> u64 {
115        self.len
116    }
117
118    /// Returns true if the bitmap is empty.
119    #[inline]
120    pub const fn is_empty(&self) -> bool {
121        self.len() == 0
122    }
123
124    /// Returns true if the bitmap length is aligned to a chunk boundary.
125    #[inline]
126    pub const fn is_chunk_aligned(&self) -> bool {
127        self.len.is_multiple_of(Self::CHUNK_SIZE_BITS)
128    }
129
130    // Get the number of chunks currently in the bitmap.
131    fn chunks_len(&self) -> usize {
132        self.chunks.len()
133    }
134
135    /* Getters */
136
137    /// Get the value of the bit at the given index.
138    ///
139    /// # Warning
140    ///
141    /// Panics if the bit doesn't exist.
142    #[inline]
143    pub fn get(&self, bit: u64) -> bool {
144        let chunk = self.get_chunk_containing(bit);
145        Self::get_bit_from_chunk(chunk, bit)
146    }
147
148    /// Returns the bitmap chunk containing the given bit.
149    ///
150    /// # Warning
151    ///
152    /// Panics if the bit doesn't exist.
153    #[inline]
154    fn get_chunk_containing(&self, bit: u64) -> &[u8; N] {
155        assert!(
156            bit < self.len(),
157            "bit {} out of bounds (len: {})",
158            bit,
159            self.len()
160        );
161        &self.chunks[Self::to_chunk_index(bit)]
162    }
163
164    /// Get a reference to a chunk by its index in the current bitmap.
165    /// Note this is an index into the chunks, not a bit.
166    ///
167    /// # Warning
168    ///
169    /// Panics if the `chunk` is out of bounds.
170    #[inline]
171    pub(super) fn get_chunk(&self, chunk: usize) -> &[u8; N] {
172        assert!(
173            chunk < self.chunks.len(),
174            "chunk {} out of bounds (chunks: {})",
175            chunk,
176            self.chunks.len()
177        );
178        &self.chunks[chunk]
179    }
180
181    /// Get the value at the given `bit` from the `chunk`.
182    /// `bit` is an index into the entire bitmap, not just the chunk.
183    #[inline]
184    pub const fn get_bit_from_chunk(chunk: &[u8; N], bit: u64) -> bool {
185        let byte = Self::chunk_byte_offset(bit);
186        let byte = chunk[byte];
187        let mask = Self::chunk_byte_bitmask(bit);
188        (byte & mask) != 0
189    }
190
191    /// Return the last chunk of the bitmap and its size in bits.
192    ///
193    /// # Panics
194    ///
195    /// Panics if bitmap is empty.
196    #[inline]
197    fn last_chunk(&self) -> (&[u8; N], u64) {
198        let rem = self.len % Self::CHUNK_SIZE_BITS;
199        let bits_in_last_chunk = if rem == 0 { Self::CHUNK_SIZE_BITS } else { rem };
200        (self.chunks.back().unwrap(), bits_in_last_chunk)
201    }
202
203    /* Setters */
204
205    /// Extend the bitmap to `new_len` bits, filling new positions with zero.
206    /// No-op if `new_len <= self.len`.
207    pub fn extend_to(&mut self, new_len: u64) {
208        if new_len <= self.len {
209            return;
210        }
211        // Allocate any needed new chunks (all zeroed).
212        let new_chunks_needed = new_len.div_ceil(Self::CHUNK_SIZE_BITS) as usize;
213        let current_chunks = self.chunks.len();
214        for _ in current_chunks..new_chunks_needed {
215            self.chunks.push_back(Self::EMPTY_CHUNK);
216        }
217        self.len = new_len;
218    }
219
220    /// Add a single bit to the bitmap.
221    pub fn push(&mut self, bit: bool) {
222        // Check if we need a new chunk
223        if self.is_chunk_aligned() {
224            self.chunks.push_back(Self::EMPTY_CHUNK);
225        }
226
227        // Append to the last chunk
228        if bit {
229            let last_chunk = self.chunks.back_mut().unwrap();
230            let chunk_byte = Self::chunk_byte_offset(self.len);
231            last_chunk[chunk_byte] |= Self::chunk_byte_bitmask(self.len);
232        }
233        // If bit is false, just advance len -- the bit is already 0
234        self.len += 1;
235    }
236
237    /// Remove and return the last bit from the bitmap.
238    ///
239    /// # Warning
240    ///
241    /// Panics if the bitmap is empty.
242    pub fn pop(&mut self) -> bool {
243        assert!(!self.is_empty(), "Cannot pop from empty bitmap");
244
245        // Get the bit value at the last position
246        let last_bit_pos = self.len - 1;
247        let bit = Self::get_bit_from_chunk(self.chunks.back().unwrap(), last_bit_pos);
248
249        // Decrement length
250        self.len -= 1;
251
252        // Clear the bit we just popped to maintain invariant (if it was 1)
253        if bit {
254            let chunk_byte = Self::chunk_byte_offset(last_bit_pos);
255            let mask = Self::chunk_byte_bitmask(last_bit_pos);
256            self.chunks.back_mut().unwrap()[chunk_byte] &= !mask;
257        }
258
259        // Remove the last chunk if it's now empty
260        if self.is_chunk_aligned() {
261            self.chunks.pop_back();
262        }
263
264        bit
265    }
266
267    /// Shrink the bitmap to `new_len` bits, discarding trailing bits.
268    ///
269    /// # Panics
270    ///
271    /// Panics if `new_len > self.len()`.
272    pub fn truncate(&mut self, new_len: u64) {
273        assert!(new_len <= self.len(), "cannot truncate to a larger size");
274
275        // Pop single bits until we can remove full chunks.
276        while self.len > new_len && !self.is_chunk_aligned() {
277            self.pop();
278        }
279
280        // Pop full chunks from the back.
281        while self.len - new_len >= Self::CHUNK_SIZE_BITS {
282            self.pop_chunk();
283        }
284
285        // Pop remaining individual bits.
286        while self.len > new_len {
287            self.pop();
288        }
289    }
290
291    /// Remove and return the last complete chunk from the bitmap.
292    ///
293    /// # Warning
294    ///
295    /// Panics if the bitmap has fewer than `CHUNK_SIZE_BITS` bits or if not chunk-aligned.
296    pub(super) fn pop_chunk(&mut self) -> [u8; N] {
297        assert!(
298            self.len() >= Self::CHUNK_SIZE_BITS,
299            "cannot pop chunk: bitmap has fewer than CHUNK_SIZE_BITS bits"
300        );
301        assert!(
302            self.is_chunk_aligned(),
303            "cannot pop chunk when not chunk aligned"
304        );
305
306        // Remove and return the last data chunk
307        let chunk = self.chunks.pop_back().expect("chunk must exist");
308        self.len -= Self::CHUNK_SIZE_BITS;
309        chunk
310    }
311
312    /// Flips the given bit.
313    ///
314    /// # Panics
315    ///
316    /// Panics if `bit` is out of bounds.
317    #[inline]
318    pub fn flip(&mut self, bit: u64) {
319        self.assert_bit(bit);
320        let chunk = Self::to_chunk_index(bit);
321        let byte = Self::chunk_byte_offset(bit);
322        let mask = Self::chunk_byte_bitmask(bit);
323        self.chunks[chunk][byte] ^= mask;
324    }
325
326    /// Flips all bits (1s become 0s and vice versa).
327    pub fn flip_all(&mut self) {
328        for chunk in &mut self.chunks {
329            for byte in chunk {
330                *byte = !*byte;
331            }
332        }
333        // Clear trailing bits to maintain invariant
334        self.clear_trailing_bits();
335    }
336
337    /// Set the value of the referenced bit.
338    ///
339    /// # Warning
340    ///
341    /// Panics if the bit doesn't exist.
342    pub fn set(&mut self, bit: u64, value: bool) {
343        assert!(
344            bit < self.len(),
345            "bit {} out of bounds (len: {})",
346            bit,
347            self.len()
348        );
349
350        let chunk = &mut self.chunks[Self::to_chunk_index(bit)];
351        let byte = Self::chunk_byte_offset(bit);
352        let mask = Self::chunk_byte_bitmask(bit);
353        if value {
354            chunk[byte] |= mask;
355        } else {
356            chunk[byte] &= !mask;
357        }
358    }
359
360    /// Sets all bits to the specified value.
361    #[inline]
362    pub fn set_all(&mut self, bit: bool) {
363        let value = if bit { u8::MAX } else { 0 };
364        for chunk in &mut self.chunks {
365            chunk.fill(value);
366        }
367        // Clear trailing bits to maintain invariant
368        if bit {
369            self.clear_trailing_bits();
370        }
371    }
372
373    // Add a byte's worth of bits to the bitmap.
374    //
375    // # Warning
376    //
377    // Panics if self.len is not byte aligned.
378    fn push_byte(&mut self, byte: u8) {
379        assert!(
380            self.len.is_multiple_of(8),
381            "cannot add byte when not byte aligned"
382        );
383
384        // Check if we need a new chunk
385        if self.is_chunk_aligned() {
386            self.chunks.push_back(Self::EMPTY_CHUNK);
387        }
388
389        let chunk_byte = Self::chunk_byte_offset(self.len);
390        self.chunks.back_mut().unwrap()[chunk_byte] = byte;
391        self.len += 8;
392    }
393
394    /// Add a chunk of bits to the bitmap.
395    ///
396    /// # Warning
397    ///
398    /// Panics if self.len is not chunk aligned.
399    pub fn push_chunk(&mut self, chunk: &[u8; N]) {
400        assert!(
401            self.is_chunk_aligned(),
402            "cannot add chunk when not chunk aligned"
403        );
404        self.chunks.push_back(*chunk);
405        self.len += Self::CHUNK_SIZE_BITS;
406    }
407
408    /* Invariant Maintenance */
409
410    /// Clear all bits in the last chunk that are >= self.len to maintain the invariant.
411    /// Returns true if any bits were flipped from 1 to 0.
412    fn clear_trailing_bits(&mut self) -> bool {
413        if self.chunks.is_empty() {
414            return false;
415        }
416
417        let pos_in_chunk = self.len % Self::CHUNK_SIZE_BITS;
418        if pos_in_chunk == 0 {
419            // Chunk is full -- there are no trailing bits to clear.
420            return false;
421        }
422
423        let mut flipped_any = false;
424        let last_chunk = self.chunks.back_mut().unwrap();
425
426        // Clear whole bytes after the last valid bit
427        let last_byte_index = ((pos_in_chunk - 1) / 8) as usize;
428        for byte in last_chunk.iter_mut().skip(last_byte_index + 1) {
429            if *byte != 0 {
430                flipped_any = true;
431                *byte = 0;
432            }
433        }
434
435        // Clear the trailing bits in the last partial byte
436        let bits_in_last_byte = pos_in_chunk % 8;
437        if bits_in_last_byte != 0 {
438            let mask = (1u8 << bits_in_last_byte) - 1;
439            let old_byte = last_chunk[last_byte_index];
440            let new_byte = old_byte & mask;
441            if old_byte != new_byte {
442                flipped_any = true;
443                last_chunk[last_byte_index] = new_byte;
444            }
445        }
446
447        flipped_any
448    }
449
450    /* Pruning */
451
452    /// Remove the first `chunks` chunks from the bitmap.
453    ///
454    /// # Warning
455    ///
456    /// Panics if trying to prune more chunks than exist.
457    fn prune_chunks(&mut self, chunks: usize) {
458        assert!(
459            chunks <= self.chunks.len(),
460            "cannot prune {chunks} chunks, only {} available",
461            self.chunks.len()
462        );
463        self.chunks.drain(..chunks);
464        // Update len to reflect the removed chunks
465        let bits_removed = (chunks as u64) * Self::CHUNK_SIZE_BITS;
466        self.len = self.len.saturating_sub(bits_removed);
467    }
468
469    /// Prepend a chunk to the beginning of the bitmap.
470    pub(super) fn prepend_chunk(&mut self, chunk: &[u8; N]) {
471        self.chunks.push_front(*chunk);
472        self.len += Self::CHUNK_SIZE_BITS;
473    }
474
475    /// Overwrite a chunk's data at the given index.
476    ///
477    /// Replaces the entire chunk data, including any bits beyond `len()` in the last chunk.
478    /// The caller is responsible for ensuring `chunk_data` has the correct bit pattern
479    /// (e.g., zeros beyond the valid length if this is a partial last chunk).
480    ///
481    /// # Panics
482    ///
483    /// Panics if chunk_index is out of bounds.
484    pub(super) fn set_chunk_by_index(&mut self, chunk_index: usize, chunk_data: &[u8; N]) {
485        assert!(
486            chunk_index < self.chunks.len(),
487            "chunk index {chunk_index} out of bounds (chunks_len: {})",
488            self.chunks.len()
489        );
490        self.chunks[chunk_index].copy_from_slice(chunk_data);
491    }
492
493    /* Counting */
494
495    /// Returns the number of bits set to 1.
496    #[inline]
497    pub fn count_ones(&self) -> u64 {
498        // Thanks to the invariant that trailing bits are always 0,
499        // we can simply count all set bits in all chunks.
500        // Iterate over both contiguous deque segments and count 64-bit words first.
501        let (front, back) = self.chunks.as_slices();
502        Self::count_ones_in_chunk_slice(front) + Self::count_ones_in_chunk_slice(back)
503    }
504
505    #[inline]
506    fn count_ones_in_chunk_slice(chunks: &[[u8; N]]) -> u64 {
507        let mut total = 0u64;
508        let mut words = chunks.as_flattened().chunks_exact(8);
509        for word in &mut words {
510            total += u64::from_le_bytes(word.try_into().unwrap()).count_ones() as u64;
511        }
512        for byte in words.remainder() {
513            total += byte.count_ones() as u64;
514        }
515        total
516    }
517
518    /// Returns the number of bits set to 0.
519    #[inline]
520    pub fn count_zeros(&self) -> u64 {
521        self.len() - self.count_ones()
522    }
523
524    /* Indexing Helpers */
525
526    /// Convert a bit offset into a bitmask for the byte containing that bit.
527    #[inline]
528    pub(super) const fn chunk_byte_bitmask(bit: u64) -> u8 {
529        1 << (bit % 8)
530    }
531
532    /// Convert a bit into the index of the byte within a chunk containing the bit.
533    #[inline]
534    pub(super) const fn chunk_byte_offset(bit: u64) -> usize {
535        ((bit / 8) % N as u64) as usize
536    }
537
538    /// Convert a bit into the index of the chunk it belongs to.
539    ///
540    /// # Panics
541    ///
542    /// Panics if the chunk index overflows `usize`.
543    #[inline]
544    pub(super) fn to_chunk_index(bit: u64) -> usize {
545        let chunk = bit / Self::CHUNK_SIZE_BITS;
546        assert!(
547            chunk <= usize::MAX as u64,
548            "chunk overflow: {chunk} exceeds usize::MAX",
549        );
550        chunk as usize
551    }
552
553    /* Iterator */
554
555    /// Creates an iterator over the bits.
556    pub const fn iter(&self) -> Iterator<'_, N> {
557        Iterator {
558            bitmap: self,
559            pos: 0,
560        }
561    }
562
563    /// Returns an iterator over the indices of set bits.
564    pub fn ones_iter(&self) -> OnesIter<'_, Self, N> {
565        Readable::ones_iter_from(self, 0)
566    }
567
568    /* Bitwise Operations */
569
570    /// Helper for binary operations
571    #[inline]
572    fn binary_op<F: Fn(u8, u8) -> u8>(&mut self, other: &Self, op: F) {
573        self.assert_eq_len(other);
574        for (a_chunk, b_chunk) in self.chunks.iter_mut().zip(other.chunks.iter()) {
575            for (a_byte, b_byte) in a_chunk.iter_mut().zip(b_chunk.iter()) {
576                *a_byte = op(*a_byte, *b_byte);
577            }
578        }
579        // Clear trailing bits to maintain invariant
580        self.clear_trailing_bits();
581    }
582
583    /// Performs a bitwise AND with another BitMap.
584    ///
585    /// # Panics
586    ///
587    /// Panics if the lengths don't match.
588    pub fn and(&mut self, other: &Self) {
589        self.binary_op(other, |a, b| a & b);
590    }
591
592    /// Performs a bitwise OR with another BitMap.
593    ///
594    /// # Panics
595    ///
596    /// Panics if the lengths don't match.
597    pub fn or(&mut self, other: &Self) {
598        self.binary_op(other, |a, b| a | b);
599    }
600
601    /// Performs a bitwise XOR with another BitMap.
602    ///
603    /// # Panics
604    ///
605    /// Panics if the lengths don't match.
606    pub fn xor(&mut self, other: &Self) {
607        self.binary_op(other, |a, b| a ^ b);
608    }
609
610    /* Assertions */
611
612    /// Asserts that the bit is within bounds.
613    #[inline(always)]
614    fn assert_bit(&self, bit: u64) {
615        assert!(
616            bit < self.len(),
617            "bit {} out of bounds (len: {})",
618            bit,
619            self.len()
620        );
621    }
622
623    /// Asserts that the lengths of two [BitMap]s match.
624    #[inline(always)]
625    fn assert_eq_len(&self, other: &Self) {
626        assert_eq!(
627            self.len(),
628            other.len(),
629            "BitMap lengths don't match: {} vs {}",
630            self.len(),
631            other.len()
632        );
633    }
634
635    /// Check if all the bits in a given range are 0.
636    ///
637    /// Returns `true` if every index in the range is unset (i.e.
638    /// [`Self::get`] returns `false`). Returns `true` if the range
639    /// is empty.
640    ///
641    /// # Panics
642    ///
643    /// Panics if `range.end` exceeds the length of the bitmap.
644    ///
645    /// # Examples
646    ///
647    /// ```
648    /// use commonware_utils::bitmap::BitMap;
649    ///
650    /// let mut bitmap = BitMap::<8>::zeroes(128);
651    /// assert!(bitmap.is_unset(0..128));
652    ///
653    /// bitmap.set(64, true);
654    /// assert!(bitmap.is_unset(0..64));
655    /// assert!(!bitmap.is_unset(0..65));
656    /// ```
657    pub fn is_unset(&self, range: Range<u64>) -> bool {
658        assert!(
659            range.end <= self.len(),
660            "range end {} out of bounds (len: {})",
661            range.end,
662            self.len()
663        );
664        if range.start >= range.end {
665            return true;
666        }
667        let start = range.start;
668        let end = range.end;
669
670        // We know this can't underflow, because start < end.
671        //
672        // We now want "end" to represent the last bit we want to consider.
673        let end = end - 1;
674
675        // Get the chunks containing the start and end bits.
676        let first_chunk = Self::to_chunk_index(start);
677        let last_chunk = Self::to_chunk_index(end);
678
679        // All of these chunks require all of their bits to be checked.
680        // If first_chunk == last_chunk, we skip the loop.
681        let zero = [0u8; N];
682        for full_chunk in (first_chunk + 1)..last_chunk {
683            if self.chunks[full_chunk] != zero {
684                return false;
685            }
686        }
687
688        // Check first chunk tail (or whole range if first_chunk == last_chunk).
689        let start_byte = Self::chunk_byte_offset(start);
690        let end_byte = Self::chunk_byte_offset(end);
691        let start_mask = (0xFFu16 << ((start & 0b111) as u32)) as u8;
692        let end_mask = (0xFFu16 >> (7 - ((end & 0b111) as u32))) as u8;
693        let first = &self.chunks[first_chunk];
694        let first_end_byte = if first_chunk == last_chunk {
695            end_byte
696        } else {
697            N - 1
698        };
699        for (i, &byte) in first
700            .iter()
701            .enumerate()
702            .take(first_end_byte + 1)
703            .skip(start_byte)
704        {
705            let mut mask = 0xFFu8;
706            if i == start_byte {
707                mask &= start_mask;
708            }
709            if first_chunk == last_chunk && i == end_byte {
710                mask &= end_mask;
711            }
712            if (byte & mask) != 0 {
713                return false;
714            }
715        }
716        if first_chunk == last_chunk {
717            return true;
718        }
719
720        // Check last chunk head.
721        let last = &self.chunks[last_chunk];
722        for (i, &byte) in last.iter().enumerate().take(end_byte + 1) {
723            let mask = if i == end_byte { end_mask } else { 0xFF };
724            if (byte & mask) != 0 {
725                return false;
726            }
727        }
728
729        true
730    }
731}
732
733impl<const N: usize> Default for BitMap<N> {
734    fn default() -> Self {
735        Self::new()
736    }
737}
738
739impl<T: AsRef<[bool]>, const N: usize> From<T> for BitMap<N> {
740    fn from(t: T) -> Self {
741        let bools = t.as_ref();
742        let mut bv = Self::with_capacity(bools.len() as u64);
743        for &b in bools {
744            bv.push(b);
745        }
746        bv
747    }
748}
749
750impl<const N: usize> From<BitMap<N>> for Vec<bool> {
751    fn from(bv: BitMap<N>) -> Self {
752        bv.iter().collect()
753    }
754}
755
756impl<const N: usize> fmt::Debug for BitMap<N> {
757    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
758        // For very large BitMaps, only show a preview
759        const MAX_DISPLAY: u64 = 64;
760        const HALF_DISPLAY: u64 = MAX_DISPLAY / 2;
761
762        // Closure for writing a bit
763        let write_bit = |formatter: &mut Formatter<'_>, bit: u64| -> core::fmt::Result {
764            formatter.write_char(if self.get(bit) { '1' } else { '0' })
765        };
766
767        f.write_str("BitMap[")?;
768        let len = self.len();
769        if len <= MAX_DISPLAY {
770            // Show all bits
771            for i in 0..len {
772                write_bit(f, i)?;
773            }
774        } else {
775            // Show first and last bits with ellipsis
776            for i in 0..HALF_DISPLAY {
777                write_bit(f, i)?;
778            }
779
780            f.write_str("...")?;
781
782            for i in (len - HALF_DISPLAY)..len {
783                write_bit(f, i)?;
784            }
785        }
786        f.write_str("]")
787    }
788}
789
790impl<const N: usize> Index<u64> for BitMap<N> {
791    type Output = bool;
792
793    /// Allows accessing bits using the `[]` operator.
794    ///
795    /// Panics if out of bounds.
796    #[inline]
797    fn index(&self, bit: u64) -> &Self::Output {
798        self.assert_bit(bit);
799        let value = self.get(bit);
800        if value {
801            &true
802        } else {
803            &false
804        }
805    }
806}
807
808impl<const N: usize> BitAnd for &BitMap<N> {
809    type Output = BitMap<N>;
810
811    fn bitand(self, rhs: Self) -> Self::Output {
812        self.assert_eq_len(rhs);
813        let mut result = self.clone();
814        result.and(rhs);
815        result
816    }
817}
818
819impl<const N: usize> BitOr for &BitMap<N> {
820    type Output = BitMap<N>;
821
822    fn bitor(self, rhs: Self) -> Self::Output {
823        self.assert_eq_len(rhs);
824        let mut result = self.clone();
825        result.or(rhs);
826        result
827    }
828}
829
830impl<const N: usize> BitXor for &BitMap<N> {
831    type Output = BitMap<N>;
832
833    fn bitxor(self, rhs: Self) -> Self::Output {
834        self.assert_eq_len(rhs);
835        let mut result = self.clone();
836        result.xor(rhs);
837        result
838    }
839}
840
841impl<const N: usize> Write for BitMap<N> {
842    fn write(&self, buf: &mut impl BufMut) {
843        // Prefix with the number of bits
844        self.len().write(buf);
845
846        // Write all chunks
847        for chunk in &self.chunks {
848            for &byte in chunk {
849                byte.write(buf);
850            }
851        }
852    }
853}
854
855impl<const N: usize> Read for BitMap<N> {
856    type Cfg = u64; // Max bitmap length
857
858    fn read_cfg(buf: &mut impl Buf, max_len: &Self::Cfg) -> Result<Self, CodecError> {
859        // Parse length in bits
860        let len = u64::read(buf)?;
861        if len > *max_len {
862            return Err(CodecError::InvalidLength(len as usize));
863        }
864
865        // Calculate how many chunks we need to read
866        let num_chunks = len.div_ceil(Self::CHUNK_SIZE_BITS) as usize;
867
868        // Parse chunks
869        let mut chunks = VecDeque::with_capacity(num_chunks);
870        for _ in 0..num_chunks {
871            at_least(buf, N)?;
872            let mut chunk = [0u8; N];
873            buf.copy_to_slice(&mut chunk);
874            chunks.push_back(chunk);
875        }
876
877        let mut result = Self { chunks, len };
878
879        // Verify trailing bits are zero (maintain invariant)
880        if result.clear_trailing_bits() {
881            return Err(CodecError::Invalid(
882                "BitMap",
883                "Invalid trailing bits in encoded data",
884            ));
885        }
886
887        Ok(result)
888    }
889}
890
891impl<const N: usize> EncodeSize for BitMap<N> {
892    fn encode_size(&self) -> usize {
893        // Size of length prefix + all chunks
894        self.len().encode_size() + (self.chunks.len() * N)
895    }
896}
897
898/// Iterator over bits in a [BitMap].
899pub struct Iterator<'a, const N: usize> {
900    /// Reference to the BitMap being iterated over
901    bitmap: &'a BitMap<N>,
902
903    /// Current index in the BitMap
904    pos: u64,
905}
906
907impl<const N: usize> iter::Iterator for Iterator<'_, N> {
908    type Item = bool;
909
910    fn next(&mut self) -> Option<Self::Item> {
911        if self.pos >= self.bitmap.len() {
912            return None;
913        }
914
915        let bit = self.bitmap.get(self.pos);
916        self.pos += 1;
917        Some(bit)
918    }
919
920    fn size_hint(&self) -> (usize, Option<usize>) {
921        let remaining = self.bitmap.len().saturating_sub(self.pos);
922        let capped = remaining.min(usize::MAX as u64) as usize;
923        (capped, Some(capped))
924    }
925}
926
927impl<const N: usize> ExactSizeIterator for Iterator<'_, N> {}
928
929/// Read-only access to a bitmap's chunks and metadata.
930pub trait Readable<const N: usize> {
931    /// Return the number of complete (fully filled) chunks.
932    fn complete_chunks(&self) -> usize;
933
934    /// Return the chunk data at the given absolute chunk index.
935    fn get_chunk(&self, chunk: usize) -> [u8; N];
936
937    /// Return the last chunk and its size in bits.
938    fn last_chunk(&self) -> ([u8; N], u64);
939
940    /// Return the number of pruned chunks.
941    fn pruned_chunks(&self) -> usize;
942
943    /// Return the total number of bits.
944    fn len(&self) -> u64;
945
946    /// Returns true if the bitmap is empty.
947    fn is_empty(&self) -> bool {
948        self.len() == 0
949    }
950
951    /// Return the number of pruned bits (i.e. pruned chunks * bits per chunk).
952    fn pruned_bits(&self) -> u64 {
953        (self.pruned_chunks() as u64) * BitMap::<N>::CHUNK_SIZE_BITS
954    }
955
956    /// Return the value of a single bit.
957    fn get_bit(&self, bit: u64) -> bool {
958        let chunk = self.get_chunk(BitMap::<N>::to_chunk_index(bit));
959        BitMap::<N>::get_bit_from_chunk(&chunk, bit % BitMap::<N>::CHUNK_SIZE_BITS)
960    }
961
962    /// Returns an iterator over the indices of set bits starting from `pos`.
963    ///
964    /// If `pos` falls within a pruned region, iteration starts at the first
965    /// unpruned bit instead.
966    fn ones_iter_from(&self, pos: u64) -> OnesIter<'_, Self, N>
967    where
968        Self: Sized,
969    {
970        OnesIter { bitmap: self, pos }
971    }
972}
973
974impl<const N: usize> Readable<N> for BitMap<N> {
975    fn complete_chunks(&self) -> usize {
976        self.chunks_len()
977            .saturating_sub(if self.is_chunk_aligned() { 0 } else { 1 })
978    }
979
980    fn get_chunk(&self, chunk: usize) -> [u8; N] {
981        *Self::get_chunk(self, chunk)
982    }
983
984    fn last_chunk(&self) -> ([u8; N], u64) {
985        let (c, n) = Self::last_chunk(self);
986        (*c, n)
987    }
988
989    fn pruned_chunks(&self) -> usize {
990        0
991    }
992
993    fn len(&self) -> u64 {
994        self.len
995    }
996}
997
998/// Iterator over the indices of set (1) bits in a bitmap.
999///
1000/// If the starting position falls within a pruned region, iteration
1001/// begins at the first unpruned bit.
1002pub struct OnesIter<'a, B, const N: usize> {
1003    bitmap: &'a B,
1004    pos: u64,
1005}
1006
1007impl<B: Readable<N>, const N: usize> iter::Iterator for OnesIter<'_, B, N> {
1008    type Item = u64;
1009
1010    fn next(&mut self) -> Option<u64> {
1011        let len = self.bitmap.len();
1012        let chunk_bits = BitMap::<N>::CHUNK_SIZE_BITS;
1013        let pruned_start = self.bitmap.pruned_chunks() as u64 * chunk_bits;
1014        if self.pos < pruned_start {
1015            self.pos = pruned_start;
1016        }
1017        while self.pos < len {
1018            let chunk_idx = BitMap::<N>::to_chunk_index(self.pos);
1019            let chunk = self.bitmap.get_chunk(chunk_idx);
1020            let chunk_start = chunk_idx as u64 * chunk_bits;
1021            let rel = (self.pos - chunk_start) as usize;
1022            let mut byte_idx = rel / 8;
1023            let mut bit_in_byte = rel % 8;
1024            while byte_idx < N {
1025                let masked = chunk[byte_idx] >> bit_in_byte;
1026                if masked != 0 {
1027                    let found = chunk_start
1028                        + (byte_idx * 8 + bit_in_byte) as u64
1029                        + masked.trailing_zeros() as u64;
1030                    if found >= len {
1031                        return None;
1032                    }
1033                    self.pos = found + 1;
1034                    return Some(found);
1035                }
1036                byte_idx += 1;
1037                bit_in_byte = 0;
1038            }
1039            self.pos = chunk_start + chunk_bits;
1040        }
1041        None
1042    }
1043}
1044
1045#[cfg(feature = "arbitrary")]
1046impl<const N: usize> arbitrary::Arbitrary<'_> for BitMap<N> {
1047    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
1048        let size = u.int_in_range(0..=1024)?;
1049        let mut bits = Self::with_capacity(size);
1050        for _ in 0..size {
1051            bits.push(u.arbitrary::<bool>()?);
1052        }
1053        Ok(bits)
1054    }
1055}
1056
1057#[cfg(test)]
1058mod tests {
1059    use super::*;
1060    use crate::hex;
1061    use bytes::BytesMut;
1062    use commonware_codec::{Decode, Encode};
1063
1064    #[test]
1065    fn test_constructors() {
1066        // Test new()
1067        let bv: BitMap<4> = BitMap::new();
1068        assert_eq!(bv.len(), 0);
1069        assert!(bv.is_empty());
1070
1071        // Test default()
1072        let bv: BitMap<4> = Default::default();
1073        assert_eq!(bv.len(), 0);
1074        assert!(bv.is_empty());
1075
1076        // Test with_capacity()
1077        let bv: BitMap<4> = BitMap::with_capacity(0);
1078        assert_eq!(bv.len(), 0);
1079        assert!(bv.is_empty());
1080
1081        let bv: BitMap<4> = BitMap::with_capacity(10);
1082        assert_eq!(bv.len(), 0);
1083        assert!(bv.is_empty());
1084    }
1085
1086    #[test]
1087    fn test_zeroes() {
1088        let bv: BitMap<1> = BitMap::zeroes(0);
1089        assert_eq!(bv.len(), 0);
1090        assert!(bv.is_empty());
1091        assert_eq!(bv.count_ones(), 0);
1092        assert_eq!(bv.count_zeros(), 0);
1093
1094        let bv: BitMap<1> = BitMap::zeroes(1);
1095        assert_eq!(bv.len(), 1);
1096        assert!(!bv.is_empty());
1097        assert_eq!(bv.len(), 1);
1098        assert!(!bv.get(0));
1099        assert_eq!(bv.count_ones(), 0);
1100        assert_eq!(bv.count_zeros(), 1);
1101
1102        let bv: BitMap<1> = BitMap::zeroes(10);
1103        assert_eq!(bv.len(), 10);
1104        assert!(!bv.is_empty());
1105        assert_eq!(bv.len(), 10);
1106        for i in 0..10 {
1107            assert!(!bv.get(i as u64));
1108        }
1109        assert_eq!(bv.count_ones(), 0);
1110        assert_eq!(bv.count_zeros(), 10);
1111    }
1112
1113    #[test]
1114    fn test_ones() {
1115        let bv: BitMap<1> = BitMap::ones(0);
1116        assert_eq!(bv.len(), 0);
1117        assert!(bv.is_empty());
1118        assert_eq!(bv.count_ones(), 0);
1119        assert_eq!(bv.count_zeros(), 0);
1120
1121        let bv: BitMap<1> = BitMap::ones(1);
1122        assert_eq!(bv.len(), 1);
1123        assert!(!bv.is_empty());
1124        assert_eq!(bv.len(), 1);
1125        assert!(bv.get(0));
1126        assert_eq!(bv.count_ones(), 1);
1127        assert_eq!(bv.count_zeros(), 0);
1128
1129        let bv: BitMap<1> = BitMap::ones(10);
1130        assert_eq!(bv.len(), 10);
1131        assert!(!bv.is_empty());
1132        assert_eq!(bv.len(), 10);
1133        for i in 0..10 {
1134            assert!(bv.get(i as u64));
1135        }
1136        assert_eq!(bv.count_ones(), 10);
1137        assert_eq!(bv.count_zeros(), 0);
1138    }
1139
1140    #[test]
1141    fn test_invariant_trailing_bits_are_zero() {
1142        // Helper function to check the invariant
1143        fn check_trailing_bits_zero<const N: usize>(bitmap: &BitMap<N>) {
1144            let (last_chunk, next_bit) = bitmap.last_chunk();
1145
1146            // Check that all bits >= next_bit in the last chunk are 0
1147            for bit_idx in next_bit..((N * 8) as u64) {
1148                let byte_idx = (bit_idx / 8) as usize;
1149                let bit_in_byte = bit_idx % 8;
1150                let mask = 1u8 << bit_in_byte;
1151                assert_eq!(last_chunk[byte_idx] & mask, 0);
1152            }
1153        }
1154
1155        // Test ones() constructor
1156        let bv: BitMap<4> = BitMap::ones(15);
1157        check_trailing_bits_zero(&bv);
1158
1159        let bv: BitMap<4> = BitMap::ones(33);
1160        check_trailing_bits_zero(&bv);
1161
1162        // Test after push operations
1163        let mut bv: BitMap<4> = BitMap::new();
1164        for i in 0..37 {
1165            bv.push(i % 2 == 0);
1166            check_trailing_bits_zero(&bv);
1167        }
1168
1169        // Test after pop operations
1170        let mut bv: BitMap<4> = BitMap::ones(40);
1171        check_trailing_bits_zero(&bv);
1172        for _ in 0..15 {
1173            bv.pop();
1174            check_trailing_bits_zero(&bv);
1175        }
1176
1177        // Test after flip_all
1178        let mut bv: BitMap<4> = BitMap::ones(25);
1179        bv.flip_all();
1180        check_trailing_bits_zero(&bv);
1181
1182        // Test after binary operations
1183        let bv1: BitMap<4> = BitMap::ones(20);
1184        let bv2: BitMap<4> = BitMap::zeroes(20);
1185
1186        let mut bv_and = bv1.clone();
1187        bv_and.and(&bv2);
1188        check_trailing_bits_zero(&bv_and);
1189
1190        let mut bv_or = bv1.clone();
1191        bv_or.or(&bv2);
1192        check_trailing_bits_zero(&bv_or);
1193
1194        let mut bv_xor = bv1;
1195        bv_xor.xor(&bv2);
1196        check_trailing_bits_zero(&bv_xor);
1197
1198        // Test after deserialization
1199        let original: BitMap<4> = BitMap::ones(27);
1200        let encoded = original.encode();
1201        let decoded: BitMap<4> =
1202            BitMap::decode_cfg(&mut encoded.as_ref(), &(usize::MAX as u64)).unwrap();
1203        check_trailing_bits_zero(&decoded);
1204
1205        // Test clear_trailing_bits return value
1206        let mut bv_clean: BitMap<4> = BitMap::ones(20);
1207        // Should return false since ones() already clears trailing bits
1208        assert!(!bv_clean.clear_trailing_bits());
1209
1210        // Create a bitmap with invalid trailing bits by manually setting them
1211        let mut bv_dirty: BitMap<4> = BitMap::ones(20);
1212        // Manually corrupt the last chunk to have trailing bits set
1213        let last_chunk = bv_dirty.chunks.back_mut().unwrap();
1214        last_chunk[3] |= 0xF0; // Set some high bits in the last byte
1215                               // Should return true since we had invalid trailing bits
1216        assert!(bv_dirty.clear_trailing_bits());
1217        // After clearing, should return false
1218        assert!(!bv_dirty.clear_trailing_bits());
1219        check_trailing_bits_zero(&bv_dirty);
1220    }
1221
1222    #[test]
1223    fn test_get_set() {
1224        let mut bv: BitMap<4> = BitMap::new();
1225
1226        // Test initial state
1227        assert_eq!(bv.len(), 0);
1228        assert!(bv.is_empty());
1229
1230        // Test push
1231        bv.push(true);
1232        bv.push(false);
1233        bv.push(true);
1234        assert_eq!(bv.len(), 3);
1235        assert!(!bv.is_empty());
1236
1237        // Test get
1238        assert!(bv.get(0));
1239        assert!(!bv.get(1));
1240        assert!(bv.get(2));
1241
1242        bv.set(1, true);
1243        assert!(bv.get(1));
1244        bv.set(2, false);
1245        assert!(!bv.get(2));
1246
1247        // Test flip
1248        bv.flip(0); // true -> false
1249        assert!(!bv.get(0));
1250        bv.flip(0); // false -> true
1251        assert!(bv.get(0));
1252    }
1253
1254    #[test]
1255    fn test_chunk_operations() {
1256        let mut bv: BitMap<4> = BitMap::new();
1257        let test_chunk = hex!("0xABCDEF12");
1258
1259        // Test push_chunk
1260        bv.push_chunk(&test_chunk);
1261        assert_eq!(bv.len(), 32); // 4 bytes * 8 bits
1262
1263        // Test get_chunk
1264        let chunk = bv.get_chunk(0);
1265        assert_eq!(chunk, &test_chunk);
1266
1267        // Test get_chunk_containing
1268        let chunk = bv.get_chunk_containing(0);
1269        assert_eq!(chunk, &test_chunk);
1270
1271        // Test last_chunk
1272        let (last_chunk, next_bit) = bv.last_chunk();
1273        assert_eq!(next_bit, BitMap::<4>::CHUNK_SIZE_BITS); // Should be at chunk boundary
1274        assert_eq!(last_chunk, &test_chunk); // The chunk we just pushed
1275    }
1276
1277    #[test]
1278    fn test_pop() {
1279        let mut bv: BitMap<3> = BitMap::new();
1280        bv.push(true);
1281        assert!(bv.pop());
1282        assert_eq!(bv.len(), 0);
1283
1284        bv.push(false);
1285        assert!(!bv.pop());
1286        assert_eq!(bv.len(), 0);
1287
1288        bv.push(true);
1289        bv.push(false);
1290        bv.push(true);
1291        assert!(bv.pop());
1292        assert_eq!(bv.len(), 2);
1293        assert!(!bv.pop());
1294        assert_eq!(bv.len(), 1);
1295        assert!(bv.pop());
1296        assert_eq!(bv.len(), 0);
1297
1298        for i in 0..100 {
1299            bv.push(i % 2 == 0);
1300        }
1301        assert_eq!(bv.len(), 100);
1302        for i in (0..100).rev() {
1303            assert_eq!(bv.pop(), i % 2 == 0);
1304        }
1305        assert_eq!(bv.len(), 0);
1306        assert!(bv.is_empty());
1307    }
1308
1309    #[test]
1310    fn test_truncate() {
1311        let mut bv: BitMap<4> = BitMap::new();
1312        let expected: Vec<bool> = (0..70).map(|i| i % 3 == 0).collect();
1313        for &bit in &expected {
1314            bv.push(bit);
1315        }
1316
1317        bv.truncate(65);
1318        assert_eq!(bv.len(), 65);
1319        for i in 0..65 {
1320            assert_eq!(bv.get(i), expected[i as usize]);
1321        }
1322
1323        bv.truncate(32);
1324        assert_eq!(bv.len(), 32);
1325        for i in 0..32 {
1326            assert_eq!(bv.get(i), expected[i as usize]);
1327        }
1328
1329        bv.truncate(0);
1330        assert_eq!(bv.len(), 0);
1331        assert!(bv.is_empty());
1332    }
1333
1334    #[test]
1335    #[should_panic(expected = "cannot truncate to a larger size")]
1336    fn test_truncate_larger_size_panics() {
1337        let mut bv: BitMap<4> = BitMap::new();
1338        bv.push(true);
1339        bv.truncate(2);
1340    }
1341
1342    #[test]
1343    fn test_pop_chunk() {
1344        let mut bv: BitMap<3> = BitMap::new();
1345        const CHUNK_SIZE: u64 = BitMap::<3>::CHUNK_SIZE_BITS;
1346
1347        // Test 1: Pop a single chunk and verify it returns the correct data
1348        let chunk1 = hex!("0xAABBCC");
1349        bv.push_chunk(&chunk1);
1350        assert_eq!(bv.len(), CHUNK_SIZE);
1351        let popped = bv.pop_chunk();
1352        assert_eq!(popped, chunk1);
1353        assert_eq!(bv.len(), 0);
1354        assert!(bv.is_empty());
1355
1356        // Test 2: Pop multiple chunks in reverse order
1357        let chunk2 = hex!("0x112233");
1358        let chunk3 = hex!("0x445566");
1359        let chunk4 = hex!("0x778899");
1360
1361        bv.push_chunk(&chunk2);
1362        bv.push_chunk(&chunk3);
1363        bv.push_chunk(&chunk4);
1364        assert_eq!(bv.len(), CHUNK_SIZE * 3);
1365
1366        assert_eq!(bv.pop_chunk(), chunk4);
1367        assert_eq!(bv.len(), CHUNK_SIZE * 2);
1368
1369        assert_eq!(bv.pop_chunk(), chunk3);
1370        assert_eq!(bv.len(), CHUNK_SIZE);
1371
1372        assert_eq!(bv.pop_chunk(), chunk2);
1373        assert_eq!(bv.len(), 0);
1374
1375        // Test 3: Verify data integrity when popping chunks
1376        let first_chunk = hex!("0xAABBCC");
1377        let second_chunk = hex!("0x112233");
1378        bv.push_chunk(&first_chunk);
1379        bv.push_chunk(&second_chunk);
1380
1381        // Pop the second chunk, verify it and that first chunk is intact
1382        assert_eq!(bv.pop_chunk(), second_chunk);
1383        assert_eq!(bv.len(), CHUNK_SIZE);
1384
1385        for i in 0..CHUNK_SIZE {
1386            let byte_idx = (i / 8) as usize;
1387            let bit_idx = i % 8;
1388            let expected = (first_chunk[byte_idx] >> bit_idx) & 1 == 1;
1389            assert_eq!(bv.get(i), expected);
1390        }
1391
1392        assert_eq!(bv.pop_chunk(), first_chunk);
1393        assert_eq!(bv.len(), 0);
1394    }
1395
1396    #[test]
1397    #[should_panic(expected = "cannot pop chunk when not chunk aligned")]
1398    fn test_pop_chunk_not_aligned() {
1399        let mut bv: BitMap<3> = BitMap::new();
1400
1401        // Push a full chunk plus one bit
1402        bv.push_chunk(&[0xFF; 3]);
1403        bv.push(true);
1404
1405        // Should panic because not chunk-aligned
1406        bv.pop_chunk();
1407    }
1408
1409    #[test]
1410    #[should_panic(expected = "cannot pop chunk: bitmap has fewer than CHUNK_SIZE_BITS bits")]
1411    fn test_pop_chunk_insufficient_bits() {
1412        let mut bv: BitMap<3> = BitMap::new();
1413
1414        // Push only a few bits (less than a full chunk)
1415        bv.push(true);
1416        bv.push(false);
1417
1418        // Should panic because we don't have a full chunk to pop
1419        bv.pop_chunk();
1420    }
1421
1422    #[test]
1423    fn test_byte_operations() {
1424        let mut bv: BitMap<4> = BitMap::new();
1425
1426        // Test push_byte
1427        bv.push_byte(0xFF);
1428        assert_eq!(bv.len(), 8);
1429
1430        // All bits in the byte should be set
1431        for i in 0..8 {
1432            assert!(bv.get(i as u64));
1433        }
1434
1435        bv.push_byte(0x00);
1436        assert_eq!(bv.len(), 16);
1437
1438        // All bits in the second byte should be clear
1439        for i in 8..16 {
1440            assert!(!bv.get(i as u64));
1441        }
1442    }
1443
1444    #[test]
1445    fn test_count_operations() {
1446        let mut bv: BitMap<4> = BitMap::new();
1447
1448        // Empty bitmap
1449        assert_eq!(bv.count_ones(), 0);
1450        assert_eq!(bv.count_zeros(), 0);
1451
1452        // Add some bits
1453        bv.push(true);
1454        bv.push(false);
1455        bv.push(true);
1456        bv.push(true);
1457        bv.push(false);
1458
1459        assert_eq!(bv.count_ones(), 3);
1460        assert_eq!(bv.count_zeros(), 2);
1461        assert_eq!(bv.len(), 5);
1462
1463        // Test with full bytes
1464        let mut bv2: BitMap<4> = BitMap::new();
1465        bv2.push_byte(0xFF); // 8 ones
1466        bv2.push_byte(0x00); // 8 zeros
1467        bv2.push_byte(0xAA); // 4 ones, 4 zeros (10101010)
1468
1469        assert_eq!(bv2.count_ones(), 12);
1470        assert_eq!(bv2.count_zeros(), 12);
1471        assert_eq!(bv2.len(), 24);
1472    }
1473
1474    #[test]
1475    fn test_set_all() {
1476        let mut bv: BitMap<1> = BitMap::new();
1477
1478        // Add some bits
1479        bv.push(true);
1480        bv.push(false);
1481        bv.push(true);
1482        bv.push(false);
1483        bv.push(true);
1484        bv.push(false);
1485        bv.push(true);
1486        bv.push(false);
1487        bv.push(true);
1488        bv.push(false);
1489
1490        assert_eq!(bv.len(), 10);
1491        assert_eq!(bv.count_ones(), 5);
1492        assert_eq!(bv.count_zeros(), 5);
1493
1494        // Test set_all(true)
1495        bv.set_all(true);
1496        assert_eq!(bv.len(), 10);
1497        assert_eq!(bv.count_ones(), 10);
1498        assert_eq!(bv.count_zeros(), 0);
1499
1500        // Test set_all(false)
1501        bv.set_all(false);
1502        assert_eq!(bv.len(), 10);
1503        assert_eq!(bv.count_ones(), 0);
1504        assert_eq!(bv.count_zeros(), 10);
1505    }
1506
1507    #[test]
1508    fn test_flip_all() {
1509        let mut bv: BitMap<4> = BitMap::new();
1510
1511        bv.push(true);
1512        bv.push(false);
1513        bv.push(true);
1514        bv.push(false);
1515        bv.push(true);
1516
1517        let original_ones = bv.count_ones();
1518        let original_zeros = bv.count_zeros();
1519        let original_len = bv.len();
1520
1521        bv.flip_all();
1522
1523        // Length should not change
1524        assert_eq!(bv.len(), original_len);
1525
1526        // Ones and zeros should be swapped
1527        assert_eq!(bv.count_ones(), original_zeros);
1528        assert_eq!(bv.count_zeros(), original_ones);
1529
1530        // Check bits
1531        assert!(!bv.get(0));
1532        assert!(bv.get(1));
1533        assert!(!bv.get(2));
1534        assert!(bv.get(3));
1535        assert!(!bv.get(4));
1536    }
1537
1538    #[test]
1539    fn test_bitwise_and() {
1540        let mut bv1: BitMap<4> = BitMap::new();
1541        let mut bv2: BitMap<4> = BitMap::new();
1542
1543        // Create test patterns: 10110 & 11010 = 10010
1544        let pattern1 = [true, false, true, true, false];
1545        let pattern2 = [true, true, false, true, false];
1546        let expected = [true, false, false, true, false];
1547
1548        for &bit in &pattern1 {
1549            bv1.push(bit);
1550        }
1551        for &bit in &pattern2 {
1552            bv2.push(bit);
1553        }
1554
1555        bv1.and(&bv2);
1556
1557        assert_eq!(bv1.len(), 5);
1558        for (i, &expected_bit) in expected.iter().enumerate() {
1559            assert_eq!(bv1.get(i as u64), expected_bit);
1560        }
1561    }
1562
1563    #[test]
1564    fn test_bitwise_or() {
1565        let mut bv1: BitMap<4> = BitMap::new();
1566        let mut bv2: BitMap<4> = BitMap::new();
1567
1568        // Create test patterns: 10110 | 11010 = 11110
1569        let pattern1 = [true, false, true, true, false];
1570        let pattern2 = [true, true, false, true, false];
1571        let expected = [true, true, true, true, false];
1572
1573        for &bit in &pattern1 {
1574            bv1.push(bit);
1575        }
1576        for &bit in &pattern2 {
1577            bv2.push(bit);
1578        }
1579
1580        bv1.or(&bv2);
1581
1582        assert_eq!(bv1.len(), 5);
1583        for (i, &expected_bit) in expected.iter().enumerate() {
1584            assert_eq!(bv1.get(i as u64), expected_bit);
1585        }
1586    }
1587
1588    #[test]
1589    fn test_bitwise_xor() {
1590        let mut bv1: BitMap<4> = BitMap::new();
1591        let mut bv2: BitMap<4> = BitMap::new();
1592
1593        // Create test patterns: 10110 ^ 11010 = 01100
1594        let pattern1 = [true, false, true, true, false];
1595        let pattern2 = [true, true, false, true, false];
1596        let expected = [false, true, true, false, false];
1597
1598        for &bit in &pattern1 {
1599            bv1.push(bit);
1600        }
1601        for &bit in &pattern2 {
1602            bv2.push(bit);
1603        }
1604
1605        bv1.xor(&bv2);
1606
1607        assert_eq!(bv1.len(), 5);
1608        for (i, &expected_bit) in expected.iter().enumerate() {
1609            assert_eq!(bv1.get(i as u64), expected_bit);
1610        }
1611    }
1612
1613    #[test]
1614    fn test_multi_chunk_operations() {
1615        let mut bv1: BitMap<4> = BitMap::new();
1616        let mut bv2: BitMap<4> = BitMap::new();
1617
1618        // Fill multiple chunks
1619        let chunk1 = hex!("0xAABBCCDD"); // 10101010 10111011 11001100 11011101
1620        let chunk2 = hex!("0x55667788"); // 01010101 01100110 01110111 10001000
1621
1622        bv1.push_chunk(&chunk1);
1623        bv1.push_chunk(&chunk1);
1624        bv2.push_chunk(&chunk2);
1625        bv2.push_chunk(&chunk2);
1626
1627        assert_eq!(bv1.len(), 64);
1628        assert_eq!(bv2.len(), 64);
1629
1630        // Test AND operation
1631        let mut bv_and = bv1.clone();
1632        bv_and.and(&bv2);
1633
1634        // Test OR operation
1635        let mut bv_or = bv1.clone();
1636        bv_or.or(&bv2);
1637
1638        // Test XOR operation
1639        let mut bv_xor = bv1.clone();
1640        bv_xor.xor(&bv2);
1641
1642        // Verify results make sense
1643        assert_eq!(bv_and.len(), 64);
1644        assert_eq!(bv_or.len(), 64);
1645        assert_eq!(bv_xor.len(), 64);
1646
1647        // AND should have fewer or equal ones than either operand
1648        assert!(bv_and.count_ones() <= bv1.count_ones());
1649        assert!(bv_and.count_ones() <= bv2.count_ones());
1650
1651        // OR should have more or equal ones than either operand
1652        assert!(bv_or.count_ones() >= bv1.count_ones());
1653        assert!(bv_or.count_ones() >= bv2.count_ones());
1654    }
1655
1656    #[test]
1657    fn test_partial_chunk_operations() {
1658        let mut bv1: BitMap<4> = BitMap::new();
1659        let mut bv2: BitMap<4> = BitMap::new();
1660
1661        // Add partial chunks (not aligned to chunk boundaries)
1662        for i in 0..35 {
1663            // 35 bits = 4 bytes + 3 bits
1664            bv1.push(i % 2 == 0);
1665            bv2.push(i % 3 == 0);
1666        }
1667
1668        assert_eq!(bv1.len(), 35);
1669        assert_eq!(bv2.len(), 35);
1670
1671        // Test operations with partial chunks
1672        let mut bv_and = bv1.clone();
1673        bv_and.and(&bv2);
1674
1675        let mut bv_or = bv1.clone();
1676        bv_or.or(&bv2);
1677
1678        let mut bv_xor = bv1.clone();
1679        bv_xor.xor(&bv2);
1680
1681        // All should maintain the same length
1682        assert_eq!(bv_and.len(), 35);
1683        assert_eq!(bv_or.len(), 35);
1684        assert_eq!(bv_xor.len(), 35);
1685
1686        // Test flip_all with partial chunk
1687        let mut bv_inv = bv1.clone();
1688        let original_ones = bv_inv.count_ones();
1689        let original_zeros = bv_inv.count_zeros();
1690        bv_inv.flip_all();
1691        assert_eq!(bv_inv.count_ones(), original_zeros);
1692        assert_eq!(bv_inv.count_zeros(), original_ones);
1693    }
1694
1695    #[test]
1696    #[should_panic(expected = "bit 1 out of bounds (len: 1)")]
1697    fn test_flip_out_of_bounds() {
1698        let mut bv: BitMap<4> = BitMap::new();
1699        bv.push(true);
1700        bv.flip(1); // Only bit 0 exists
1701    }
1702
1703    #[test]
1704    #[should_panic(expected = "BitMap lengths don't match: 2 vs 1")]
1705    fn test_and_length_mismatch() {
1706        let mut bv1: BitMap<4> = BitMap::new();
1707        let mut bv2: BitMap<4> = BitMap::new();
1708
1709        bv1.push(true);
1710        bv1.push(false);
1711        bv2.push(true); // Different length
1712
1713        bv1.and(&bv2);
1714    }
1715
1716    #[test]
1717    #[should_panic(expected = "BitMap lengths don't match: 1 vs 2")]
1718    fn test_or_length_mismatch() {
1719        let mut bv1: BitMap<4> = BitMap::new();
1720        let mut bv2: BitMap<4> = BitMap::new();
1721
1722        bv1.push(true);
1723        bv2.push(true);
1724        bv2.push(false); // Different length
1725
1726        bv1.or(&bv2);
1727    }
1728
1729    #[test]
1730    #[should_panic(expected = "BitMap lengths don't match: 3 vs 2")]
1731    fn test_xor_length_mismatch() {
1732        let mut bv1: BitMap<4> = BitMap::new();
1733        let mut bv2: BitMap<4> = BitMap::new();
1734
1735        bv1.push(true);
1736        bv1.push(false);
1737        bv1.push(true);
1738        bv2.push(true);
1739        bv2.push(false); // Different length
1740
1741        bv1.xor(&bv2);
1742    }
1743
1744    #[test]
1745    fn test_equality() {
1746        // Test empty bitmaps
1747        assert_eq!(BitMap::<4>::new(), BitMap::<4>::new());
1748        assert_eq!(BitMap::<8>::new(), BitMap::<8>::new());
1749
1750        // Test non-empty bitmaps from constructors
1751        let pattern = [true, false, true, true, false, false, true, false, true];
1752        let bv4: BitMap<4> = pattern.as_ref().into();
1753        assert_eq!(bv4, BitMap::<4>::from(pattern.as_ref()));
1754        let bv8: BitMap<8> = pattern.as_ref().into();
1755        assert_eq!(bv8, BitMap::<8>::from(pattern.as_ref()));
1756
1757        // Test non-empty bitmaps from push operations
1758        let mut bv1: BitMap<4> = BitMap::new();
1759        let mut bv2: BitMap<4> = BitMap::new();
1760        for i in 0..33 {
1761            let bit = i % 3 == 0;
1762            bv1.push(bit);
1763            bv2.push(bit);
1764        }
1765        assert_eq!(bv1, bv2);
1766
1767        // Test inequality: different lengths
1768        bv1.push(true);
1769        assert_ne!(bv1, bv2);
1770        bv1.pop(); // Restore equality
1771        assert_eq!(bv1, bv2);
1772
1773        // Test inequality: different content
1774        bv1.flip(15);
1775        assert_ne!(bv1, bv2);
1776        bv1.flip(15); // Restore equality
1777        assert_eq!(bv1, bv2);
1778
1779        // Test equality after operations
1780        let mut bv_ops1 = BitMap::<16>::ones(25);
1781        let mut bv_ops2 = BitMap::<16>::ones(25);
1782        bv_ops1.flip_all();
1783        bv_ops2.flip_all();
1784        assert_eq!(bv_ops1, bv_ops2);
1785
1786        let mask_bits: Vec<bool> = (0..33).map(|i| i % 3 == 0).collect();
1787        let mask = BitMap::<4>::from(mask_bits);
1788        bv1.and(&mask);
1789        bv2.and(&mask);
1790        assert_eq!(bv1, bv2);
1791    }
1792
1793    #[test]
1794    fn test_different_chunk_sizes() {
1795        // Test with different chunk sizes
1796        let mut bv8: BitMap<8> = BitMap::new();
1797        let mut bv16: BitMap<16> = BitMap::new();
1798        let mut bv32: BitMap<32> = BitMap::new();
1799
1800        // Test chunk operations first (must be chunk-aligned)
1801        let chunk8 = [0xFF; 8];
1802        let chunk16 = [0xAA; 16];
1803        let chunk32 = [0x55; 32];
1804
1805        bv8.push_chunk(&chunk8);
1806        bv16.push_chunk(&chunk16);
1807        bv32.push_chunk(&chunk32);
1808
1809        // Test basic operations work with different sizes
1810        bv8.push(true);
1811        bv8.push(false);
1812        assert_eq!(bv8.len(), 64 + 2);
1813        assert_eq!(bv8.count_ones(), 64 + 1); // chunk8 is all 0xFF + 1 true bit
1814        assert_eq!(bv8.count_zeros(), 1);
1815
1816        bv16.push(true);
1817        bv16.push(false);
1818        assert_eq!(bv16.len(), 128 + 2);
1819        assert_eq!(bv16.count_ones(), 64 + 1); // chunk16 is 0xAA pattern + 1 true bit
1820        assert_eq!(bv16.count_zeros(), 64 + 1);
1821
1822        bv32.push(true);
1823        bv32.push(false);
1824        assert_eq!(bv32.len(), 256 + 2);
1825        assert_eq!(bv32.count_ones(), 128 + 1); // chunk32 is 0x55 pattern + 1 true bit
1826        assert_eq!(bv32.count_zeros(), 128 + 1);
1827    }
1828
1829    #[test]
1830    fn test_iterator() {
1831        // Test empty iterator
1832        let bv: BitMap<4> = BitMap::new();
1833        let mut iter = bv.iter();
1834        assert_eq!(iter.next(), None);
1835        assert_eq!(iter.size_hint(), (0, Some(0)));
1836
1837        // Test iterator with some bits
1838        let pattern = [true, false, true, false, true];
1839        let bv: BitMap<4> = pattern.as_ref().into();
1840
1841        // Collect all bits via iterator
1842        let collected: Vec<bool> = bv.iter().collect();
1843        assert_eq!(collected, pattern);
1844
1845        // Test size_hint
1846        let mut iter = bv.iter();
1847        assert_eq!(iter.size_hint(), (5, Some(5)));
1848
1849        // Consume one element and check size_hint again
1850        assert_eq!(iter.next(), Some(true));
1851        assert_eq!(iter.size_hint(), (4, Some(4)));
1852
1853        // Test ExactSizeIterator
1854        let iter = bv.iter();
1855        assert_eq!(iter.len(), 5);
1856
1857        // Test iterator with larger bitmap
1858        let mut large_bv: BitMap<8> = BitMap::new();
1859        for i in 0..100 {
1860            large_bv.push(i % 3 == 0);
1861        }
1862
1863        let collected: Vec<bool> = large_bv.iter().collect();
1864        assert_eq!(collected.len(), 100);
1865        for (i, &bit) in collected.iter().enumerate() {
1866            assert_eq!(bit, i % 3 == 0);
1867        }
1868    }
1869
1870    #[test]
1871    fn test_iterator_edge_cases() {
1872        // Test iterator with single bit
1873        let mut bv: BitMap<4> = BitMap::new();
1874        bv.push(true);
1875
1876        let collected: Vec<bool> = bv.iter().collect();
1877        assert_eq!(collected, vec![true]);
1878
1879        // Test iterator across chunk boundaries
1880        let mut bv: BitMap<4> = BitMap::new();
1881        // Fill exactly one chunk (32 bits)
1882        for i in 0..32 {
1883            bv.push(i % 2 == 0);
1884        }
1885        // Add a few more bits in the next chunk
1886        bv.push(true);
1887        bv.push(false);
1888        bv.push(true);
1889
1890        let collected: Vec<bool> = bv.iter().collect();
1891        assert_eq!(collected.len(), 35);
1892
1893        // Verify the pattern
1894        for (i, &bit) in collected.iter().enumerate().take(32) {
1895            assert_eq!(bit, i % 2 == 0);
1896        }
1897        assert!(collected[32]);
1898        assert!(!collected[33]);
1899        assert!(collected[34]);
1900    }
1901
1902    #[test]
1903    fn test_ones_iter_empty() {
1904        let bv: BitMap<4> = BitMap::new();
1905        let ones: Vec<u64> = bv.ones_iter().collect();
1906        assert!(ones.is_empty());
1907    }
1908
1909    #[test]
1910    fn test_ones_iter_all_zeros() {
1911        let bv = BitMap::<4>::zeroes(100);
1912        let ones: Vec<u64> = bv.ones_iter().collect();
1913        assert!(ones.is_empty());
1914    }
1915
1916    #[test]
1917    fn test_ones_iter_all_ones() {
1918        let bv = BitMap::<4>::ones(100);
1919        let ones: Vec<u64> = bv.ones_iter().collect();
1920        let expected: Vec<u64> = (0..100).collect();
1921        assert_eq!(ones, expected);
1922    }
1923
1924    #[test]
1925    fn test_ones_iter_sparse() {
1926        let mut bv = BitMap::<4>::zeroes(64);
1927        bv.set(0, true);
1928        bv.set(31, true);
1929        bv.set(32, true);
1930        bv.set(63, true);
1931
1932        let ones: Vec<u64> = bv.ones_iter().collect();
1933        assert_eq!(ones, vec![0, 31, 32, 63]);
1934    }
1935
1936    #[test]
1937    fn test_ones_iter_single_bit() {
1938        let mut bv: BitMap<4> = BitMap::new();
1939        bv.push(true);
1940        assert_eq!(bv.ones_iter().collect::<Vec<_>>(), vec![0]);
1941
1942        let mut bv: BitMap<4> = BitMap::new();
1943        bv.push(false);
1944        assert!(bv.ones_iter().collect::<Vec<_>>().is_empty());
1945    }
1946
1947    #[test]
1948    fn test_ones_iter_multi_chunk() {
1949        // Use small chunks (4 bytes = 32 bits) to ensure multi-chunk coverage.
1950        let mut bv = BitMap::<4>::zeroes(96);
1951        // Set one bit per chunk.
1952        bv.set(7, true); // chunk 0
1953        bv.set(40, true); // chunk 1
1954        bv.set(95, true); // chunk 2
1955
1956        let ones: Vec<u64> = bv.ones_iter().collect();
1957        assert_eq!(ones, vec![7, 40, 95]);
1958    }
1959
1960    #[test]
1961    fn test_ones_iter_partial_chunk() {
1962        // 35 bits = 1 full chunk (32 bits) + 3 bits in a partial chunk.
1963        let mut bv = BitMap::<4>::zeroes(35);
1964        bv.set(31, true); // last bit of full chunk
1965        bv.set(32, true); // first bit of partial chunk
1966        bv.set(34, true); // last bit
1967
1968        let ones: Vec<u64> = bv.ones_iter().collect();
1969        assert_eq!(ones, vec![31, 32, 34]);
1970    }
1971
1972    #[test]
1973    fn test_ones_iter_from_midway() {
1974        let mut bv = BitMap::<4>::zeroes(64);
1975        bv.set(5, true);
1976        bv.set(20, true);
1977        bv.set(40, true);
1978        bv.set(60, true);
1979
1980        // Start from position 20 -- should skip bit 5.
1981        let ones: Vec<u64> = Readable::ones_iter_from(&bv, 20).collect();
1982        assert_eq!(ones, vec![20, 40, 60]);
1983
1984        // Start from position 21 -- should skip bits 5 and 20.
1985        let ones: Vec<u64> = Readable::ones_iter_from(&bv, 21).collect();
1986        assert_eq!(ones, vec![40, 60]);
1987
1988        // Start past all set bits.
1989        let ones: Vec<u64> = Readable::ones_iter_from(&bv, 61).collect();
1990        assert!(ones.is_empty());
1991    }
1992
1993    #[test]
1994    fn test_ones_iter_matches_count_ones() {
1995        let mut bv: BitMap<8> = BitMap::new();
1996        for i in 0..200 {
1997            bv.push(i % 7 == 0);
1998        }
1999        assert_eq!(bv.ones_iter().count() as u64, bv.count_ones());
2000    }
2001
2002    #[test]
2003    fn test_ones_iter_different_chunk_sizes() {
2004        let pattern: Vec<bool> = (0..100).map(|i| i % 5 == 0).collect();
2005        let expected: Vec<u64> = (0..100).filter(|i| i % 5 == 0).collect();
2006
2007        let bv4: BitMap<4> = pattern.as_slice().into();
2008        let bv8: BitMap<8> = pattern.as_slice().into();
2009        let bv16: BitMap<16> = pattern.as_slice().into();
2010
2011        assert_eq!(bv4.ones_iter().collect::<Vec<_>>(), expected);
2012        assert_eq!(bv8.ones_iter().collect::<Vec<_>>(), expected);
2013        assert_eq!(bv16.ones_iter().collect::<Vec<_>>(), expected);
2014    }
2015
2016    #[test]
2017    fn test_codec_roundtrip() {
2018        // Test empty bitmap
2019        let original: BitMap<4> = BitMap::new();
2020        let encoded = original.encode();
2021        let decoded = BitMap::decode_cfg(&mut encoded.as_ref(), &(usize::MAX as u64)).unwrap();
2022        assert_eq!(original, decoded);
2023
2024        // Test small bitmap
2025        let pattern = [true, false, true, false, true];
2026        let original: BitMap<4> = pattern.as_ref().into();
2027        let encoded = original.encode();
2028        let decoded = BitMap::decode_cfg(&mut encoded.as_ref(), &(usize::MAX as u64)).unwrap();
2029        assert_eq!(original, decoded);
2030
2031        // Verify the decoded bitmap has the same bits
2032        for (i, &expected) in pattern.iter().enumerate() {
2033            assert_eq!(decoded.get(i as u64), expected);
2034        }
2035
2036        // Test larger bitmap across multiple chunks
2037        let mut large_original: BitMap<8> = BitMap::new();
2038        for i in 0..100 {
2039            large_original.push(i % 7 == 0);
2040        }
2041
2042        let encoded = large_original.encode();
2043        let decoded = BitMap::decode_cfg(&mut encoded.as_ref(), &(usize::MAX as u64)).unwrap();
2044        assert_eq!(large_original, decoded);
2045
2046        // Verify all bits match
2047        assert_eq!(decoded.len(), 100);
2048        for i in 0..100 {
2049            assert_eq!(decoded.get(i as u64), i % 7 == 0);
2050        }
2051    }
2052
2053    #[test]
2054    fn test_codec_different_chunk_sizes() {
2055        let pattern = [true, false, true, true, false, false, true];
2056
2057        // Test with different chunk sizes
2058        let bv4: BitMap<4> = pattern.as_ref().into();
2059        let bv8: BitMap<8> = pattern.as_ref().into();
2060        let bv16: BitMap<16> = pattern.as_ref().into();
2061
2062        // Encode and decode each
2063        let encoded4 = bv4.encode();
2064        let decoded4 = BitMap::decode_cfg(&mut encoded4.as_ref(), &(usize::MAX as u64)).unwrap();
2065        assert_eq!(bv4, decoded4);
2066
2067        let encoded8 = bv8.encode();
2068        let decoded8 = BitMap::decode_cfg(&mut encoded8.as_ref(), &(usize::MAX as u64)).unwrap();
2069        assert_eq!(bv8, decoded8);
2070
2071        let encoded16 = bv16.encode();
2072        let decoded16 = BitMap::decode_cfg(&mut encoded16.as_ref(), &(usize::MAX as u64)).unwrap();
2073        assert_eq!(bv16, decoded16);
2074
2075        // All should have the same logical content
2076        for (i, &expected) in pattern.iter().enumerate() {
2077            let i = i as u64;
2078            assert_eq!(decoded4.get(i), expected);
2079            assert_eq!(decoded8.get(i), expected);
2080            assert_eq!(decoded16.get(i), expected);
2081        }
2082    }
2083
2084    #[test]
2085    fn test_codec_edge_cases() {
2086        // Test bitmap with exactly one chunk filled
2087        let mut bv: BitMap<4> = BitMap::new();
2088        for i in 0..32 {
2089            bv.push(i % 2 == 0);
2090        }
2091
2092        let encoded = bv.encode();
2093        let decoded = BitMap::decode_cfg(&mut encoded.as_ref(), &(usize::MAX as u64)).unwrap();
2094        assert_eq!(bv, decoded);
2095        assert_eq!(decoded.len(), 32);
2096
2097        // Test bitmap with partial chunk
2098        let mut bv2: BitMap<4> = BitMap::new();
2099        for i in 0..35 {
2100            // 32 + 3 bits
2101            bv2.push(i % 3 == 0);
2102        }
2103
2104        let encoded2 = bv2.encode();
2105        let decoded2 = BitMap::decode_cfg(&mut encoded2.as_ref(), &(usize::MAX as u64)).unwrap();
2106        assert_eq!(bv2, decoded2);
2107        assert_eq!(decoded2.len(), 35);
2108    }
2109
2110    #[test]
2111    fn test_encode_size() {
2112        // Test encode size calculation
2113        let bv: BitMap<4> = BitMap::new();
2114        let encoded = bv.encode();
2115        assert_eq!(bv.encode_size(), encoded.len());
2116
2117        // Test with some data
2118        let pattern = [true, false, true, false, true];
2119        let bv: BitMap<4> = pattern.as_ref().into();
2120        let encoded = bv.encode();
2121        assert_eq!(bv.encode_size(), encoded.len());
2122
2123        // Test with larger data
2124        let mut large_bv: BitMap<8> = BitMap::new();
2125        for i in 0..100 {
2126            large_bv.push(i % 2 == 0);
2127        }
2128        let encoded = large_bv.encode();
2129        assert_eq!(large_bv.encode_size(), encoded.len());
2130    }
2131
2132    #[test]
2133    fn test_codec_empty_chunk_optimization() {
2134        // Test that empty last chunks are not serialized
2135
2136        // Case 1: Empty bitmap (omits the only empty chunk)
2137        let bv_empty: BitMap<4> = BitMap::new();
2138        let encoded_empty = bv_empty.encode();
2139        let decoded_empty: BitMap<4> =
2140            BitMap::decode_cfg(&mut encoded_empty.as_ref(), &(usize::MAX as u64)).unwrap();
2141        assert_eq!(bv_empty, decoded_empty);
2142        assert_eq!(bv_empty.len(), decoded_empty.len());
2143        // Should only encode the length, no chunks
2144        assert_eq!(encoded_empty.len(), bv_empty.len().encode_size());
2145
2146        // Case 2: Bitmap ending exactly at chunk boundary (omits empty last chunk)
2147        let mut bv_exact: BitMap<4> = BitMap::new();
2148        for _ in 0..32 {
2149            bv_exact.push(true);
2150        }
2151        let encoded_exact = bv_exact.encode();
2152        let decoded_exact: BitMap<4> =
2153            BitMap::decode_cfg(&mut encoded_exact.as_ref(), &(usize::MAX as u64)).unwrap();
2154        assert_eq!(bv_exact, decoded_exact);
2155
2156        // Case 3: Bitmap with partial last chunk (includes last chunk)
2157        let mut bv_partial: BitMap<4> = BitMap::new();
2158        for _ in 0..35 {
2159            bv_partial.push(true);
2160        }
2161        let encoded_partial = bv_partial.encode();
2162        let decoded_partial: BitMap<4> =
2163            BitMap::decode_cfg(&mut encoded_partial.as_ref(), &(usize::MAX as u64)).unwrap();
2164        assert_eq!(bv_partial, decoded_partial);
2165        assert_eq!(bv_partial.len(), decoded_partial.len());
2166
2167        // Verify optimization works correctly
2168        assert!(encoded_exact.len() < encoded_partial.len());
2169        assert_eq!(encoded_exact.len(), bv_exact.len().encode_size() + 4); // length + 1 chunk
2170        assert_eq!(encoded_partial.len(), bv_partial.len().encode_size() + 8); // length + 2 chunks
2171    }
2172
2173    #[test]
2174    fn test_codec_error_cases() {
2175        // Test invalid length with range check
2176        let mut buf = BytesMut::new();
2177        100u64.write(&mut buf); // bits length
2178
2179        // 100 bits requires 4 chunks (3 full + partially filled)
2180        for _ in 0..4 {
2181            [0u8; 4].write(&mut buf);
2182        }
2183
2184        // Test with a restricted range that excludes 100
2185        let result = BitMap::<4>::decode_cfg(&mut buf, &99);
2186        assert!(matches!(result, Err(CodecError::InvalidLength(100))));
2187
2188        // Test truncated buffer (not enough chunks)
2189        let mut buf = BytesMut::new();
2190        100u64.write(&mut buf); // bits length requiring 4 chunks (3 full + partially filled)
2191                                // Only write 3 chunks
2192        [0u8; 4].write(&mut buf);
2193        [0u8; 4].write(&mut buf);
2194        [0u8; 4].write(&mut buf);
2195
2196        let result = BitMap::<4>::decode_cfg(&mut buf, &(usize::MAX as u64));
2197        // Should fail when trying to read missing chunks
2198        assert!(result.is_err());
2199
2200        // Test invalid trailing bits
2201
2202        // Create a valid bitmap and encode it
2203        let original: BitMap<4> = BitMap::ones(20);
2204        let mut buf = BytesMut::new();
2205        original.write(&mut buf);
2206
2207        // Manually corrupt the encoded data by setting trailing bits
2208        let corrupted_data = buf.freeze();
2209        let mut corrupted_bytes = corrupted_data.to_vec();
2210
2211        // The last byte should have some trailing bits set to 1
2212        // For 20 bits with 4-byte chunks: 20 bits = 2.5 bytes, so last byte should have 4 valid bits
2213        // Set the high 4 bits of the last byte to 1 (these should be 0)
2214        let last_byte_idx = corrupted_bytes.len() - 1;
2215        corrupted_bytes[last_byte_idx] |= 0xF0;
2216
2217        // Read should fail
2218        let result = BitMap::<4>::read_cfg(&mut corrupted_bytes.as_slice(), &(usize::MAX as u64));
2219        assert!(matches!(
2220            result,
2221            Err(CodecError::Invalid(
2222                "BitMap",
2223                "Invalid trailing bits in encoded data"
2224            ))
2225        ));
2226    }
2227
2228    #[test]
2229    fn test_codec_range_config() {
2230        // Test RangeCfg validation in read_cfg
2231
2232        // Create a bitmap with 100 bits
2233        let mut original: BitMap<4> = BitMap::new();
2234        for i in 0..100 {
2235            original.push(i % 3 == 0);
2236        }
2237
2238        // Write to a buffer
2239        let mut buf = BytesMut::new();
2240        original.write(&mut buf);
2241
2242        // Test with max length < actual size (should fail)
2243        let result = BitMap::<4>::decode_cfg(&mut buf.as_ref(), &50);
2244        assert!(matches!(result, Err(CodecError::InvalidLength(100))));
2245
2246        // Test with max length == actual size (should succeed)
2247        let decoded = BitMap::<4>::decode_cfg(&mut buf.as_ref(), &100).unwrap();
2248        assert_eq!(decoded.len(), 100);
2249        assert_eq!(decoded, original);
2250
2251        // Test with max length > actual size (should succeed)
2252        let decoded = BitMap::<4>::decode_cfg(&mut buf.as_ref(), &101).unwrap();
2253        assert_eq!(decoded.len(), 100);
2254        assert_eq!(decoded, original);
2255
2256        // Test empty bitmap
2257        let empty = BitMap::<4>::new();
2258        let mut buf = BytesMut::new();
2259        empty.write(&mut buf);
2260
2261        // Empty bitmap should work with max length 0
2262        let decoded = BitMap::<4>::decode_cfg(&mut buf.as_ref(), &0).unwrap();
2263        assert_eq!(decoded.len(), 0);
2264        assert!(decoded.is_empty());
2265
2266        // Empty bitmap should work with max length > 0
2267        let decoded = BitMap::<4>::decode_cfg(&mut buf.as_ref(), &1).unwrap();
2268        assert_eq!(decoded.len(), 0);
2269        assert!(decoded.is_empty());
2270    }
2271
2272    #[test]
2273    fn test_from() {
2274        // Test From trait with different input types
2275
2276        // Test with Vec<bool>
2277        let vec_bool = vec![true, false, true, false, true];
2278        let bv: BitMap<4> = vec_bool.into();
2279        assert_eq!(bv.len(), 5);
2280        assert_eq!(bv.count_ones(), 3);
2281        assert_eq!(bv.count_zeros(), 2);
2282        for (i, &expected) in [true, false, true, false, true].iter().enumerate() {
2283            assert_eq!(bv.get(i as u64), expected);
2284        }
2285
2286        // Test with array slice
2287        let array = [false, true, true, false];
2288        let bv: BitMap<4> = (&array).into();
2289        assert_eq!(bv.len(), 4);
2290        assert_eq!(bv.count_ones(), 2);
2291        assert_eq!(bv.count_zeros(), 2);
2292        for (i, &expected) in array.iter().enumerate() {
2293            assert_eq!(bv.get(i as u64), expected);
2294        }
2295
2296        // Test with empty slice
2297        let empty: Vec<bool> = vec![];
2298        let bv: BitMap<4> = empty.into();
2299        assert_eq!(bv.len(), 0);
2300        assert!(bv.is_empty());
2301
2302        // Test with large slice
2303        let large: Vec<bool> = (0..100).map(|i| i % 3 == 0).collect();
2304        let bv: BitMap<8> = large.clone().into();
2305        assert_eq!(bv.len(), 100);
2306        for (i, &expected) in large.iter().enumerate() {
2307            assert_eq!(bv.get(i as u64), expected);
2308        }
2309    }
2310
2311    #[test]
2312    fn test_debug_formatting() {
2313        // Test Debug formatting for different sizes
2314
2315        // Test empty bitmap
2316        let bv: BitMap<4> = BitMap::new();
2317        let debug_str = format!("{bv:?}");
2318        assert_eq!(debug_str, "BitMap[]");
2319
2320        // Test small bitmap (should show all bits)
2321        let bv: BitMap<4> = [true, false, true, false, true].as_ref().into();
2322        let debug_str = format!("{bv:?}");
2323        assert_eq!(debug_str, "BitMap[10101]");
2324
2325        // Test bitmap at the display limit (64 bits)
2326        let pattern: Vec<bool> = (0..64).map(|i| i % 2 == 0).collect();
2327        let bv: BitMap<8> = pattern.into();
2328        let debug_str = format!("{bv:?}");
2329        let expected_pattern = "1010".repeat(16); // 64 bits alternating
2330        assert_eq!(debug_str, format!("BitMap[{expected_pattern}]"));
2331
2332        // Test large bitmap (should show ellipsis)
2333        let large_pattern: Vec<bool> = (0..100).map(|i| i % 2 == 0).collect();
2334        let bv: BitMap<16> = large_pattern.into();
2335        let debug_str = format!("{bv:?}");
2336
2337        // Should show first 32 bits + "..." + last 32 bits
2338        let first_32 = "10".repeat(16); // First 32 bits: 1010...
2339        let last_32 = "10".repeat(16); // Last 32 bits: ...1010
2340        let expected = format!("BitMap[{first_32}...{last_32}]");
2341        assert_eq!(debug_str, expected);
2342
2343        // Test single bit
2344        let bv: BitMap<4> = [true].as_ref().into();
2345        assert_eq!(format!("{bv:?}"), "BitMap[1]");
2346
2347        let bv: BitMap<4> = [false].as_ref().into();
2348        assert_eq!(format!("{bv:?}"), "BitMap[0]");
2349
2350        // Test exactly at boundary (65 bits - should show ellipsis)
2351        let pattern: Vec<bool> = (0..65).map(|i| i == 0 || i == 64).collect(); // First and last bits are true
2352        let bv: BitMap<16> = pattern.into();
2353        let debug_str = format!("{bv:?}");
2354
2355        // Should show first 32 bits (100000...) + "..." + last 32 bits (...000001)
2356        let first_32 = "1".to_string() + &"0".repeat(31);
2357        let last_32 = "0".repeat(31) + "1";
2358        let expected = format!("BitMap[{first_32}...{last_32}]");
2359        assert_eq!(debug_str, expected);
2360    }
2361
2362    #[test]
2363    fn test_from_different_chunk_sizes() {
2364        // Test From trait works with different chunk sizes
2365        let pattern = [true, false, true, true, false, false, true];
2366
2367        let bv4: BitMap<4> = pattern.as_ref().into();
2368        let bv8: BitMap<8> = pattern.as_ref().into();
2369        let bv16: BitMap<16> = pattern.as_ref().into();
2370
2371        // All should have the same content regardless of chunk size
2372        // Test each bitmap separately since they have different types
2373        for bv in [&bv4] {
2374            assert_eq!(bv.len(), 7);
2375            assert_eq!(bv.count_ones(), 4);
2376            assert_eq!(bv.count_zeros(), 3);
2377            for (i, &expected) in pattern.iter().enumerate() {
2378                assert_eq!(bv.get(i as u64), expected);
2379            }
2380        }
2381
2382        assert_eq!(bv8.len(), 7);
2383        assert_eq!(bv8.count_ones(), 4);
2384        assert_eq!(bv8.count_zeros(), 3);
2385        for (i, &expected) in pattern.iter().enumerate() {
2386            assert_eq!(bv8.get(i as u64), expected);
2387        }
2388
2389        assert_eq!(bv16.len(), 7);
2390        assert_eq!(bv16.count_ones(), 4);
2391        assert_eq!(bv16.count_zeros(), 3);
2392        for (i, &expected) in pattern.iter().enumerate() {
2393            assert_eq!(bv16.get(i as u64), expected);
2394        }
2395    }
2396
2397    #[test]
2398    fn test_prune_chunks() {
2399        let mut bv: BitMap<4> = BitMap::new();
2400        bv.push_chunk(&[1, 2, 3, 4]);
2401        bv.push_chunk(&[5, 6, 7, 8]);
2402        bv.push_chunk(&[9, 10, 11, 12]);
2403
2404        assert_eq!(bv.len(), 96);
2405        assert_eq!(bv.get_chunk(0), &[1, 2, 3, 4]);
2406
2407        // Prune first chunk
2408        bv.prune_chunks(1);
2409        assert_eq!(bv.len(), 64);
2410        assert_eq!(bv.get_chunk(0), &[5, 6, 7, 8]);
2411        assert_eq!(bv.get_chunk(1), &[9, 10, 11, 12]);
2412
2413        // Prune another chunk
2414        bv.prune_chunks(1);
2415        assert_eq!(bv.len(), 32);
2416        assert_eq!(bv.get_chunk(0), &[9, 10, 11, 12]);
2417    }
2418
2419    #[test]
2420    #[should_panic(expected = "cannot prune")]
2421    fn test_prune_too_many_chunks() {
2422        let mut bv: BitMap<4> = BitMap::new();
2423        bv.push_chunk(&[1, 2, 3, 4]);
2424        bv.push_chunk(&[5, 6, 7, 8]);
2425        bv.push(true);
2426
2427        // Try to prune 4 chunks when only 3 are available
2428        bv.prune_chunks(4);
2429    }
2430
2431    #[test]
2432    fn test_prune_with_partial_last_chunk() {
2433        let mut bv: BitMap<4> = BitMap::new();
2434        bv.push_chunk(&[1, 2, 3, 4]);
2435        bv.push_chunk(&[5, 6, 7, 8]);
2436        bv.push(true);
2437        bv.push(false);
2438
2439        assert_eq!(bv.len(), 66);
2440
2441        // Can prune first chunk
2442        bv.prune_chunks(1);
2443        assert_eq!(bv.len(), 34);
2444        assert_eq!(bv.get_chunk(0), &[5, 6, 7, 8]);
2445
2446        // Last partial chunk still has the appended bits
2447        assert!(bv.get(32));
2448        assert!(!bv.get(33));
2449    }
2450
2451    #[test]
2452    fn test_prune_all_chunks_resets_next_bit() {
2453        let mut bv: BitMap<4> = BitMap::new();
2454        bv.push_chunk(&[1, 2, 3, 4]);
2455        bv.push_chunk(&[5, 6, 7, 8]);
2456        bv.push(true);
2457        bv.push(false);
2458        bv.push(true);
2459
2460        // Bitmap has 2 full chunks + 3 bits in partial chunk
2461        assert_eq!(bv.len(), 67);
2462
2463        // Prune all chunks (this leaves chunks empty, triggering the reset path)
2464        bv.prune_chunks(3);
2465
2466        // Regression test: len() should be 0, not the old next_bit value (3)
2467        assert_eq!(bv.len(), 0);
2468        assert!(bv.is_empty());
2469
2470        // Bitmap should behave as freshly created
2471        bv.push(true);
2472        assert_eq!(bv.len(), 1);
2473        assert!(bv.get(0));
2474    }
2475
2476    #[test]
2477    fn test_is_chunk_aligned() {
2478        // Empty bitmap is chunk aligned
2479        let bv: BitMap<4> = BitMap::new();
2480        assert!(bv.is_chunk_aligned());
2481
2482        // Test with various chunk sizes
2483        let mut bv4: BitMap<4> = BitMap::new();
2484        assert!(bv4.is_chunk_aligned());
2485
2486        // Add bits one at a time and check alignment
2487        for i in 1..=32 {
2488            bv4.push(i % 2 == 0);
2489            if i == 32 {
2490                assert!(bv4.is_chunk_aligned()); // Exactly one chunk
2491            } else {
2492                assert!(!bv4.is_chunk_aligned()); // Partial chunk
2493            }
2494        }
2495
2496        // Add more bits
2497        for i in 33..=64 {
2498            bv4.push(i % 2 == 0);
2499            if i == 64 {
2500                assert!(bv4.is_chunk_aligned()); // Exactly two chunks
2501            } else {
2502                assert!(!bv4.is_chunk_aligned()); // Partial chunk
2503            }
2504        }
2505
2506        // Test with push_chunk
2507        let mut bv: BitMap<8> = BitMap::new();
2508        assert!(bv.is_chunk_aligned());
2509        bv.push_chunk(&[0xFF; 8]);
2510        assert!(bv.is_chunk_aligned()); // 64 bits = 1 chunk for N=8
2511        bv.push_chunk(&[0xAA; 8]);
2512        assert!(bv.is_chunk_aligned()); // 128 bits = 2 chunks
2513        bv.push(true);
2514        assert!(!bv.is_chunk_aligned()); // 129 bits = partial chunk
2515
2516        // Test with push_byte
2517        let mut bv: BitMap<4> = BitMap::new();
2518        for _ in 0..4 {
2519            bv.push_byte(0xFF);
2520        }
2521        assert!(bv.is_chunk_aligned()); // 32 bits = 1 chunk for N=4
2522
2523        // Test after pop
2524        bv.pop();
2525        assert!(!bv.is_chunk_aligned()); // 31 bits = partial chunk
2526
2527        // Test with zeroes and ones constructors
2528        let bv_zeroes: BitMap<4> = BitMap::zeroes(64);
2529        assert!(bv_zeroes.is_chunk_aligned());
2530
2531        let bv_ones: BitMap<4> = BitMap::ones(96);
2532        assert!(bv_ones.is_chunk_aligned());
2533
2534        let bv_partial: BitMap<4> = BitMap::zeroes(65);
2535        assert!(!bv_partial.is_chunk_aligned());
2536    }
2537
2538    #[test]
2539    fn test_unprune_restores_length() {
2540        let mut prunable: Prunable<4> = Prunable::new_with_pruned_chunks(1).unwrap();
2541        assert_eq!(prunable.len(), Prunable::<4>::CHUNK_SIZE_BITS);
2542        assert_eq!(prunable.pruned_chunks(), 1);
2543        let chunk = [0xDE, 0xAD, 0xBE, 0xEF];
2544
2545        prunable.unprune_chunks(&[chunk]);
2546
2547        assert_eq!(prunable.pruned_chunks(), 0);
2548        assert_eq!(prunable.len(), Prunable::<4>::CHUNK_SIZE_BITS);
2549        assert_eq!(prunable.get_chunk_containing(0), &chunk);
2550    }
2551
2552    mod proptests {
2553        use super::*;
2554        use proptest::prelude::*;
2555
2556        proptest! {
2557            #[test]
2558            fn is_unset_matches_naive(
2559                bits in prop::collection::vec(any::<bool>(), 1..=512usize),
2560                start in 0u64..=512,
2561                end in 0u64..=512,
2562            ) {
2563                let bitmap: BitMap = BitMap::from(bits.as_slice());
2564                let len = bitmap.len();
2565                let start = start.min(len);
2566                let end = end.max(start).min(len);
2567                let range = start..end;
2568
2569                let expected = range.clone().all(|i| !bitmap.get(i));
2570
2571                prop_assert_eq!(bitmap.is_unset(range), expected);
2572            }
2573        }
2574    }
2575
2576    #[test]
2577    fn is_unset_all_zeros() {
2578        let bitmap = BitMap::<8>::zeroes(256);
2579        assert!(bitmap.is_unset(0..256));
2580    }
2581
2582    #[test]
2583    fn is_unset_all_ones() {
2584        let bitmap = BitMap::<8>::ones(256);
2585        assert!(!bitmap.is_unset(0..256));
2586    }
2587
2588    #[test]
2589    fn is_unset_single_bit() {
2590        let mut bitmap = BitMap::<8>::zeroes(64);
2591        bitmap.set(31, true);
2592        assert!(bitmap.is_unset(0..31));
2593        assert!(!bitmap.is_unset(0..32));
2594        assert!(!bitmap.is_unset(31..32));
2595        assert!(bitmap.is_unset(32..64));
2596    }
2597
2598    #[test]
2599    fn is_unset_empty_range() {
2600        let bitmap = BitMap::<8>::ones(64);
2601        assert!(bitmap.is_unset(0..0));
2602        assert!(bitmap.is_unset(32..32));
2603        assert!(bitmap.is_unset(64..64));
2604    }
2605
2606    #[test]
2607    fn is_unset_chunk_boundaries() {
2608        // N=1 means 8 bits per chunk, so boundaries are more frequent
2609        let mut bitmap = BitMap::<1>::zeroes(32);
2610        bitmap.set(7, true);
2611        assert!(bitmap.is_unset(0..7));
2612        assert!(!bitmap.is_unset(0..8));
2613        assert!(bitmap.is_unset(8..32));
2614    }
2615
2616    #[test]
2617    fn is_unset_small_chunk_multi_span() {
2618        // N=4 means 32 bits per chunk, test spanning 3 chunks
2619        let mut bitmap = BitMap::<4>::zeroes(128);
2620        bitmap.set(96, true);
2621        assert!(bitmap.is_unset(0..96));
2622        assert!(!bitmap.is_unset(0..97));
2623        assert!(bitmap.is_unset(97..128));
2624    }
2625
2626    #[test]
2627    #[should_panic(expected = "out of bounds")]
2628    fn is_unset_out_of_bounds() {
2629        let bitmap = BitMap::<8>::zeroes(64);
2630        bitmap.is_unset(0..65);
2631    }
2632
2633    #[cfg(feature = "arbitrary")]
2634    mod conformance {
2635        use super::*;
2636        use commonware_codec::conformance::CodecConformance;
2637
2638        commonware_conformance::conformance_tests! {
2639            CodecConformance<BitMap>
2640        }
2641    }
2642}