ic_stable_roaring/
bitmap.rs

1//! Roaring bitmap types backed by the
2//! [`Memory`](https://docs.rs/ic-stable-structures/latest/ic_stable_structures/trait.Memory.html)
3//! trait.
4//!
5//! Application code should construct values with [`RoaringBitmap::init`]. This module documents
6//! the stable layout; [`RoaringBitmap`] documents runtime behavior and per-method costs.
7//!
8//! # V1 layout
9//!
10//! ```text
11//! ---------------------------------------- <- Address 0
12//! Magic `RSB`                 ↕ 3 bytes
13//! ----------------------------------------
14//! Layout version              ↕ 1 byte
15//! ----------------------------------------
16//! Logical length (`len_bits`) ↕ 8 bytes
17//! ----------------------------------------
18//! Journal slot count          ↕ 8 bytes
19//! ----------------------------------------
20//! Snapshot length             ↕ 8 bytes
21//! ----------------------------------------
22//! Reserved space              ↕ 36 bytes
23//! ---------------------------------------- <- Address 64
24//! Mutation record 0           ↕ 5 bytes
25//! ----------------------------------------
26//! Mutation record 1           ↕ 5 bytes
27//! ----------------------------------------
28//! ...
29//! ----------------------------------------
30//! Mutation record N - 1       ↕ 5 bytes
31//! ---------------------------------------- <- 64 + JOURNAL_CAP_SLOTS * 5
32//! Zero padding                ↕ 0..7 bytes
33//! ---------------------------------------- <- snapshot_base = align_up(journal end, 8)
34//! Serialized Roaring snapshot ↕ variable length
35//! ```
36//!
37//! The header is 64 bytes. Journal records are packed into five bytes, and the snapshot begins at
38//! the next eight-byte boundary. The snapshot uses the standard
39//! [`roaring::RoaringBitmap`](https://docs.rs/roaring/latest/roaring/bitmap/struct.RoaringBitmap.html)
40//! serialization format.
41//!
42//! # Compatibility and recovery
43//!
44//! `JOURNAL_CAP_SLOTS` is stored in the header. A build with a different capacity has a
45//! different journal and snapshot offset, so it cannot reopen existing memory without an
46//! application-level migration; [`RoaringBitmap::init`] returns [`InitError::InvalidLayout`].
47//!
48//! Recovery validates the reachable header and snapshot, then replays journal records until the
49//! first empty record. It deliberately does not inspect unreachable bytes after that point. The
50//! caller must therefore keep the memory region isolated from untrusted writers.
51//!
52//! The header version describes this crate's layout, not the `roaring` crate version. Compatibility
53//! with the supported Roaring serialization format is covered by a checked-in historical fixture.
54
55use crate::journal::{JournalRecord, JournalTag};
56use crate::memory::{
57    MemoryReader, MemoryWriter, grow_memory_to_at_least_bytes, read_bytes, safe_write,
58    write_5_bytes_preallocated, write_u64, write_zero_bytes,
59};
60#[cfg(test)]
61use crate::memory::{read_u64, write_5_bytes};
62use core::cell::{Cell, Ref, RefCell};
63use core::fmt;
64use ic_stable_structures::Memory;
65use roaring::RoaringBitmap as RoaringHeap;
66
67const MAGIC: [u8; 3] = *b"RSB";
68const VERSION: u8 = 1;
69const HEADER_SIZE: u64 = 64;
70const JOURNAL_RECORD_SIZE: u64 = 5;
71
72const MAGIC_OFFSET: u64 = 0;
73const VERSION_OFFSET: u64 = 3;
74const LEN_OFFSET: u64 = 4;
75/// Header field: must equal the build-time journal capacity as `u64` (fixed journal size on disk).
76const JOURNAL_SLOTS_METADATA_OFFSET: u64 = 12;
77const SNAPSHOT_LEN_OFFSET: u64 = 20;
78
79#[derive(Clone, Debug)]
80struct HeapState {
81    len_bits: u64,
82    bitmap: RoaringHeap,
83}
84
85impl HeapState {
86    fn new() -> Self {
87        Self {
88            len_bits: 0,
89            bitmap: RoaringHeap::new(),
90        }
91    }
92}
93
94/// Clears all set bits with index `>= start_exclusive` (indices are `u32`).
95fn remove_suffix_bits(bitmap: &mut RoaringHeap, start_exclusive: u64) {
96    if start_exclusive > u32::MAX as u64 {
97        return;
98    }
99    bitmap.remove_range(start_exclusive as u32..=u32::MAX);
100}
101
102/// Error returned when [`RoaringBitmap::init`] rejects stable memory contents.
103#[derive(Debug, PartialEq, Eq)]
104pub enum InitError {
105    /// The first three bytes were not the expected `RSB` magic (see the [`crate::bitmap`] layout).
106    BadMagic { actual: [u8; 3], expected: [u8; 3] },
107    /// Header layout version is not supported by this build.
108    IncompatibleVersion(u8),
109    /// Catch-all for inconsistent header fields, snapshot length vs. memory size, corrupted
110    /// snapshot bytes, or journal records that fail validation during replay.
111    ///
112    /// This includes a **journal slot count** in the header (offset `12`) that does not equal
113    /// the build-time journal capacity—for example opening stable memory written by a build compiled
114    /// with a different journal capacity.
115    InvalidLayout,
116    /// [`RoaringBitmap::init`] on empty memory calls [`RoaringBitmap::new`]; bootstrap failures there
117    /// (usually [`BitmapError::GrowFailed`]) are returned as this variant.
118    OutOfMemory,
119}
120
121impl fmt::Display for InitError {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        match self {
124            Self::BadMagic { actual, expected } => {
125                write!(f, "bad magic number {actual:?}, expected {expected:?}")
126            }
127            Self::IncompatibleVersion(version) => write!(
128                f,
129                "unsupported layout version {version}; supported version number is {VERSION}"
130            ),
131            Self::InvalidLayout => write!(f, "invalid stable roaring bitmap layout"),
132            Self::OutOfMemory => write!(f, "failed to allocate memory for stable roaring bitmap"),
133        }
134    }
135}
136
137impl std::error::Error for InitError {}
138
139/// Error returned by [`RoaringBitmap::new`] and mutating methods (`set`, `ensure_len`, checkpoint I/O, …).
140#[derive(Debug, PartialEq, Eq)]
141pub enum BitmapError {
142    /// `len` or `index + 1` exceeds the supported logical bit length.
143    LimitsExceeded { value: u64, max: u64 },
144    /// Stable memory could not be grown for a write or checkpoint.
145    GrowFailed(crate::GrowFailed),
146    /// A [`ContainsView`] still borrows the heap mirror, so a state-changing operation cannot run.
147    BorrowConflict,
148    /// Snapshot serialization or stable write failed (`roaring` / [`std::io::Write`] path).
149    Io(String),
150}
151
152impl fmt::Display for BitmapError {
153    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
154        match self {
155            Self::LimitsExceeded { value, max } => write!(
156                f,
157                "value {value} exceeds supported limit {max} (JOURNAL_LEN_MAX; u32 index space)"
158            ),
159            Self::GrowFailed(e) => write!(f, "{e}"),
160            Self::BorrowConflict => write!(f, "bitmap is borrowed by an active ContainsView"),
161            Self::Io(msg) => write!(f, "snapshot I/O: {msg}"),
162        }
163    }
164}
165
166impl std::error::Error for BitmapError {
167    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
168        match self {
169            Self::GrowFailed(e) => Some(e),
170            _ => None,
171        }
172    }
173}
174
175impl From<crate::GrowFailed> for BitmapError {
176    fn from(value: crate::GrowFailed) -> Self {
177        Self::GrowFailed(value)
178    }
179}
180
181impl From<std::io::Error> for BitmapError {
182    fn from(value: std::io::Error) -> Self {
183        Self::Io(value.to_string())
184    }
185}
186
187/// Start offset of the journal region in stable memory.
188fn journal_offset() -> u64 {
189    HEADER_SIZE
190}
191
192/// Byte offset just past the last journal slot: `64 + JOURNAL_CAP_SLOTS * 5` (not necessarily 8-aligned).
193fn journal_end_bytes() -> u64 {
194    crate::JOURNAL_END_BYTES
195}
196
197/// Start of the serialized Roaring snapshot; always 8-byte aligned after zero padding.
198fn snapshot_base() -> u64 {
199    crate::JOURNAL_SNAPSHOT_BASE
200}
201
202/// On-disk header for the current layout.
203///
204/// `journal_slots` (offset `12`, `u64`) records the build-time journal capacity this image was
205/// created with. [`RoaringBitmap::init`] rejects a capacity mismatch,
206/// preventing an upgraded canister with a different capacity from misinterpreting the
207/// journal/snapshot layout.
208#[repr(C)]
209#[derive(Debug)]
210struct Header {
211    magic: [u8; 3],
212    version: u8,
213    len_bits: u64,
214    journal_slots: u64,
215    snapshot_len_bytes: u64,
216}
217
218impl Header {
219    fn new(len_bits: u64, snapshot_len_bytes: u64) -> Self {
220        Self {
221            magic: MAGIC,
222            version: VERSION,
223            len_bits,
224            journal_slots: crate::JOURNAL_CAP_SLOTS as u64,
225            snapshot_len_bytes,
226        }
227    }
228
229    fn write<M: Memory>(&self, memory: &M) -> Result<(), crate::GrowFailed> {
230        safe_write(memory, MAGIC_OFFSET, &self.magic)?;
231        safe_write(memory, VERSION_OFFSET, &[self.version])?;
232        write_u64(memory, LEN_OFFSET, self.len_bits)?;
233        write_u64(memory, JOURNAL_SLOTS_METADATA_OFFSET, self.journal_slots)?;
234        write_u64(memory, SNAPSHOT_LEN_OFFSET, self.snapshot_len_bytes)?;
235        Ok(())
236    }
237}
238
239fn read_header<M: Memory>(memory: &M) -> Result<Header, InitError> {
240    let mut bytes = [0u8; 28];
241    memory.read(MAGIC_OFFSET, &mut bytes);
242    let magic: [u8; 3] = bytes[0..3].try_into().expect("header magic width");
243    let version = bytes[3];
244    if magic != MAGIC {
245        return Err(InitError::BadMagic {
246            actual: magic,
247            expected: MAGIC,
248        });
249    }
250    if version != VERSION {
251        return Err(InitError::IncompatibleVersion(version));
252    }
253    Ok(Header {
254        magic,
255        version,
256        len_bits: u64::from_le_bytes(bytes[4..12].try_into().expect("header length width")),
257        journal_slots: u64::from_le_bytes(bytes[12..20].try_into().expect("header capacity width")),
258        snapshot_len_bytes: u64::from_le_bytes(
259            bytes[20..28].try_into().expect("header snapshot width"),
260        ),
261    })
262}
263
264fn write_header<M: Memory>(
265    memory: &M,
266    len_bits: u64,
267    snapshot_len_bytes: u64,
268) -> Result<(), crate::GrowFailed> {
269    Header::new(len_bits, snapshot_len_bytes).write(memory)
270}
271
272/// Stable roaring bitmap with a heap mirror and a durable journal.
273///
274/// # Documentation split
275///
276/// - **[`crate::bitmap`]**: on-disk layout, packed journal record format, and compatibility rules.
277/// - **`RoaringBitmap` (this type)**: logical length semantics (`len`, out-of-range `contains`),
278///   what is persisted when, [`Self::init`] as the normal entry point (canister code), checkpoint
279///   behavior, and method-level complexity.
280///
281/// # Storage model
282///
283/// Reads use a heap-backed
284/// [`roaring::RoaringBitmap`](https://docs.rs/roaring/latest/roaring/bitmap/struct.RoaringBitmap.html).
285/// Writes append **5-byte** journal records (see [`crate::bitmap`]) and update that mirror. A
286/// serialized roaring snapshot in stable memory starts at an 8-byte aligned offset after the
287/// journal region (with up to 7 bytes of zero padding).
288///
289/// # Checkpointing and amortization
290///
291/// The journal holds a fixed build-time number of records. To ensure a mutation that returns
292/// an error has not been journaled, the implementation checkpoints before an append would consume
293/// the final slot. A full journal from a previous build is checkpointed before any further append.
294/// That **checkpoint** costs **Θ(S)** time and I/O where **S** is the serialized snapshot size in
295/// bytes (and may grow stable memory). Between checkpoints, bit mutations cost **O(1)** amortized
296/// typical roaring work plus **O(1)** journal I/O per state-changing operation.
297///
298/// A few methods call the overflow check even when no journal append occurs, so **rare Θ(S)**
299/// work can still run when opening a legacy full journal (see [`Self::set`]).
300///
301/// # Failure atomicity
302///
303/// On ICP, mutation and checkpoint writes execute synchronously within one message execution. The
304/// platform commits heap and stable-memory changes only on success and rolls them back on a trap or
305/// panic; see [ICP Message Execution Property 5](https://docs.internetcomputer.org/references/message-execution-properties/).
306///
307/// Checkpoint serialization uses multiple [`Memory::write`] calls. This type therefore does not
308/// provide generic process-crash atomicity for a custom [`Memory`] whose individual writes persist
309/// across an interruption. Such an implementation must supply an equivalent rollback or
310/// transactional boundary.
311///
312/// # Concurrency
313///
314/// Interior mutability backs the heap mirror; treat this type as **single-writer**. Do not alias
315/// the same [`Memory`] through another API while an instance is live.
316pub struct RoaringBitmap<M: Memory> {
317    memory: M,
318    state: RefCell<HeapState>,
319    journal_len: Cell<u64>,
320}
321
322/// Borrowed guard over the heap mirror for batched [`Self::contains`] calls.
323///
324/// Obtained from [`RoaringBitmap::contains_view`]. Dropping the view ends the underlying
325/// [`RefCell`] borrow acquired from [`RoaringBitmap`].
326///
327/// While this view is alive, state-changing methods on the same [`RoaringBitmap`] return
328/// [`BitmapError::BorrowConflict`].
329pub struct ContainsView<'a> {
330    state: Ref<'a, HeapState>,
331}
332
333impl ContainsView<'_> {
334    /// Tests membership using the [`RoaringBitmap`] length and heap mirror captured in
335    /// [`RoaringBitmap::contains_view`].
336    ///
337    /// Out-of-range indices yield `false` (same as [`RoaringBitmap::contains`]).
338    ///
339    /// # Time complexity
340    ///
341    /// **O(1)**.
342    #[inline]
343    pub fn contains(&self, index: u32) -> bool {
344        if u64::from(index) >= self.state.len_bits {
345            return false;
346        }
347        self.state.bitmap.contains(index)
348    }
349}
350
351impl<M: Memory> fmt::Debug for RoaringBitmap<M> {
352    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
353        let st = self.state.borrow();
354        f.debug_struct("RoaringBitmap")
355            .field("len_bits", &st.len_bits)
356            .field("cardinality", &st.bitmap.len())
357            .field("journal_len", &self.journal_len.get())
358            .field("journal_cap_slots", &crate::JOURNAL_CAP_SLOTS)
359            .finish()
360    }
361}
362
363impl<M: Memory> RoaringBitmap<M> {
364    /// Writes an empty layout: header, zeroed journal padding up to the aligned snapshot base, and
365    /// an empty heap mirror.
366    ///
367    /// **Canister code should call [`Self::init`] instead**—that is the single supported entry
368    /// point for opening stable memory (including the first time, when `memory.size() == 0`, which
369    /// forwards here). [`Self::new`] remains public for tests and advanced callers that need
370    /// [`InitError`]-free bootstrap errors ([`BitmapError`]).
371    ///
372    /// # Errors
373    ///
374    /// Returns [`BitmapError`] when header setup or padding writes fail.
375    ///
376    /// # Time complexity
377    ///
378    /// **O(1)** with respect to bitmap contents (fixed header and journal layout).
379    pub fn new(memory: M) -> Result<Self, BitmapError> {
380        let snap = snapshot_base();
381        grow_memory_to_at_least_bytes(&memory, snap)?;
382        let journal_end = journal_end_bytes();
383        write_zero_bytes(&memory, journal_offset(), journal_end - journal_offset())?;
384        if journal_end < snap {
385            write_zero_bytes(&memory, journal_end, snap - journal_end)?;
386        }
387        write_header(&memory, 0, 0)?;
388        Ok(Self {
389            memory,
390            state: RefCell::new(HeapState::new()),
391            journal_len: Cell::new(0),
392        })
393    }
394
395    /// **Primary entry point:** open the bitmap from stable memory on every canister entry /
396    /// reload (cold start with no pages yet, upgrade with existing bytes, or steady state).
397    ///
398    /// Validates the header, deserializes the roaring snapshot, and replays journal records up to
399    /// the first all-zero slot (see [`crate`]).
400    ///
401    /// When `memory.size() == 0`, this forwards to [`Self::new`] and maps any [`BitmapError`] to
402    /// [`InitError::OutOfMemory`] (storage bootstrap failure).
403    ///
404    /// # Errors
405    ///
406    /// Returns [`InitError`] when magic/version/journal metadata disagree, lengths are inconsistent
407    /// with the backing memory, the snapshot deserialize fails, or a journal record fails
408    /// validation. In particular, the header **`journal_slots` field must match
409    /// the build-time journal capacity**: otherwise [`InitError::InvalidLayout`] is returned—stable
410    /// memory laid out under a **different compile-time journal capacity cannot be reused** without an
411    /// application-level migration.
412    ///
413    /// # Time complexity
414    ///
415    /// Let **S** be the stored snapshot length in bytes and **K** the number of contiguous journal
416    /// records before the first all-zero slot. Decoding the
417    /// snapshot costs **Θ(S)**. Recovery reads journal chunks through the chunk containing that
418    /// first empty slot. Replaying **K** records costs **Σ** per-record work: typically **O(1)**
419    /// amortized for `SetBit` replay, while a shrinking `SetLen` applies
420    /// `remove_range`/`remove_suffix`-style work **O(C)** over roaring containers intersecting the
421    /// dropped suffix (`C` ≤ number of containers in the map).
422    pub fn init(memory: M) -> Result<Self, InitError> {
423        if memory.size() == 0 {
424            return Self::new(memory).map_err(|_| InitError::OutOfMemory);
425        }
426        let header = read_header(&memory)?;
427        if header.journal_slots != crate::JOURNAL_CAP_SLOTS as u64 {
428            return Err(InitError::InvalidLayout);
429        }
430        let len_bits = header.len_bits;
431        let snapshot_len_bytes = header.snapshot_len_bytes;
432        let need = snapshot_base()
433            .checked_add(snapshot_len_bytes)
434            .ok_or(InitError::InvalidLayout)?;
435        let size_bytes = memory
436            .size()
437            .checked_mul(crate::memory::WASM_PAGE_SIZE)
438            .expect("address overflow");
439        if size_bytes < need {
440            return Err(InitError::InvalidLayout);
441        }
442        if len_bits > crate::JOURNAL_LEN_MAX {
443            return Err(InitError::InvalidLayout);
444        }
445
446        let bitmap = if snapshot_len_bytes == 0 {
447            RoaringHeap::new()
448        } else {
449            let mut reader = MemoryReader::new(&memory, snapshot_base(), snapshot_len_bytes);
450            let bitmap =
451                RoaringHeap::deserialize_from(&mut reader).map_err(|_| InitError::InvalidLayout)?;
452            if !reader.is_exhausted() {
453                return Err(InitError::InvalidLayout);
454            }
455            bitmap
456        };
457        if bitmap
458            .max()
459            .is_some_and(|max_index| u64::from(max_index) >= len_bits)
460        {
461            return Err(InitError::InvalidLayout);
462        }
463        let mut state = HeapState { len_bits, bitmap };
464
465        let mut journal_len = 0u64;
466        let mut chunk_buf = [0u8; crate::JOURNAL_READ_CHUNK_BYTES];
467        let n_chunks = crate::JOURNAL_REGION_BYTES / crate::JOURNAL_READ_CHUNK_BYTES;
468        for chunk_idx in 0..n_chunks {
469            let off = journal_offset() + (chunk_idx * crate::JOURNAL_READ_CHUNK_BYTES) as u64;
470            read_bytes(&memory, off, &mut chunk_buf);
471            for slot in chunk_buf.chunks_exact(JOURNAL_RECORD_SIZE as usize) {
472                let slot: [u8; 5] = slot.try_into().expect("chunks_exact by 5");
473                if slot == [0u8; 5] {
474                    return Ok(Self {
475                        memory,
476                        state: RefCell::new(state),
477                        journal_len: Cell::new(journal_len),
478                    });
479                }
480                apply_record(&mut state, JournalRecord(slot))?;
481                journal_len += 1;
482            }
483        }
484
485        Ok(Self {
486            memory,
487            state: RefCell::new(state),
488            journal_len: Cell::new(journal_len),
489        })
490    }
491
492    /// Consumes `self` and returns the [`Memory`] handle.
493    ///
494    /// # Time complexity
495    ///
496    /// **O(1)**.
497    pub fn into_memory(self) -> M {
498        self.memory
499    }
500
501    /// Returns the **exclusive** logical bit length (`len_bits`): valid indices are `0..len()`.
502    ///
503    /// This is **not** the count of set bits (cardinality).
504    ///
505    /// # Time complexity
506    ///
507    /// **O(1)**.
508    pub fn len(&self) -> u64 {
509        self.state.borrow().len_bits
510    }
511
512    /// Returns `true` when [`Self::len`] is zero.
513    ///
514    /// # Time complexity
515    ///
516    /// **O(1)**.
517    pub fn is_empty(&self) -> bool {
518        self.len() == 0
519    }
520
521    /// Tests whether the bit is set using the heap mirror only (stable memory is not consulted).
522    ///
523    /// Indices `>= len()` yield `false` without extending the bitmap.
524    ///
525    /// # Time complexity
526    ///
527    /// **O(1)**.
528    #[inline]
529    pub fn contains(&self, index: u32) -> bool {
530        let st = self.state.borrow();
531        if u64::from(index) >= st.len_bits {
532            return false;
533        }
534        st.bitmap.contains(index)
535    }
536
537    /// Returns a borrowed view for repeated [`ContainsView::contains`] calls (see [`ContainsView`]).
538    ///
539    /// # Time complexity
540    ///
541    /// **O(1)** (one `RefCell` borrow).
542    #[inline]
543    pub fn contains_view(&self) -> ContainsView<'_> {
544        ContainsView {
545            state: self.state.borrow(),
546        }
547    }
548
549    /// Grows the exclusive logical length to `min_len` if needed, without materializing unset bits.
550    ///
551    /// No-op when `min_len <= len()`.
552    ///
553    /// # Errors
554    ///
555    /// Returns [`BitmapError::LimitsExceeded`] when `min_len` exceeds the supported logical bit length,
556    /// or other [`BitmapError`] variants when journaling or checkpointing fails. On any error, this
557    /// call has not changed the logical bitmap.
558    ///
559    /// # Time complexity
560    ///
561    /// **O(1)** on the no-op path.
562    ///
563    /// When `min_len` is larger, this appends a `SetLen` journal record and may checkpoint; work is
564    /// **O(1)** amortized for steady mutations between checkpoints, with rare **Θ(S)** checkpoints
565    /// for serialized snapshot size **S**. Only growing the length does **not** enumerate unset
566    /// bits in the roaring structure.
567    pub fn ensure_len(&self, min_len: u64) -> Result<(), BitmapError> {
568        if min_len > crate::JOURNAL_LEN_MAX {
569            return Err(BitmapError::LimitsExceeded {
570                value: min_len,
571                max: crate::JOURNAL_LEN_MAX,
572            });
573        }
574        let current = self.len();
575        if min_len <= current {
576            return Ok(());
577        }
578        self.ensure_mutation_not_borrowed()?;
579        self.append_record(JournalRecord::set_len(min_len))?;
580        {
581            let mut st = self.state.borrow_mut();
582            st.len_bits = min_len;
583        }
584        Ok(())
585    }
586
587    /// Sets or clears a bit, journaling only when the logical value changes.
588    ///
589    /// Idempotent: if the bit already equals `value`, no journal record is appended. The
590    /// `value == true` path still performs the journal-full overflow check and may therefore run a
591    /// **Θ(S)** checkpoint even when the bit was already set.
592    ///
593    /// Setting `true` can extend [`Self::len`] to `index + 1` without a separate `SetLen` record.
594    ///
595    /// # Errors
596    ///
597    /// Returns [`BitmapError::LimitsExceeded`] if `index + 1` as `u64` would exceed the supported logical bit length
598    /// (this is unreachable for any `u32` index, but kept for API symmetry), or other [`BitmapError`]
599    /// variants when journaling or checkpointing fails. On any error, this call has not changed the
600    /// logical bitmap.
601    ///
602    /// # Time complexity
603    ///
604    /// **O(1)** to test [`Self::contains`] on the no-append paths.
605    ///
606    /// When the stored bit changes, expect **O(1)** amortized roaring updates and journal I/O, plus
607    /// **Θ(S)** when a checkpoint runs (**S** = serialized snapshot size).
608    pub fn set(&self, index: u32, value: bool) -> Result<(), BitmapError> {
609        let need_len = u64::from(index).saturating_add(1);
610        if need_len > crate::JOURNAL_LEN_MAX {
611            return Err(BitmapError::LimitsExceeded {
612                value: need_len,
613                max: crate::JOURNAL_LEN_MAX,
614            });
615        }
616        if value {
617            if !self.contains(index) {
618                self.ensure_mutation_not_borrowed()?;
619                self.append_record(JournalRecord::set_bit(index, true))?;
620                {
621                    let mut st = self.state.borrow_mut();
622                    if need_len > st.len_bits {
623                        st.len_bits = need_len;
624                    }
625                    st.bitmap.insert(index);
626                }
627                return Ok(());
628            }
629            self.maybe_checkpoint()?;
630            return Ok(());
631        }
632
633        if !self.contains(index) {
634            return Ok(());
635        }
636        self.ensure_mutation_not_borrowed()?;
637        self.append_record(JournalRecord::set_bit(index, false))?;
638        {
639            let mut st = self.state.borrow_mut();
640            st.bitmap.remove(index);
641        }
642        Ok(())
643    }
644
645    /// Equivalent to `self.set(index, true)`. See [`Self::set`] for journaling rules and complexity.
646    pub fn insert(&self, index: u32) -> Result<(), BitmapError> {
647        self.set(index, true)
648    }
649
650    /// Equivalent to `self.set(index, false)`. See [`Self::set`] for journaling rules and complexity.
651    pub fn clear(&self, index: u32) -> Result<(), BitmapError> {
652        self.set(index, false)
653    }
654
655    /// Shrinks the exclusive logical length to `new_len`, clearing set bits at indices `>= new_len`.
656    ///
657    /// No-op when `new_len >= len()`.
658    ///
659    /// # Errors
660    ///
661    /// Returns [`BitmapError::LimitsExceeded`] when `new_len` exceeds the supported logical bit length,
662    /// or other [`BitmapError`] variants when journaling or checkpointing fails. On any error, this
663    /// call has not changed the logical bitmap.
664    ///
665    /// # Time complexity
666    ///
667    /// **O(1)** on the no-op path.
668    ///
669    /// Otherwise **O(C)** to clear the suffix where **C** is the number of roaring containers
670    /// overlapping the removed range, plus journal append and **Θ(S)** checkpoints (**S** =
671    /// serialized snapshot size) when the journal is full.
672    pub fn truncate(&self, new_len: u64) -> Result<(), BitmapError> {
673        if new_len > crate::JOURNAL_LEN_MAX {
674            return Err(BitmapError::LimitsExceeded {
675                value: new_len,
676                max: crate::JOURNAL_LEN_MAX,
677            });
678        }
679        if new_len >= self.len() {
680            return Ok(());
681        }
682        self.ensure_mutation_not_borrowed()?;
683        self.append_record(JournalRecord::set_len(new_len))?;
684        {
685            let mut st = self.state.borrow_mut();
686            st.len_bits = new_len;
687            remove_suffix_bits(&mut st.bitmap, new_len);
688        }
689        Ok(())
690    }
691
692    /// Appends a packed mutation record to the journal.
693    fn append_record(&self, record: JournalRecord) -> Result<(), BitmapError> {
694        let checkpoint_before_append = crate::JOURNAL_CAP_SLOTS as u64 - 1;
695        if self.journal_len.get() >= checkpoint_before_append {
696            self.checkpoint()?;
697        }
698        let idx = self.journal_len.get();
699        let base = journal_offset() + idx * JOURNAL_RECORD_SIZE;
700        write_5_bytes_preallocated(&self.memory, base, &record.0);
701        self.journal_len.set(idx + 1);
702        Ok(())
703    }
704
705    /// Returns an error before journaling when a [`ContainsView`] holds the heap mirror.
706    fn ensure_mutation_not_borrowed(&self) -> Result<(), BitmapError> {
707        if self.state.try_borrow_mut().is_err() {
708            return Err(BitmapError::BorrowConflict);
709        }
710        Ok(())
711    }
712
713    /// Checkpoints a full journal left by an older build or an idempotent operation.
714    fn maybe_checkpoint(&self) -> Result<(), BitmapError> {
715        if self.journal_len.get() >= crate::JOURNAL_CAP_SLOTS as u64 {
716            self.checkpoint()?;
717        }
718        Ok(())
719    }
720
721    /// Writes the heap mirror back into stable memory and clears the journal.
722    fn checkpoint(&self) -> Result<(), BitmapError> {
723        let (len_bits, snapshot_len_bytes) = {
724            let st = self.state.borrow();
725            (st.len_bits, st.bitmap.serialized_size() as u64)
726        };
727        let need_bytes = snapshot_base()
728            .checked_add(snapshot_len_bytes)
729            .ok_or_else(|| BitmapError::Io("address overflow computing snapshot end".into()))?;
730        grow_memory_to_at_least_bytes(&self.memory, need_bytes)?;
731
732        {
733            let st = self.state.borrow();
734            let mut writer = MemoryWriter::new(&self.memory, snapshot_base());
735            st.bitmap.serialize_into(&mut writer)?;
736        }
737
738        write_header(&self.memory, len_bits, snapshot_len_bytes)?;
739        write_zero_bytes(
740            &self.memory,
741            journal_offset(),
742            self.journal_len.get() * JOURNAL_RECORD_SIZE,
743        )?;
744        self.journal_len.set(0);
745        Ok(())
746    }
747
748    #[cfg(feature = "canbench")]
749    pub(crate) fn checkpoint_for_bench(&self) -> Result<(), BitmapError> {
750        self.checkpoint()
751    }
752}
753
754#[inline]
755fn apply_record(state: &mut HeapState, record: JournalRecord) -> Result<(), InitError> {
756    let (tag, value, payload) = record.unpack().map_err(|_| InitError::InvalidLayout)?;
757    match tag {
758        JournalTag::Empty => return Err(InitError::InvalidLayout),
759        JournalTag::SetLen => {
760            let new_len = payload;
761            if new_len > crate::JOURNAL_LEN_MAX || new_len == state.len_bits {
762                return Err(InitError::InvalidLayout);
763            }
764            if new_len < state.len_bits {
765                state.len_bits = new_len;
766                remove_suffix_bits(&mut state.bitmap, new_len);
767            } else {
768                state.len_bits = new_len;
769            }
770        }
771        JournalTag::SetBit => {
772            if payload > u32::MAX as u64 {
773                return Err(InitError::InvalidLayout);
774            }
775            let index = payload as u32;
776            if value {
777                if state.bitmap.contains(index) {
778                    return Err(InitError::InvalidLayout);
779                }
780                let need_len = u64::from(index).saturating_add(1);
781                if need_len > crate::JOURNAL_LEN_MAX {
782                    return Err(InitError::InvalidLayout);
783                }
784                if need_len > state.len_bits {
785                    state.len_bits = need_len;
786                }
787                state.bitmap.insert(index);
788            } else {
789                if u64::from(index) >= state.len_bits || !state.bitmap.contains(index) {
790                    return Err(InitError::InvalidLayout);
791                }
792                state.bitmap.remove(index);
793            }
794        }
795    }
796    Ok(())
797}
798
799#[cfg(test)]
800mod tests {
801    use super::*;
802    use ic_stable_structures::{Memory, vec_mem::VectorMemory};
803    use proptest::prelude::*;
804    use std::collections::BTreeSet;
805    use std::rc::Rc;
806
807    const HISTORICAL_SNAPSHOT_LEN: u64 = 262_295;
808    const OPERATION_DOMAIN: u32 = 255;
809
810    #[derive(Clone)]
811    struct FailOnGrowMemory {
812        inner: VectorMemory,
813        fail_grows: Rc<Cell<bool>>,
814    }
815
816    impl FailOnGrowMemory {
817        fn new() -> Self {
818            Self {
819                inner: VectorMemory::default(),
820                fail_grows: Rc::new(Cell::new(false)),
821            }
822        }
823
824        fn fail_grows(&self) {
825            self.fail_grows.set(true);
826        }
827
828        fn allow_grows(&self) {
829            self.fail_grows.set(false);
830        }
831    }
832
833    impl Memory for FailOnGrowMemory {
834        fn size(&self) -> u64 {
835            self.inner.size()
836        }
837
838        fn grow(&self, pages: u64) -> i64 {
839            if self.fail_grows.get() {
840                -1
841            } else {
842                self.inner.grow(pages)
843            }
844        }
845
846        fn read(&self, offset: u64, dst: &mut [u8]) {
847            self.inner.read(offset, dst);
848        }
849
850        fn write(&self, offset: u64, src: &[u8]) {
851            self.inner.write(offset, src);
852        }
853    }
854
855    /// Test-only memory that records a deep copy after each production `Memory::write` call.
856    /// Recording is opt-in so construction and setup writes do not enter a checkpoint trace.
857    #[derive(Clone)]
858    struct RecordingMemory {
859        inner: VectorMemory,
860        recording: Rc<Cell<bool>>,
861        captures: Rc<RefCell<Vec<VectorMemory>>>,
862    }
863
864    impl RecordingMemory {
865        fn new() -> Self {
866            Self {
867                inner: VectorMemory::default(),
868                recording: Rc::new(Cell::new(false)),
869                captures: Rc::new(RefCell::new(Vec::new())),
870            }
871        }
872
873        fn start_recording(&self) {
874            self.captures.borrow_mut().clear();
875            self.recording.set(true);
876        }
877
878        fn stop_recording(&self) -> Vec<VectorMemory> {
879            self.recording.set(false);
880            core::mem::take(&mut *self.captures.borrow_mut())
881        }
882    }
883
884    impl Memory for RecordingMemory {
885        fn size(&self) -> u64 {
886            self.inner.size()
887        }
888
889        fn grow(&self, pages: u64) -> i64 {
890            self.inner.grow(pages)
891        }
892
893        fn read(&self, offset: u64, dst: &mut [u8]) {
894            self.inner.read(offset, dst);
895        }
896
897        fn write(&self, offset: u64, src: &[u8]) {
898            self.inner.write(offset, src);
899            if self.recording.get() {
900                let bytes = self.inner.borrow().clone();
901                self.captures
902                    .borrow_mut()
903                    .push(Rc::new(RefCell::new(bytes)));
904            }
905        }
906    }
907
908    #[cfg(journal_slots_ge_1024)]
909    fn fill_journal_for_checkpoint_failure(bs: &RoaringBitmap<FailOnGrowMemory>) -> u32 {
910        const CONTAINER_STRIDE: u64 = 1 << 16;
911        let cap = crate::JOURNAL_CAP_SLOTS as u64;
912        let mut next_container = 0u64;
913
914        loop {
915            while bs.journal_len.get() < cap - 1 {
916                let index = next_container * CONTAINER_STRIDE;
917                bs.insert(index as u32).unwrap();
918                next_container += 1;
919            }
920
921            let snapshot_len = bs.state.borrow().bitmap.serialized_size() as u64;
922            let snapshot_end = snapshot_base().checked_add(snapshot_len).unwrap();
923            let allocated_bytes = bs.memory.size() * crate::memory::WASM_PAGE_SIZE;
924            if snapshot_end > allocated_bytes {
925                let index = next_container * CONTAINER_STRIDE;
926                return index as u32;
927            }
928
929            let index = next_container * CONTAINER_STRIDE;
930            bs.insert(index as u32).unwrap();
931            next_container += 1;
932        }
933    }
934
935    fn reopen<M: Memory>(memory: M) -> RoaringBitmap<M> {
936        RoaringBitmap::init(memory).unwrap()
937    }
938
939    /// Decodes the checked-in, whitespace-separated hex/RLE fixture format.
940    ///
941    /// Each token is either an even-length hex byte string or `hh*count`, where `hh` is one byte.
942    /// Keeping the bitmap container as RLE makes the immutable 8 KiB standard-Roaring fixture
943    /// inspectable in source control without asking the current `roaring` writer to reproduce it.
944    fn decode_historical_snapshot() -> Vec<u8> {
945        let mut bytes = Vec::new();
946        for line in include_str!("../tests/fixtures/roaring-0.11.4-mixed.hex").lines() {
947            let token = line.split('#').next().unwrap().trim();
948            if token.is_empty() {
949                continue;
950            }
951            let (hex, repeat) = token
952                .split_once('*')
953                .map_or((token, 1usize), |(hex, repeat)| {
954                    (hex, repeat.parse::<usize>().unwrap())
955                });
956            assert!(
957                hex.len().is_multiple_of(2),
958                "fixture token must contain whole bytes"
959            );
960            let mut decoded = Vec::with_capacity(hex.len() / 2);
961            for pair in hex.as_bytes().chunks_exact(2) {
962                decoded.push(u8::from_str_radix(std::str::from_utf8(pair).unwrap(), 16).unwrap());
963            }
964            for _ in 0..repeat {
965                bytes.extend_from_slice(&decoded);
966            }
967        }
968        bytes
969    }
970
971    fn assert_matches_oracle(
972        bitmap: &RoaringBitmap<VectorMemory>,
973        expected_len: u64,
974        expected_set: &BTreeSet<u32>,
975    ) {
976        assert_eq!(bitmap.len(), expected_len);
977        for index in 0..=OPERATION_DOMAIN {
978            assert_eq!(
979                bitmap.contains(index),
980                expected_set.contains(&index),
981                "index {index}"
982            );
983        }
984    }
985
986    #[derive(Clone, Debug)]
987    enum TestOperation {
988        Insert(u32),
989        Clear(u32),
990        EnsureLen(u64),
991        Truncate(u64),
992        Reopen,
993    }
994
995    fn operation_strategy() -> impl Strategy<Value = TestOperation> {
996        prop_oneof![
997            (0..=OPERATION_DOMAIN).prop_map(TestOperation::Insert),
998            (0..=OPERATION_DOMAIN).prop_map(TestOperation::Clear),
999            (0..=u64::from(OPERATION_DOMAIN) + 1).prop_map(TestOperation::EnsureLen),
1000            (0..=u64::from(OPERATION_DOMAIN) + 1).prop_map(TestOperation::Truncate),
1001            Just(TestOperation::Reopen),
1002        ]
1003    }
1004
1005    #[test]
1006    fn roaring_snapshot_fixture_reopens() {
1007        let snapshot = decode_historical_snapshot();
1008        assert_eq!(snapshot.len(), 8_261);
1009        let memory = VectorMemory::default();
1010        safe_write(&memory, snapshot_base(), &snapshot).unwrap();
1011        write_header(&memory, HISTORICAL_SNAPSHOT_LEN, snapshot.len() as u64).unwrap();
1012
1013        let bitmap = RoaringBitmap::init(memory).unwrap();
1014        assert_eq!(bitmap.len(), HISTORICAL_SNAPSHOT_LEN);
1015        assert_eq!(bitmap.state.borrow().bitmap.len(), 6_009);
1016        for index in [
1017            1,
1018            5,
1019            1_000,
1020            1 << 16,
1021            (1 << 16) | 5_000,
1022            (2 << 16) | 100,
1023            (2 << 16) | 1_000,
1024            (3 << 16) | 7,
1025            (3 << 16) | 700,
1026            (4 << 16) | 50,
1027            (4 << 16) | 150,
1028        ] {
1029            assert!(bitmap.contains(index), "expected fixture bit {index}");
1030        }
1031        for index in [
1032            0,
1033            999,
1034            (1 << 16) | 5_001,
1035            (2 << 16) | 99,
1036            (2 << 16) | 1_001,
1037            (3 << 16) | 8,
1038            (4 << 16) | 49,
1039            (4 << 16) | 151,
1040        ] {
1041            assert!(!bitmap.contains(index), "unexpected fixture bit {index}");
1042        }
1043    }
1044
1045    proptest! {
1046        #![proptest_config(ProptestConfig::with_cases(64))]
1047
1048        #[test]
1049        fn stateful_operations_survive_reopen(operations in prop::collection::vec(operation_strategy(), 1..96)) {
1050            let mut bitmap = RoaringBitmap::new(VectorMemory::default()).unwrap();
1051            let mut expected_len = 0;
1052            let mut expected_set = BTreeSet::new();
1053
1054            for operation in operations {
1055                match operation {
1056                    TestOperation::Insert(index) => {
1057                        bitmap.insert(index).unwrap();
1058                        expected_set.insert(index);
1059                        expected_len = expected_len.max(u64::from(index) + 1);
1060                    }
1061                    TestOperation::Clear(index) => {
1062                        bitmap.clear(index).unwrap();
1063                        expected_set.remove(&index);
1064                    }
1065                    TestOperation::EnsureLen(len) => {
1066                        bitmap.ensure_len(len).unwrap();
1067                        expected_len = expected_len.max(len);
1068                    }
1069                    TestOperation::Truncate(len) => {
1070                        bitmap.truncate(len).unwrap();
1071                        if len < expected_len {
1072                            expected_len = len;
1073                            expected_set.retain(|index| u64::from(*index) < len);
1074                        }
1075                    }
1076                    TestOperation::Reopen => {
1077                        bitmap = reopen(bitmap.into_memory());
1078                    }
1079                }
1080                assert_matches_oracle(&bitmap, expected_len, &expected_set);
1081            }
1082
1083            bitmap = reopen(bitmap.into_memory());
1084            assert_matches_oracle(&bitmap, expected_len, &expected_set);
1085        }
1086    }
1087
1088    #[test]
1089    fn limits_exceeded_returns_error() {
1090        let mem = VectorMemory::default();
1091        let bs = RoaringBitmap::new(mem).unwrap();
1092        assert_eq!(
1093            bs.ensure_len(crate::JOURNAL_LEN_MAX + 1),
1094            Err(BitmapError::LimitsExceeded {
1095                value: crate::JOURNAL_LEN_MAX + 1,
1096                max: crate::JOURNAL_LEN_MAX,
1097            })
1098        );
1099        assert_eq!(
1100            bs.truncate(crate::JOURNAL_LEN_MAX + 1),
1101            Err(BitmapError::LimitsExceeded {
1102                value: crate::JOURNAL_LEN_MAX + 1,
1103                max: crate::JOURNAL_LEN_MAX,
1104            })
1105        );
1106    }
1107
1108    #[test]
1109    fn fresh_create_and_reopen_roundtrip() {
1110        let mem = VectorMemory::default();
1111        let bs = RoaringBitmap::new(mem).unwrap();
1112        assert_eq!(bs.len(), 0);
1113        assert!(bs.is_empty());
1114        let mem = bs.into_memory();
1115        let bs = reopen(mem);
1116        assert_eq!(bs.len(), 0);
1117        assert!(bs.is_empty());
1118    }
1119
1120    #[test]
1121    fn initialization_reports_grow_failure_without_creating_a_layout() {
1122        let memory = FailOnGrowMemory::new();
1123        memory.fail_grows();
1124        assert!(matches!(
1125            RoaringBitmap::new(memory.clone()),
1126            Err(BitmapError::GrowFailed(_))
1127        ));
1128        assert!(matches!(
1129            RoaringBitmap::init(memory),
1130            Err(InitError::OutOfMemory)
1131        ));
1132    }
1133
1134    #[test]
1135    fn journal_append_uses_preallocated_memory() {
1136        let memory = FailOnGrowMemory::new();
1137        let bs = RoaringBitmap::new(memory.clone()).unwrap();
1138        memory.fail_grows();
1139        bs.insert(0).unwrap();
1140        memory.allow_grows();
1141        assert!(bs.contains(0));
1142        assert!(reopen(bs.into_memory()).contains(0));
1143    }
1144
1145    #[test]
1146    fn insert_clear_contains_roundtrip() {
1147        let mem = VectorMemory::default();
1148        let bs = RoaringBitmap::new(mem).unwrap();
1149        bs.insert(0).unwrap();
1150        bs.insert(3).unwrap();
1151        bs.insert(10).unwrap();
1152        bs.clear(3).unwrap();
1153        assert!(bs.contains(0));
1154        assert!(!bs.contains(3));
1155        assert!(bs.contains(10));
1156        assert_eq!(bs.len(), 11);
1157        let mem = bs.into_memory();
1158        let bs = reopen(mem);
1159        assert!(bs.contains(0));
1160        assert!(!bs.contains(3));
1161        assert!(bs.contains(10));
1162        assert_eq!(bs.len(), 11);
1163    }
1164
1165    #[test]
1166    fn ensure_len_preserves_zero_suffix_across_reopen() {
1167        let mem = VectorMemory::default();
1168        let bs = RoaringBitmap::new(mem).unwrap();
1169        bs.ensure_len(16).unwrap();
1170        assert_eq!(bs.len(), 16);
1171        assert!(!bs.contains(0));
1172        assert!(!bs.contains(15));
1173        let mem = bs.into_memory();
1174        let bs = reopen(mem);
1175        assert_eq!(bs.len(), 16);
1176        assert!(!bs.contains(0));
1177        assert!(!bs.contains(15));
1178    }
1179
1180    #[test]
1181    fn truncate_clears_suffix_across_reopen() {
1182        let mem = VectorMemory::default();
1183        let bs = RoaringBitmap::new(mem).unwrap();
1184        bs.insert(1).unwrap();
1185        bs.insert(70).unwrap();
1186        bs.insert(130).unwrap();
1187        bs.truncate(64).unwrap();
1188        assert_eq!(bs.len(), 64);
1189        assert!(bs.contains(1));
1190        assert!(!bs.contains(70));
1191        assert!(!bs.contains(130));
1192        let mem = bs.into_memory();
1193        let bs = reopen(mem);
1194        assert_eq!(bs.len(), 64);
1195        assert!(bs.contains(1));
1196        assert!(!bs.contains(70));
1197        assert!(!bs.contains(130));
1198    }
1199
1200    #[test]
1201    fn checkpoint_after_full_journal_preserves_state() {
1202        let mem = VectorMemory::default();
1203        let bs = RoaringBitmap::new(mem).unwrap();
1204        for i in 0..crate::JOURNAL_CAP_SLOTS {
1205            bs.insert(i as u32).unwrap();
1206        }
1207        let cleared_index = if crate::JOURNAL_CAP_SLOTS > 1 { 1 } else { 0 };
1208        bs.clear(cleared_index as u32).unwrap();
1209        bs.insert(crate::JOURNAL_CAP_SLOTS as u32).unwrap();
1210        assert_eq!(bs.contains(0), cleared_index != 0);
1211        assert!(!bs.contains(cleared_index as u32));
1212        assert!(bs.contains(crate::JOURNAL_CAP_SLOTS as u32));
1213        assert_eq!(bs.len(), (crate::JOURNAL_CAP_SLOTS + 1) as u64);
1214        let mem = bs.into_memory();
1215        let bs = reopen(mem);
1216        assert_eq!(bs.contains(0), cleared_index != 0);
1217        assert!(!bs.contains(cleared_index as u32));
1218        assert!(bs.contains(crate::JOURNAL_CAP_SLOTS as u32));
1219        assert_eq!(bs.len(), (crate::JOURNAL_CAP_SLOTS + 1) as u64);
1220    }
1221
1222    #[derive(Debug)]
1223    enum CheckpointObservation {
1224        Rejected,
1225        Accepted { len: u64, bitmap: RoaringHeap },
1226    }
1227
1228    struct CheckpointTrace {
1229        expected_len: u64,
1230        expected_bitmap: RoaringHeap,
1231        observations: Vec<CheckpointObservation>,
1232    }
1233
1234    fn checkpoint_trace<F>(seed: RoaringHeap, seed_len: u64, mutate: F) -> CheckpointTrace
1235    where
1236        F: FnOnce(&RoaringBitmap<RecordingMemory>),
1237    {
1238        let memory = RecordingMemory::new();
1239        drop(RoaringBitmap::new(memory.clone()).unwrap());
1240        let mut snapshot = Vec::with_capacity(seed.serialized_size());
1241        seed.serialize_into(&mut snapshot).unwrap();
1242        safe_write(&memory, snapshot_base(), &snapshot).unwrap();
1243        write_header(&memory, seed_len, snapshot.len() as u64).unwrap();
1244
1245        let bs = RoaringBitmap::init(memory.clone()).unwrap();
1246        assert_eq!(bs.state.borrow().bitmap, seed, "bad checkpoint seed");
1247        mutate(&bs);
1248        assert!(bs.journal_len.get() > 0, "mutation was not journaled");
1249
1250        let expected_len = bs.len();
1251        let expected_bitmap = bs.state.borrow().bitmap.clone();
1252        memory.start_recording();
1253        bs.checkpoint().unwrap();
1254        let captures = memory.stop_recording();
1255
1256        // More than the five header fields plus journal clear proves the trace includes writes
1257        // issued by the real streaming Roaring serializer.
1258        assert!(
1259            captures.len() > 6,
1260            "checkpoint trace contained only {captures_len} writes",
1261            captures_len = captures.len()
1262        );
1263
1264        let observations = captures
1265            .into_iter()
1266            .map(|captured| match RoaringBitmap::init(captured) {
1267                Ok(reopened) => CheckpointObservation::Accepted {
1268                    len: reopened.len(),
1269                    bitmap: reopened.state.borrow().bitmap.clone(),
1270                },
1271                Err(_) => CheckpointObservation::Rejected,
1272            })
1273            .collect();
1274        CheckpointTrace {
1275            expected_len,
1276            expected_bitmap,
1277            observations,
1278        }
1279    }
1280
1281    fn assert_checkpoint_boundaries<F>(
1282        case: &str,
1283        seed: RoaringHeap,
1284        seed_len: u64,
1285        mutate: F,
1286    ) -> usize
1287    where
1288        F: FnOnce(&RoaringBitmap<RecordingMemory>),
1289    {
1290        let trace = checkpoint_trace(seed, seed_len, mutate);
1291        let mut accepted = 0;
1292        let mut rejected = 0;
1293        for (boundary, observation) in trace.observations.into_iter().enumerate() {
1294            match observation {
1295                CheckpointObservation::Accepted { len, bitmap } => {
1296                    accepted += 1;
1297                    assert_eq!(
1298                        len, trace.expected_len,
1299                        "{case}: write boundary {boundary} recovered a third logical length"
1300                    );
1301                    assert_eq!(
1302                        bitmap, trace.expected_bitmap,
1303                        "{case}: write boundary {boundary} recovered a third bitmap"
1304                    );
1305                }
1306                CheckpointObservation::Rejected => rejected += 1,
1307            }
1308        }
1309        assert!(accepted > 0, "{case}: no recoverable boundary");
1310        rejected
1311    }
1312
1313    #[test]
1314    fn checkpoint_container_transition_boundaries_recover_current_or_reject() {
1315        let first = 1;
1316        let second = (1 << 16) | 10;
1317        let third = (2 << 16) | 100;
1318        let array_seed = RoaringHeap::from_iter([first, second, third]);
1319        let mut rejected = assert_checkpoint_boundaries(
1320            "array-to-array",
1321            array_seed,
1322            u64::from(third) + 1,
1323            |bs| {
1324                bs.clear(first).unwrap();
1325                bs.insert(first + 1).unwrap();
1326                bs.clear(second).unwrap();
1327                bs.insert(second + 1).unwrap();
1328                bs.ensure_len(u64::from(third) + 1_000).unwrap();
1329            },
1330        );
1331
1332        let array_limit_seed = RoaringHeap::from_iter((0..8192).step_by(2));
1333        rejected += assert_checkpoint_boundaries("array-to-bitmap", array_limit_seed, 8191, |bs| {
1334            bs.insert(1).unwrap()
1335        });
1336
1337        let bitmap_seed = RoaringHeap::from_iter(0..4097);
1338        rejected += assert_checkpoint_boundaries("bitmap-to-array", bitmap_seed, 4097, |bs| {
1339            bs.clear(4096).unwrap()
1340        });
1341
1342        let mut run_seed = RoaringHeap::new();
1343        run_seed.insert_range(0..10_000);
1344        let mut run_probe = run_seed.clone();
1345        assert!(
1346            run_probe.remove_run_compression(),
1347            "run seed was not run-compressed"
1348        );
1349        rejected += assert_checkpoint_boundaries("run-to-run", run_seed, 10_000, |bs| {
1350            bs.clear(5_000).unwrap();
1351            bs.insert(12_000).unwrap();
1352        });
1353
1354        assert!(rejected > 0, "suite did not exercise a rejected boundary");
1355    }
1356
1357    #[test]
1358    fn checkpoint_cross_container_splice_recovers_third_state() {
1359        // With a single journal slot, the second mutation checkpoints before this
1360        // streaming checkpoint begins, so the cross-container splice boundary
1361        // exercised below does not exist. `capacity_one_mutation_avoids_double_checkpoint`
1362        // covers the dedicated capacity-one behavior.
1363        if crate::JOURNAL_CAP_SLOTS == 1 {
1364            return;
1365        }
1366
1367        let second_key = 1u32 << 16;
1368        let seed = RoaringHeap::from_iter([1, 2, second_key | 10, second_key | 20]);
1369        let trace = checkpoint_trace(seed, u64::from(second_key | 20) + 1, |bs| {
1370            bs.clear(1).unwrap();
1371            bs.insert(second_key | 30).unwrap();
1372        });
1373
1374        let third_bitmap = RoaringHeap::from_iter([
1375            second_key | 2,
1376            second_key | 10,
1377            second_key | 20,
1378            second_key | 30,
1379        ]);
1380        assert_ne!(third_bitmap, trace.expected_bitmap);
1381        let witness_boundary = trace.observations.iter().position(|observation| {
1382            matches!(
1383                observation,
1384                CheckpointObservation::Accepted { len, bitmap }
1385                    if *len == trace.expected_len && *bitmap == third_bitmap
1386            )
1387        });
1388        // After both new container descriptions are written, the decoder ignores the still-old
1389        // offset table and partitions the old payload using the new cardinalities.
1390        assert_eq!(witness_boundary, Some(5));
1391    }
1392
1393    #[test]
1394    fn capacity_one_mutation_avoids_double_checkpoint() {
1395        if crate::JOURNAL_CAP_SLOTS != 1 {
1396            return;
1397        }
1398
1399        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1400        bs.insert(7).unwrap();
1401        let memory = bs.into_memory();
1402        assert_eq!(
1403            read_u64(&memory, SNAPSHOT_LEN_OFFSET),
1404            RoaringHeap::new().serialized_size() as u64
1405        );
1406        let mut record = [0u8; JOURNAL_RECORD_SIZE as usize];
1407        memory.read(journal_offset(), &mut record);
1408        assert_ne!(record, [0u8; JOURNAL_RECORD_SIZE as usize]);
1409        assert!(reopen(memory).contains(7));
1410    }
1411
1412    #[test]
1413    fn idempotent_insert_checkpoints_a_legacy_full_journal() {
1414        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1415        let memory = bs.into_memory();
1416        for slot in 0..crate::JOURNAL_CAP_SLOTS as u64 {
1417            let record = JournalRecord::set_bit(slot as u32, true);
1418            write_5_bytes(
1419                &memory,
1420                journal_offset() + slot * JOURNAL_RECORD_SIZE,
1421                &record.0,
1422            )
1423            .unwrap();
1424        }
1425
1426        let bs = reopen(memory);
1427        assert_eq!(bs.journal_len.get(), crate::JOURNAL_CAP_SLOTS as u64);
1428        bs.insert(0).unwrap();
1429        assert_eq!(bs.journal_len.get(), 0);
1430        assert!(reopen(bs.into_memory()).contains(0));
1431    }
1432
1433    #[test]
1434    fn mixed_operations_replay_deterministically() {
1435        let mem = VectorMemory::default();
1436        let bs = RoaringBitmap::new(mem).unwrap();
1437        bs.insert(0).unwrap();
1438        bs.insert(9).unwrap();
1439        bs.clear(0).unwrap();
1440        bs.ensure_len(20).unwrap();
1441        bs.insert(19).unwrap();
1442        bs.truncate(12).unwrap();
1443        assert_eq!(bs.len(), 12);
1444        assert!(!bs.contains(0));
1445        assert!(bs.contains(9));
1446        assert!(!bs.contains(19));
1447        let mem = bs.into_memory();
1448        let bs = reopen(mem);
1449        assert_eq!(bs.len(), 12);
1450        assert!(!bs.contains(0));
1451        assert!(bs.contains(9));
1452        assert!(!bs.contains(19));
1453    }
1454
1455    #[test]
1456    fn sparse_high_u32_index_inserts_roundtrip() {
1457        let mem = VectorMemory::default();
1458        let bs = RoaringBitmap::new(mem).unwrap();
1459        let a = (1u32 << 31) + 123;
1460        let b = u32::MAX;
1461        bs.insert(a).unwrap();
1462        bs.insert(b).unwrap();
1463        assert!(bs.contains(a));
1464        assert!(bs.contains(b));
1465        assert_eq!(bs.len(), u32::MAX as u64 + 1);
1466        let mem = bs.into_memory();
1467        let bs = reopen(mem);
1468        assert!(bs.contains(a));
1469        assert!(bs.contains(b));
1470        assert_eq!(bs.len(), u32::MAX as u64 + 1);
1471    }
1472
1473    #[test]
1474    fn new_over_existing_memory_does_not_replay_stale_journal() {
1475        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1476        bs.insert(9).unwrap();
1477        let memory = bs.into_memory();
1478
1479        let bs = RoaringBitmap::new(memory).unwrap();
1480        assert_eq!(bs.len(), 0);
1481        let memory = bs.into_memory();
1482        let bs = reopen(memory);
1483        assert_eq!(bs.len(), 0);
1484        assert!(!bs.contains(9));
1485    }
1486
1487    #[test]
1488    fn init_rejects_snapshot_bits_outside_logical_length() {
1489        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1490        bs.insert(9).unwrap();
1491        bs.checkpoint().unwrap();
1492        let memory = bs.into_memory();
1493        memory.write(LEN_OFFSET, &1u64.to_le_bytes());
1494
1495        assert!(matches!(
1496            RoaringBitmap::init(memory),
1497            Err(InitError::InvalidLayout)
1498        ));
1499    }
1500
1501    #[test]
1502    fn init_rejects_snapshot_with_trailing_bytes() {
1503        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1504        bs.insert(9).unwrap();
1505        bs.checkpoint().unwrap();
1506        let memory = bs.into_memory();
1507        let snapshot_len = read_u64(&memory, SNAPSHOT_LEN_OFFSET);
1508        write_u64(&memory, SNAPSHOT_LEN_OFFSET, snapshot_len + 1).unwrap();
1509
1510        assert!(matches!(
1511            RoaringBitmap::init(memory),
1512            Err(InitError::InvalidLayout)
1513        ));
1514    }
1515
1516    #[test]
1517    fn init_stops_replay_at_first_empty_journal_slot() {
1518        if crate::JOURNAL_CAP_SLOTS < 3 {
1519            return;
1520        }
1521
1522        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1523        let memory = bs.into_memory();
1524        let record = JournalRecord::set_bit(7, true);
1525
1526        write_5_bytes(
1527            &memory,
1528            journal_offset() + JOURNAL_RECORD_SIZE * 2,
1529            &record.0,
1530        )
1531        .unwrap();
1532        let bs = reopen(memory);
1533        assert!(!bs.contains(7));
1534
1535        if crate::JOURNAL_REGION_BYTES > crate::JOURNAL_READ_CHUNK_BYTES {
1536            let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1537            let memory = bs.into_memory();
1538            write_5_bytes(
1539                &memory,
1540                journal_offset() + crate::JOURNAL_READ_CHUNK_BYTES as u64,
1541                &record.0,
1542            )
1543            .unwrap();
1544            let bs = reopen(memory);
1545            assert!(!bs.contains(7));
1546        }
1547    }
1548
1549    #[test]
1550    fn init_rejects_no_op_journal_records() {
1551        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1552        let memory = bs.into_memory();
1553        write_5_bytes(&memory, journal_offset(), &JournalRecord::set_len(0).0).unwrap();
1554        assert!(matches!(
1555            RoaringBitmap::init(memory),
1556            Err(InitError::InvalidLayout)
1557        ));
1558
1559        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1560        let memory = bs.into_memory();
1561        write_5_bytes(
1562            &memory,
1563            journal_offset(),
1564            &JournalRecord::set_bit(0, false).0,
1565        )
1566        .unwrap();
1567        assert!(matches!(
1568            RoaringBitmap::init(memory),
1569            Err(InitError::InvalidLayout)
1570        ));
1571
1572        if crate::JOURNAL_CAP_SLOTS < 2 {
1573            return;
1574        }
1575
1576        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1577        let memory = bs.into_memory();
1578        let set_bit = JournalRecord::set_bit(0, true);
1579        write_5_bytes(&memory, journal_offset(), &set_bit.0).unwrap();
1580        write_5_bytes(&memory, journal_offset() + JOURNAL_RECORD_SIZE, &set_bit.0).unwrap();
1581        assert!(matches!(
1582            RoaringBitmap::init(memory),
1583            Err(InitError::InvalidLayout)
1584        ));
1585    }
1586
1587    #[test]
1588    fn active_contains_view_returns_borrow_conflict_without_journaling() {
1589        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1590        let view = bs.contains_view();
1591        assert!(matches!(bs.ensure_len(1), Err(BitmapError::BorrowConflict)));
1592        drop(view);
1593        assert_eq!(reopen(bs.into_memory()).len(), 0);
1594
1595        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1596        let view = bs.contains_view();
1597        assert!(matches!(bs.insert(0), Err(BitmapError::BorrowConflict)));
1598        drop(view);
1599        assert!(!reopen(bs.into_memory()).contains(0));
1600
1601        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1602        bs.insert(0).unwrap();
1603        let view = bs.contains_view();
1604        assert!(matches!(bs.clear(0), Err(BitmapError::BorrowConflict)));
1605        drop(view);
1606        assert!(reopen(bs.into_memory()).contains(0));
1607
1608        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1609        bs.insert(1).unwrap();
1610        let view = bs.contains_view();
1611        assert!(matches!(bs.truncate(1), Err(BitmapError::BorrowConflict)));
1612        drop(view);
1613        let bs = reopen(bs.into_memory());
1614        assert_eq!(bs.len(), 2);
1615        assert!(bs.contains(1));
1616    }
1617
1618    #[cfg(journal_slots_ge_1024)]
1619    #[test]
1620    fn mutation_error_does_not_apply_journaled_change() {
1621        let memory = FailOnGrowMemory::new();
1622        let bs = RoaringBitmap::new(memory.clone()).unwrap();
1623        let requested_index = fill_journal_for_checkpoint_failure(&bs);
1624        let len_before = bs.len();
1625        memory.fail_grows();
1626        assert!(matches!(
1627            bs.insert(requested_index),
1628            Err(BitmapError::GrowFailed(_))
1629        ));
1630        memory.allow_grows();
1631        assert_eq!(bs.len(), len_before);
1632        assert!(!bs.contains(requested_index));
1633        let bs = reopen(bs.into_memory());
1634        assert_eq!(bs.len(), len_before);
1635        assert!(!bs.contains(requested_index));
1636
1637        let memory = FailOnGrowMemory::new();
1638        let bs = RoaringBitmap::new(memory.clone()).unwrap();
1639        let requested_index = fill_journal_for_checkpoint_failure(&bs);
1640        let existing_index = requested_index - (1u32 << 16);
1641        let len_before = bs.len();
1642        assert!(bs.contains(existing_index));
1643        memory.fail_grows();
1644        assert!(matches!(
1645            bs.clear(existing_index),
1646            Err(BitmapError::GrowFailed(_))
1647        ));
1648        memory.allow_grows();
1649        assert_eq!(bs.len(), len_before);
1650        assert!(bs.contains(existing_index));
1651        let bs = reopen(bs.into_memory());
1652        assert_eq!(bs.len(), len_before);
1653        assert!(bs.contains(existing_index));
1654
1655        let memory = FailOnGrowMemory::new();
1656        let bs = RoaringBitmap::new(memory.clone()).unwrap();
1657        fill_journal_for_checkpoint_failure(&bs);
1658        let len_before = bs.len();
1659        memory.fail_grows();
1660        assert!(matches!(
1661            bs.ensure_len(len_before + 1),
1662            Err(BitmapError::GrowFailed(_))
1663        ));
1664        memory.allow_grows();
1665        assert_eq!(bs.len(), len_before);
1666        let bs = reopen(bs.into_memory());
1667        assert_eq!(bs.len(), len_before);
1668
1669        let memory = FailOnGrowMemory::new();
1670        let bs = RoaringBitmap::new(memory.clone()).unwrap();
1671        fill_journal_for_checkpoint_failure(&bs);
1672        let len_before = bs.len();
1673        let last_index = (len_before - 1) as u32;
1674        memory.fail_grows();
1675        assert!(matches!(
1676            bs.truncate(len_before - 1),
1677            Err(BitmapError::GrowFailed(_))
1678        ));
1679        memory.allow_grows();
1680        assert_eq!(bs.len(), len_before);
1681        assert!(bs.contains(last_index));
1682        let bs = reopen(bs.into_memory());
1683        assert_eq!(bs.len(), len_before);
1684        assert!(bs.contains(last_index));
1685    }
1686
1687    #[test]
1688    fn invalid_magic_version_layout_and_snapshot_are_rejected() {
1689        let mem = VectorMemory::default();
1690        let bs = RoaringBitmap::new(mem).unwrap();
1691        bs.insert(1).unwrap();
1692        let mem = bs.into_memory();
1693        mem.write(0, b"BAD");
1694        assert!(matches!(
1695            RoaringBitmap::init(mem),
1696            Err(InitError::BadMagic { .. })
1697        ));
1698
1699        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1700        let mem = bs.into_memory();
1701        mem.write(VERSION_OFFSET, &[VERSION.wrapping_add(1)]);
1702        assert!(matches!(
1703            RoaringBitmap::init(mem),
1704            Err(InitError::IncompatibleVersion(_))
1705        ));
1706
1707        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1708        let mem = bs.into_memory();
1709        mem.write(JOURNAL_SLOTS_METADATA_OFFSET, &123u64.to_le_bytes());
1710        assert!(matches!(
1711            RoaringBitmap::init(mem),
1712            Err(InitError::InvalidLayout)
1713        ));
1714
1715        let bs = RoaringBitmap::new(VectorMemory::default()).unwrap();
1716        let mem = bs.into_memory();
1717        mem.write(SNAPSHOT_LEN_OFFSET, &10_000_000u64.to_le_bytes());
1718        assert!(matches!(
1719            RoaringBitmap::init(mem),
1720            Err(InitError::InvalidLayout)
1721        ));
1722    }
1723}
ic_stable_roaring/bitmap.rs

ic_stable_roaring/
bitmap.rs
