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}