Skip to main content

commonware_utils/bitmap/historical/
bitmap.rs

1use super::Error;
2use crate::bitmap::Prunable;
3#[cfg(not(feature = "std"))]
4use alloc::{collections::BTreeMap, vec::Vec};
5#[cfg(feature = "std")]
6use std::collections::BTreeMap;
7
8/// Sealed trait for bitmap state types.
9mod private {
10    pub trait Sealed {}
11}
12
13/// Trait for valid bitmap type states [Clean], [Dirty].
14pub trait State: private::Sealed + Sized + Send + Sync {}
15
16/// Clean bitmap type state. The bitmap has no pending mutations.
17#[derive(Clone, Debug)]
18pub struct Clean;
19
20impl private::Sealed for Clean {}
21impl State for Clean {}
22
23/// Dirty bitmap type state. The bitmap has pending mutations not yet committed.
24///
25/// # De-duplication and Cancellation
26///
27/// **The dirty state de-duplicates during operations, not at commit time.**
28///
29/// Operations that cancel out are handled automatically:
30///
31/// ```text
32/// Example 1: push + pop = no-op
33///   push(true)  -> appended_bits=[true], projected_len=11
34///   pop()       -> appended_bits=[], projected_len=10
35///   Result: dirty state unchanged from base
36///
37/// Example 2: set_bit + set_bit = last write wins
38///   set_bit(5, true)   -> modified_bits={5: true}
39///   set_bit(5, false)  -> modified_bits={5: false}
40///   Result: only final value recorded
41///
42/// Example 3: set_bit + pop = cancels modification
43///   set_bit(9, true)  -> modified_bits={9: true}
44///   pop()             -> modified_bits={} (removed), projected_len=9
45///   Result: bit 9 no longer exists, modification discarded
46/// ```
47///
48/// # Key Invariants
49///
50/// 1. **Base immutability**: `base_len` and `base_pruned_chunks` never change
51/// 2. **Appended region**: Always occupies `[projected_len - appended_bits.len(), projected_len)`
52/// 3. **Modified region**: `modified_bits` only contains offsets in `[0, projected_len - appended_bits.len())`
53///    - These are modifications to the base bitmap, never to appended bits
54///    - Appended bits are modified by directly updating the `appended_bits` vector
55/// 4. **No overlap**: A bit is either in `modified_bits` OR `appended_bits`, never both
56#[derive(Clone, Debug)]
57pub struct Dirty<const N: usize> {
58    /// Bitmap state when dirty started (immutable).
59    base_len: u64,
60    base_pruned_chunks: usize,
61
62    /// What the bitmap will look like after commit (mutable).
63    projected_len: u64,
64    projected_pruned_chunks: usize,
65
66    /// Modifications to bits that existed in the bitmap (not appended bits).
67    /// Contains offsets in [0, projected_len - appended_bits.len()).
68    /// Maps: bit -> new_value
69    modified_bits: BTreeMap<u64, bool>,
70
71    /// New bits pushed in this dirty state (in order).
72    /// Logical position: [projected_len - appended_bits.len(), projected_len)
73    appended_bits: Vec<bool>,
74
75    /// Old chunk data for chunks being pruned.
76    /// Captured eagerly during `prune_to_bit()` for historical reconstruction.
77    chunks_to_prune: BTreeMap<usize, [u8; N]>,
78}
79
80impl<const N: usize> private::Sealed for Dirty<N> {}
81impl<const N: usize> State for Dirty<N> {}
82
83/// A change to a chunk.
84#[derive(Clone, Debug)]
85pub(super) enum ChunkDiff<const N: usize> {
86    /// Chunk was modified (contains old value before the change).
87    Modified([u8; N]),
88    /// Chunk was removed from the right side (contains old value before removal).
89    Removed([u8; N]),
90    /// Chunk was added (did not exist before).
91    Added,
92    /// Chunk was pruned from the left side (contains old value before pruning).
93    Pruned([u8; N]),
94}
95
96/// A reverse diff that describes the state before a commit.
97#[derive(Clone, Debug)]
98pub(super) struct CommitDiff<const N: usize> {
99    /// Total length in bits before this commit.
100    pub(super) len: u64,
101    /// Number of pruned chunks before this commit.
102    pub(super) pruned_chunks: usize,
103    /// Chunk-level changes.
104    pub(super) chunk_diffs: BTreeMap<usize, ChunkDiff<N>>,
105}
106
107/// A historical bitmap that maintains one actual bitmap plus diffs for history.
108///
109/// Uses a type-state pattern to track whether the bitmap is clean (no pending
110/// mutations) or dirty (has pending mutations).
111///
112/// Commit numbers must be strictly monotonically increasing and < u64::MAX.
113#[derive(Clone, Debug)]
114pub struct BitMap<const N: usize, S: State = Clean> {
115    /// The current/HEAD state. This is the only full bitmap.
116    current: Prunable<N>,
117
118    /// Historical commits: commit_number -> reverse diff from that commit.
119    commits: BTreeMap<u64, CommitDiff<N>>,
120
121    /// State marker (Clean or Dirty).
122    state: S,
123}
124
125/// Type alias for a clean bitmap with no pending mutations.
126pub type CleanBitMap<const N: usize> = BitMap<N, Clean>;
127
128/// Type alias for a dirty bitmap with pending mutations.
129pub type DirtyBitMap<const N: usize> = BitMap<N, Dirty<N>>;
130
131impl<const N: usize> CleanBitMap<N> {
132    /// Create a new empty historical bitmap.
133    pub const fn new() -> Self {
134        Self {
135            current: Prunable::new(),
136            commits: BTreeMap::new(),
137            state: Clean,
138        }
139    }
140
141    /// Create a new historical bitmap with the given number of pruned chunks.
142    pub fn new_with_pruned_chunks(pruned_chunks: usize) -> Result<Self, Error> {
143        Ok(Self {
144            current: Prunable::new_with_pruned_chunks(pruned_chunks)?,
145            commits: BTreeMap::new(),
146            state: Clean,
147        })
148    }
149
150    /// Transition to dirty state to begin making mutations.
151    ///
152    /// All mutations are applied to a diff layer and do not affect the current
153    /// bitmap until commit.
154    ///
155    /// # Examples
156    ///
157    /// ```
158    /// # use commonware_utils::bitmap::historical::BitMap;
159    /// let bitmap: BitMap<4> = BitMap::new();
160    ///
161    /// let mut dirty = bitmap.into_dirty();
162    /// dirty.push(true);
163    /// dirty.push(false);
164    /// let bitmap = dirty.commit(1).unwrap();
165    ///
166    /// assert_eq!(bitmap.len(), 2);
167    /// ```
168    pub fn into_dirty(self) -> DirtyBitMap<N> {
169        DirtyBitMap {
170            state: Dirty {
171                base_len: self.current.len(),
172                base_pruned_chunks: self.current.pruned_chunks(),
173                projected_len: self.current.len(),
174                projected_pruned_chunks: self.current.pruned_chunks(),
175                modified_bits: BTreeMap::new(),
176                appended_bits: Vec::new(),
177                chunks_to_prune: BTreeMap::new(),
178            },
179            current: self.current,
180            commits: self.commits,
181        }
182    }
183
184    /// Execute a closure with a dirty bitmap and commit it at the given commit number.
185    ///
186    /// # Errors
187    ///
188    /// Returns [Error::NonMonotonicCommit] if the commit number is not
189    /// greater than the previous commit.
190    ///
191    /// Returns [Error::ReservedCommitNumber] if the commit number is `u64::MAX`.
192    ///
193    /// # Examples
194    ///
195    /// ```
196    /// # use commonware_utils::bitmap::historical::BitMap;
197    /// let mut bitmap: BitMap<4> = BitMap::new();
198    ///
199    /// bitmap = bitmap.apply_batch(1, |dirty| {
200    ///     dirty.push(true).push(false);
201    /// }).unwrap();
202    ///
203    /// assert_eq!(bitmap.len(), 2);
204    /// ```
205    pub fn apply_batch<F>(self, commit_number: u64, f: F) -> Result<Self, Error>
206    where
207        F: FnOnce(&mut DirtyBitMap<N>),
208    {
209        let mut dirty = self.into_dirty();
210        f(&mut dirty);
211        dirty.commit(commit_number)
212    }
213
214    /// Get the bitmap state as it existed at a specific commit.
215    ///
216    /// Returns `None` if the commit does not exist or if `commit_number` is `u64::MAX`
217    /// (which is reserved and cannot be used as a commit number).
218    ///
219    /// This reconstructs the historical state by applying reverse diffs backward from
220    /// the current state. Each commit's reverse diff describes the state before that
221    /// commit, so we "undo" commits one by one until we reach the target.
222    ///
223    /// # Examples
224    ///
225    /// ```
226    /// # use commonware_utils::bitmap::historical::BitMap;
227    /// let mut bitmap: BitMap<4> = BitMap::new();
228    ///
229    /// bitmap = bitmap.apply_batch(1, |dirty| {
230    ///     dirty.push(true);
231    ///     dirty.push(false);
232    /// }).unwrap();
233    ///
234    /// bitmap = bitmap.apply_batch(2, |dirty| {
235    ///     dirty.set_bit(0, false);
236    ///     dirty.push(true);
237    /// }).unwrap();
238    ///
239    /// // Get state as it was at commit 1
240    /// let state_at_1 = bitmap.get_at_commit(1).unwrap();
241    /// assert_eq!(state_at_1.len(), 2);
242    /// assert!(state_at_1.get_bit(0));
243    /// assert!(!state_at_1.get_bit(1));
244    ///
245    /// // Current state is different
246    /// assert_eq!(bitmap.len(), 3);
247    /// assert!(!bitmap.get_bit(0));
248    /// ```
249    pub fn get_at_commit(&self, commit_number: u64) -> Option<Prunable<N>> {
250        // Check if the commit exists and is valid
251        if commit_number == u64::MAX || !self.commits.contains_key(&commit_number) {
252            return None;
253        }
254
255        // Start with current state
256        let mut state = self.current.clone();
257
258        // Apply reverse diffs from newest down to target (exclusive)
259        // Each reverse diff at commit N describes the state before commit N
260        // Addition can't overflow because commit_number < u64::MAX
261        for (_commit, diff) in self.commits.range(commit_number + 1..).rev() {
262            self.apply_reverse_diff(&mut state, diff);
263        }
264
265        Some(state)
266    }
267
268    /// Check if a commit exists.
269    pub fn commit_exists(&self, commit_number: u64) -> bool {
270        self.commits.contains_key(&commit_number)
271    }
272
273    /// Get an iterator over all commit numbers in ascending order.
274    pub fn commits(&self) -> impl Iterator<Item = u64> + '_ {
275        self.commits.keys().copied()
276    }
277
278    /// Get the latest commit number, if any commits exist.
279    pub fn latest_commit(&self) -> Option<u64> {
280        self.commits.keys().next_back().copied()
281    }
282
283    /// Get the earliest commit number, if any commits exist.
284    pub fn earliest_commit(&self) -> Option<u64> {
285        self.commits.keys().next().copied()
286    }
287
288    /// Get a reference to the current bitmap state.
289    pub const fn current(&self) -> &Prunable<N> {
290        &self.current
291    }
292
293    /// Number of bits in the current bitmap.
294    #[inline]
295    pub const fn len(&self) -> u64 {
296        self.current.len()
297    }
298
299    /// Returns true if the current bitmap is empty.
300    #[inline]
301    pub const fn is_empty(&self) -> bool {
302        self.current.is_empty()
303    }
304
305    /// Get the value of a bit in the current bitmap.
306    #[inline]
307    pub fn get_bit(&self, bit: u64) -> bool {
308        self.current.get_bit(bit)
309    }
310
311    /// Get the chunk containing a bit in the current bitmap.
312    #[inline]
313    pub fn get_chunk_containing(&self, bit: u64) -> &[u8; N] {
314        self.current.get_chunk_containing(bit)
315    }
316
317    /// Number of pruned chunks in the current bitmap.
318    #[inline]
319    pub const fn pruned_chunks(&self) -> usize {
320        self.current.pruned_chunks()
321    }
322
323    /// Remove all commits with numbers below the commit number.
324    ///
325    /// Returns the number of commits removed.
326    pub fn prune_commits_before(&mut self, commit_number: u64) -> usize {
327        let count = self.commits.len();
328        self.commits = self.commits.split_off(&commit_number);
329        count - self.commits.len()
330    }
331
332    /// Clear all historical commits.
333    pub fn clear_history(&mut self) {
334        self.commits.clear();
335    }
336
337    /// Apply a reverse diff to transform newer_state into the previous state (in-place).
338    ///
339    /// Algorithm:
340    /// 1. Restore pruned chunks by prepending them back (unprune)
341    /// 2. Adjust bitmap structure to target length (extend/shrink as needed)
342    /// 3. Update chunk data for Modified and Removed chunks
343    /// 4. Set next_bit to match target length exactly
344    fn apply_reverse_diff(&self, newer_state: &mut Prunable<N>, diff: &CommitDiff<N>) {
345        let target_len = diff.len;
346        let target_pruned = diff.pruned_chunks;
347        let newer_pruned = newer_state.pruned_chunks();
348
349        // Phase 1: Restore pruned chunks
350        assert!(
351            target_pruned <= newer_pruned,
352            "invariant violation: target_pruned ({target_pruned}) > newer_pruned ({newer_pruned})"
353        );
354        let mut chunks_to_unprune = Vec::with_capacity(newer_pruned - target_pruned);
355        for chunk_index in (target_pruned..newer_pruned).rev() {
356            let Some(ChunkDiff::Pruned(chunk)) = diff.chunk_diffs.get(&chunk_index) else {
357                panic!("chunk {chunk_index} should be Pruned in diff");
358            };
359            chunks_to_unprune.push(*chunk);
360        }
361        newer_state.unprune_chunks(&chunks_to_unprune);
362
363        // Phase 2: Adjust bitmap structure to target length
364        if newer_state.len() < target_len {
365            newer_state.extend_to(target_len);
366        } else {
367            newer_state.truncate(target_len);
368        }
369
370        // Phase 3: Update chunk data
371        for (&chunk_index, change) in diff
372            .chunk_diffs
373            .iter()
374            .filter(|(chunk_index, _)| **chunk_index >= newer_pruned)
375        {
376            match change {
377                ChunkDiff::Modified(old_data) | ChunkDiff::Removed(old_data) => {
378                    // Both cases: chunk exists in target, just update its data
379                    newer_state.set_chunk_by_index(chunk_index, old_data);
380                }
381                ChunkDiff::Added => {
382                    // Chunk didn't exist in target, so phase 2 already removed it.
383                    // We can break here because there are no more modifications to apply.
384                    // Added can only occur after all Modified. If we encounter Added, we know
385                    // there are no Removed. (diff.chunk_diffs can't have both Added and Removed.)
386                    break;
387                }
388                ChunkDiff::Pruned(_) => {
389                    panic!("pruned chunk found at unexpected index {chunk_index}")
390                }
391            }
392        }
393
394        assert_eq!(newer_state.pruned_chunks(), target_pruned);
395        assert_eq!(newer_state.len(), target_len);
396    }
397}
398
399impl<const N: usize> Default for CleanBitMap<N> {
400    fn default() -> Self {
401        Self::new()
402    }
403}
404
405impl<const N: usize> DirtyBitMap<N> {
406    /// Get the length of the bitmap as it would be after committing.
407    #[inline]
408    pub const fn len(&self) -> u64 {
409        self.state.projected_len
410    }
411
412    /// Returns true if the bitmap would be empty after committing.
413    #[inline]
414    pub const fn is_empty(&self) -> bool {
415        self.len() == 0
416    }
417
418    /// Get the number of pruned chunks after committing.
419    #[inline]
420    pub const fn pruned_chunks(&self) -> usize {
421        self.state.projected_pruned_chunks
422    }
423
424    /// Get a bit value with read-through semantics.
425    ///
426    /// Returns the bit's value as it would be after committing.
427    /// Priority: appended bits > modified bits > original bitmap.
428    ///
429    /// # Panics
430    ///
431    /// Panics if the bit offset is out of bounds or if the bit has been pruned.
432    pub fn get_bit(&self, bit: u64) -> bool {
433        assert!(
434            bit < self.state.projected_len,
435            "bit offset {bit} out of bounds (len: {})",
436            self.state.projected_len
437        );
438
439        let chunk_idx = Prunable::<N>::to_chunk_index(bit);
440        assert!(
441            chunk_idx >= self.state.projected_pruned_chunks,
442            "cannot get bit {bit}: chunk {chunk_idx} is pruned (pruned up to chunk {})",
443            self.state.projected_pruned_chunks
444        );
445
446        // Priority 1: Check if bit is in appended region.
447        // Must use appended_start, not base_len, to handle net pops + appends.
448        let appended_start = self.state.projected_len - self.state.appended_bits.len() as u64;
449        if bit >= appended_start {
450            let append_offset = (bit - appended_start) as usize;
451            return self.state.appended_bits[append_offset];
452        }
453
454        // Priority 2: Check if bit was modified.
455        if let Some(&value) = self.state.modified_bits.get(&bit) {
456            return value;
457        }
458
459        // Priority 3: Fall through to original bitmap.
460        self.current.get_bit(bit)
461    }
462
463    /// Get a chunk value with read-through semantics.
464    ///
465    /// Reconstructs the chunk if it has modifications, otherwise returns from current.
466    ///
467    /// # Panics
468    ///
469    /// Panics if the bit offset is out of bounds or if the chunk has been pruned.
470    pub fn get_chunk(&self, bit: u64) -> [u8; N] {
471        // Check bounds
472        assert!(
473            bit < self.state.projected_len,
474            "bit offset {bit} out of bounds (len: {})",
475            self.state.projected_len
476        );
477
478        let chunk_idx = Prunable::<N>::to_chunk_index(bit);
479
480        // Check if chunk is in pruned range
481        assert!(
482            chunk_idx >= self.state.projected_pruned_chunks,
483            "cannot get chunk at bit offset {bit}: chunk {chunk_idx} is pruned (pruned up to chunk {})",
484            self.state.projected_pruned_chunks
485        );
486
487        let chunk_start_bit = chunk_idx as u64 * Prunable::<N>::CHUNK_SIZE_BITS;
488        let chunk_end_bit = chunk_start_bit + Prunable::<N>::CHUNK_SIZE_BITS;
489
490        // Determine if this chunk needs reconstruction.
491        let appended_start = self.state.projected_len - self.state.appended_bits.len() as u64;
492
493        // Skip reconstruction only if chunk is entirely outside modified regions
494        let chunk_entirely_past_end = chunk_start_bit >= self.state.projected_len;
495        let chunk_entirely_before_changes =
496            chunk_end_bit <= appended_start && chunk_end_bit <= self.state.projected_len;
497
498        let chunk_needs_reconstruction =
499            // Chunk overlaps with pops or appends
500            !(chunk_entirely_past_end || chunk_entirely_before_changes)
501            // OR chunk has explicit bit modifications
502            || (chunk_start_bit..chunk_end_bit.min(self.state.base_len))
503                .any(|bit| self.state.modified_bits.contains_key(&bit));
504
505        if chunk_needs_reconstruction {
506            // Reconstruct chunk from current + modifications
507            self.reconstruct_modified_chunk(chunk_start_bit)
508        } else {
509            // Fall through to current bitmap
510            *self.current.get_chunk_containing(bit)
511        }
512    }
513
514    /// Reconstruct a chunk that has modifications, appends, or pops.
515    fn reconstruct_modified_chunk(&self, chunk_start: u64) -> [u8; N] {
516        // Start with current chunk if it exists
517        let mut chunk = if chunk_start < self.current.len() {
518            *self.current.get_chunk_containing(chunk_start)
519        } else {
520            [0u8; N]
521        };
522
523        // Calculate appended region boundary
524        let appended_start = self.state.projected_len - self.state.appended_bits.len() as u64;
525
526        // Apply modifications and zero out popped bits
527        for bit_in_chunk in 0..Prunable::<N>::CHUNK_SIZE_BITS {
528            let bit = chunk_start + bit_in_chunk;
529
530            let byte_idx = (bit_in_chunk / 8) as usize;
531            let bit_idx = bit_in_chunk % 8;
532            let mask = 1u8 << bit_idx;
533
534            if bit >= self.state.projected_len {
535                // Bit is beyond projected length (popped), zero it out
536                chunk[byte_idx] &= !mask;
537            } else if let Some(&value) = self.state.modified_bits.get(&bit) {
538                // Bit was explicitly modified
539                if value {
540                    chunk[byte_idx] |= mask;
541                } else {
542                    chunk[byte_idx] &= !mask;
543                }
544            } else if bit >= appended_start {
545                // This is an appended bit
546                let append_offset = (bit - appended_start) as usize;
547                if append_offset < self.state.appended_bits.len() {
548                    let value = self.state.appended_bits[append_offset];
549                    if value {
550                        chunk[byte_idx] |= mask;
551                    } else {
552                        chunk[byte_idx] &= !mask;
553                    }
554                }
555            }
556        }
557
558        chunk
559    }
560
561    /// Set a bit value.
562    ///
563    /// # Panics
564    ///
565    /// Panics if the bit offset is out of bounds or if the bit has been pruned.
566    pub fn set_bit(&mut self, bit: u64, value: bool) -> &mut Self {
567        assert!(
568            bit < self.state.projected_len,
569            "cannot set bit {bit}: out of bounds (len: {})",
570            self.state.projected_len
571        );
572
573        let chunk_idx = Prunable::<N>::to_chunk_index(bit);
574        assert!(
575            chunk_idx >= self.state.projected_pruned_chunks,
576            "cannot set bit {bit}: chunk {chunk_idx} is pruned (pruned up to chunk {})",
577            self.state.projected_pruned_chunks
578        );
579
580        // Determine which region this bit belongs to.
581        // Appended region: bits pushed, starting at projected_len - appended_bits.len()
582        let appended_start = self.state.projected_len - self.state.appended_bits.len() as u64;
583
584        if bit >= appended_start {
585            // Bit is in the appended region: update the appended_bits vector directly.
586            let append_offset = (bit - appended_start) as usize;
587            self.state.appended_bits[append_offset] = value;
588        } else {
589            // Bit is in the base region: record as a modification.
590            self.state.modified_bits.insert(bit, value);
591        }
592
593        self
594    }
595
596    /// Push a bit to the end of the bitmap.
597    pub fn push(&mut self, bit: bool) -> &mut Self {
598        self.state.appended_bits.push(bit);
599        self.state.projected_len += 1;
600        self
601    }
602
603    /// Push a byte to the end of the bitmap.
604    pub fn push_byte(&mut self, byte: u8) -> &mut Self {
605        for i in 0..8 {
606            let bit = (byte >> i) & 1 == 1;
607            self.push(bit);
608        }
609        self
610    }
611
612    /// Push a full chunk to the end of the bitmap.
613    pub fn push_chunk(&mut self, chunk: &[u8; N]) -> &mut Self {
614        for byte in chunk {
615            self.push_byte(*byte);
616        }
617        self
618    }
619
620    /// Pop the last bit from the bitmap.
621    ///
622    /// Returns the value of the popped bit, accounting for any modifications.
623    ///
624    /// # Panics
625    ///
626    /// Panics if the bitmap is empty.
627    pub fn pop(&mut self) -> bool {
628        assert!(self.state.projected_len > 0, "cannot pop from empty bitmap");
629
630        let old_projected_len = self.state.projected_len;
631        self.state.projected_len -= 1;
632        let bit = self.state.projected_len;
633
634        // Determine which region the popped bit came from.
635        // The appended region contains bits pushed: [appended_start, old_projected_len)
636        let appended_start = old_projected_len - self.state.appended_bits.len() as u64;
637
638        if bit >= appended_start {
639            // Popping from appended region: remove from appended_bits vector.
640            self.state.appended_bits.pop().unwrap()
641        } else {
642            // Popping from base region: return the modified value, or the original if unmodified.
643            self.state
644                .modified_bits
645                .remove(&bit)
646                .unwrap_or_else(|| self.current.get_bit(bit))
647        }
648    }
649
650    /// Prune chunks up to the chunk containing the given bit offset.
651    ///
652    /// Note: `bit` can equal `projected_len` when pruning at a chunk boundary.
653    ///
654    /// # Panics
655    ///
656    /// Panics if `bit` is > the projected length.
657    pub fn prune_to_bit(&mut self, bit: u64) -> &mut Self {
658        assert!(
659            bit <= self.state.projected_len,
660            "cannot prune to bit {bit}: beyond projected length ({})",
661            self.state.projected_len
662        );
663
664        let chunk_num = Prunable::<N>::to_chunk_index(bit);
665
666        if chunk_num <= self.state.projected_pruned_chunks {
667            return self; // Already pruned
668        }
669
670        // Capture preimages of chunks being pruned
671        let current_pruned = self.current.pruned_chunks();
672        for chunk_idx in self.state.projected_pruned_chunks..chunk_num {
673            if self.state.chunks_to_prune.contains_key(&chunk_idx) {
674                continue; // Already captured
675            }
676
677            // Invariant: chunk_idx should always be >= current_pruned because
678            // projected_pruned_chunks starts at base_pruned_chunks (= current_pruned)
679            assert!(
680                chunk_idx >= current_pruned,
681                "attempting to prune chunk {chunk_idx} which is already pruned (current pruned_chunks={current_pruned})",
682            );
683
684            // Get chunk data, which may come from dirty state if it's appended
685            let chunk_data = if chunk_idx < self.current.chunks_len() {
686                *self.current.get_chunk(chunk_idx)
687            } else {
688                // Chunk only exists in appended bits
689                // Manually reconstruct it from appended_bits
690                let chunk_start_bit = chunk_idx as u64 * Prunable::<N>::CHUNK_SIZE_BITS;
691                let appended_start =
692                    self.state.projected_len - self.state.appended_bits.len() as u64;
693
694                let mut chunk = [0u8; N];
695                for bit_in_chunk in 0..Prunable::<N>::CHUNK_SIZE_BITS {
696                    let bit = chunk_start_bit + bit_in_chunk;
697                    if bit >= self.state.projected_len {
698                        break;
699                    }
700                    if bit >= appended_start {
701                        let append_idx = (bit - appended_start) as usize;
702                        if append_idx < self.state.appended_bits.len()
703                            && self.state.appended_bits[append_idx]
704                        {
705                            let byte_idx = (bit_in_chunk / 8) as usize;
706                            let bit_idx = bit_in_chunk % 8;
707                            chunk[byte_idx] |= 1u8 << bit_idx;
708                        }
709                    }
710                }
711                chunk
712            };
713
714            self.state.chunks_to_prune.insert(chunk_idx, chunk_data);
715        }
716
717        self.state.projected_pruned_chunks = chunk_num;
718
719        self
720    }
721
722    /// Commit the changes and return a clean bitmap with a historical snapshot.
723    ///
724    /// # Errors
725    ///
726    /// Returns [Error::NonMonotonicCommit] if the commit number is not
727    /// greater than the previous commit.
728    ///
729    /// Returns [Error::ReservedCommitNumber] if the commit number is `u64::MAX`.
730    pub fn commit(mut self, commit_number: u64) -> Result<CleanBitMap<N>, Error> {
731        // Validate commit number is not reserved
732        if commit_number == u64::MAX {
733            return Err(Error::ReservedCommitNumber);
734        }
735
736        // Validate commit number is monotonically increasing
737        if let Some(&max_commit) = self.commits.keys().next_back()
738            && commit_number <= max_commit
739        {
740            return Err(Error::NonMonotonicCommit {
741                previous: max_commit,
742                attempted: commit_number,
743            });
744        }
745
746        // Build reverse diff (captures OLD state before applying changes)
747        let reverse_diff = self.build_reverse_diff();
748
749        // Shrink to length before appends (handles net pops)
750        let target_len_before_appends =
751            self.state.projected_len - self.state.appended_bits.len() as u64;
752        self.current.truncate(target_len_before_appends);
753        // Grow by appending new bits
754        for &bit in &self.state.appended_bits {
755            self.current.push(bit);
756        }
757        assert_eq!(self.current.len(), self.state.projected_len);
758        // Modify existing base bits (not appended bits)
759        for (&bit, &value) in &self.state.modified_bits {
760            self.current.set_bit(bit, value);
761        }
762        // Prune chunks from the beginning
763        if self.state.projected_pruned_chunks > self.state.base_pruned_chunks {
764            let prune_to_bit =
765                self.state.projected_pruned_chunks as u64 * Prunable::<N>::CHUNK_SIZE_BITS;
766            self.current.prune_to_bit(prune_to_bit);
767        }
768
769        // Store the reverse diff
770        self.commits.insert(commit_number, reverse_diff);
771
772        Ok(CleanBitMap {
773            current: self.current,
774            commits: self.commits,
775            state: Clean,
776        })
777    }
778
779    /// Abort the changes and return to clean state.
780    ///
781    /// All pending mutations are discarded.
782    pub fn abort(self) -> CleanBitMap<N> {
783        CleanBitMap {
784            current: self.current,
785            commits: self.commits,
786            state: Clean,
787        }
788    }
789
790    /// Build a reverse diff from current dirty state.
791    fn build_reverse_diff(&self) -> CommitDiff<N> {
792        let mut changes = BTreeMap::new();
793        self.capture_modified_chunks(&mut changes);
794        self.capture_appended_chunks(&mut changes);
795        self.capture_popped_chunks(&mut changes);
796        self.capture_pruned_chunks(&mut changes);
797        CommitDiff {
798            len: self.state.base_len,
799            pruned_chunks: self.state.base_pruned_chunks,
800            chunk_diffs: changes,
801        }
802    }
803
804    /// Capture chunks affected by bit modifications.
805    ///
806    /// For each chunk containing modified bits, we store its original value so we can
807    /// restore it when reconstructing historical states.
808    fn capture_modified_chunks(&self, changes: &mut BTreeMap<usize, ChunkDiff<N>>) {
809        for &bit in self.state.modified_bits.keys() {
810            let chunk_idx = Prunable::<N>::to_chunk_index(bit);
811            changes.entry(chunk_idx).or_insert_with(|| {
812                // `modified_bits` only contains bits from the base region, so the chunk must exist.
813                let old_chunk = self
814                    .get_chunk_from_current(chunk_idx)
815                    .expect("chunk must exist for modified bit");
816                ChunkDiff::Modified(old_chunk)
817            });
818        }
819    }
820
821    /// Capture chunks affected by appended bits.
822    ///
823    /// When bits are appended, they may:
824    /// - Extend an existing partial chunk (mark as Modified with old data)
825    /// - Create entirely new chunks (mark as Added)
826    fn capture_appended_chunks(&self, changes: &mut BTreeMap<usize, ChunkDiff<N>>) {
827        if self.state.appended_bits.is_empty() {
828            return;
829        }
830
831        // Calculate which chunks will be affected by appends.
832        // Note: append_start_bit accounts for any net pops before the pushes.
833        let append_start_bit = self.state.projected_len - self.state.appended_bits.len() as u64;
834        let start_chunk = Prunable::<N>::to_chunk_index(append_start_bit);
835        let end_chunk = Prunable::<N>::to_chunk_index(self.state.projected_len.saturating_sub(1));
836
837        for chunk_idx in start_chunk..=end_chunk {
838            // Use or_insert_with so we don't overwrite chunks already captured
839            // by capture_modified_chunks (which runs first and takes precedence).
840            changes.entry(chunk_idx).or_insert_with(|| {
841                self.get_chunk_from_current(chunk_idx).map_or(
842                    // Chunk is brand new: mark as Added
843                    ChunkDiff::Added,
844                    // Chunk existed before: store its old data
845                    ChunkDiff::Modified,
846                )
847            });
848        }
849    }
850
851    /// Capture chunks affected by pop operations.
852    ///
853    /// When bits are popped (projected_len < base_len), we need to capture the original
854    /// data of chunks that will be truncated or fully removed. This allows reconstruction
855    /// to restore the bits that were popped.
856    fn capture_popped_chunks(&self, changes: &mut BTreeMap<usize, ChunkDiff<N>>) {
857        if self.state.projected_len >= self.state.base_len || self.state.base_len == 0 {
858            return; // No net pops
859        }
860
861        // Identify the range of chunks affected by length reduction.
862        let old_last_chunk = Prunable::<N>::to_chunk_index(self.state.base_len - 1);
863        let new_last_chunk = if self.state.projected_len > 0 {
864            Prunable::<N>::to_chunk_index(self.state.projected_len - 1)
865        } else {
866            0
867        };
868
869        // Capture all chunks between the new and old endpoints.
870
871        // Handle the case where we popped all unpruned bits, leaving new_last_chunk
872        // < self.state.base_pruned_chunks. For example, suppose bitmap has 10 bits per chunk,
873        // and 50 entries, where 40 are pruned. Then we pop 10 bits to make the bitmap have 40 entries,
874        // where all 40 are pruned. Then new_last_chunk is 3 and self.state.base_pruned_chunks is 4.
875        let start_chunk = self.state.base_pruned_chunks.max(new_last_chunk);
876
877        for chunk_idx in start_chunk..=old_last_chunk {
878            changes.entry(chunk_idx).or_insert_with(|| {
879                let old_chunk = self
880                    .get_chunk_from_current(chunk_idx)
881                    .expect("chunk must exist in base bitmap for popped bits");
882
883                // Determine if this chunk is partially kept or completely removed
884                let chunk_start_bit = chunk_idx as u64 * Prunable::<N>::CHUNK_SIZE_BITS;
885
886                if self.state.projected_len > chunk_start_bit {
887                    // Chunk spans the new length boundary -> partially kept (Modified)
888                    ChunkDiff::Modified(old_chunk)
889                } else {
890                    // Chunk is completely beyond the new length -> fully removed (Removed)
891                    ChunkDiff::Removed(old_chunk)
892                }
893            });
894        }
895    }
896
897    /// Capture chunks that will be pruned.
898    ///
899    /// The `prune_to_bit` method already captured the old chunk data,
900    /// so we simply copy it into the reverse diff.
901    fn capture_pruned_chunks(&self, changes: &mut BTreeMap<usize, ChunkDiff<N>>) {
902        for (&chunk_idx, &chunk_data) in &self.state.chunks_to_prune {
903            changes.insert(chunk_idx, ChunkDiff::Pruned(chunk_data));
904        }
905    }
906
907    /// Get chunk data from current state if it exists.
908    ///
909    /// Returns `Some(chunk_data)` if the chunk exists in the current bitmap,
910    /// or `None` if it's out of bounds or pruned.
911    fn get_chunk_from_current(&self, chunk_idx: usize) -> Option<[u8; N]> {
912        let current_pruned = self.current.pruned_chunks();
913        if chunk_idx >= current_pruned && chunk_idx < self.current.chunks_len() {
914            return Some(*self.current.get_chunk(chunk_idx));
915        }
916        None
917    }
918}