Skip to main content

commonware_storage/qmdb/current/
batch.rs

1//! Batch mutation API for Current QMDBs.
2//!
3//! Wraps the [`any::batch`] API.
4
5use crate::{
6    Context,
7    index::Unordered as UnorderedIndex,
8    journal::contiguous::{Contiguous, Mutable},
9    merkle::{
10        self, Graftable, Location, Position, Readable,
11        batch::MerkleizedBatch as GenericMerkleizedBatch, mem::Mem,
12        storage::Storage as MerkleStorage,
13    },
14    qmdb::{
15        Error,
16        any::{
17            self, ValueEncoding,
18            batch::{DiffCursors, DiffEntry, Staged as AnyStaged, StagedUpdates},
19            operation::{Operation, update},
20        },
21        batch_chain::Bounds,
22        bitmap::{Shared, fill_from},
23        current::{
24            db::{compute_db_root, partial_chunk, read_graft_inputs},
25            grafting,
26        },
27        operation::Key,
28    },
29};
30use ahash::AHashMap;
31use commonware_codec::Codec;
32use commonware_cryptography::{Digest, Hasher};
33use commonware_parallel::Strategy;
34use commonware_utils::bitmap::{self, Readable as _};
35use core::ops::Range;
36use std::sync::Arc;
37
38/// Speculative chunk-level bitmap overlay.
39///
40/// Instead of tracking individual pushed bits and cleared locations, maintains materialized chunk
41/// bytes for every chunk that differs from the parent bitmap. This directly produces the chunk data
42/// needed for grafted MMR leaf computation.
43#[derive(Clone, Debug, Default)]
44pub(crate) struct ChunkOverlay<const N: usize> {
45    /// Dirty chunks: chunk_idx -> materialized chunk bytes.
46    ///
47    /// Iteration order is not observed by any consumer.
48    pub(crate) chunks: AHashMap<usize, [u8; N]>,
49    /// Total number of bits (parent + new operations).
50    pub(crate) len: u64,
51    /// The parent bitmap's dimensions, captured at construction.
52    parent: Dimensions,
53}
54
55/// Parent-bitmap dimensions captured once per overlay. `chunk_mut` needs them on every newly
56/// materialized chunk, and reading each through a `Base` chain costs a lock acquisition on
57/// the shared committed bitmap, so they are read once instead of per touched chunk.
58#[derive(Clone, Copy, Debug, Default)]
59struct Dimensions {
60    len: u64,
61    complete_chunks: usize,
62    pruned_chunks: usize,
63}
64
65impl Dimensions {
66    fn of<B: bitmap::Readable<N>, const N: usize>(base: &B) -> Self {
67        Self {
68            len: base.len(),
69            complete_chunks: base.complete_chunks(),
70            pruned_chunks: base.pruned_chunks(),
71        }
72    }
73}
74
75impl<const N: usize> ChunkOverlay<N> {
76    const CHUNK_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
77
78    /// Create an overlay of `len` total bits on top of `base`. The `base` handed to later
79    /// `set_bit` / `clear_bit` / `chunk_mut` calls must be the bitmap given here.
80    fn new<B: bitmap::Readable<N>>(base: &B, len: u64, capacity: usize) -> Self {
81        Self {
82            chunks: AHashMap::with_capacity(capacity),
83            len,
84            parent: Dimensions::of(base),
85        }
86    }
87
88    /// Load-or-create a chunk: returns a mutable reference to the materialized chunk bytes. On
89    /// first access for an existing chunk, reads from `base`.
90    fn chunk_mut<B: bitmap::Readable<N>>(&mut self, base: &B, idx: usize) -> &mut [u8; N] {
91        let parent = self.parent;
92        self.chunks.entry(idx).or_insert_with(|| {
93            let base_has_partial = !parent.len.is_multiple_of(Self::CHUNK_BITS);
94            if idx < parent.complete_chunks {
95                base.get_chunk(idx)
96            } else if idx == parent.complete_chunks && base_has_partial {
97                base.last_chunk().0
98            } else {
99                bitmap::BitMap::<N>::EMPTY_CHUNK
100            }
101        })
102    }
103
104    /// Set a single bit (used for pushes and active operations).
105    fn set_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
106        let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
107        let rel = (loc % Self::CHUNK_BITS) as usize;
108        let chunk = self.chunk_mut(base, idx);
109        chunk[rel / 8] |= 1 << (rel % 8);
110    }
111
112    /// Clear a single bit (used for superseded locations). Skips locations in pruned chunks
113    /// since those bits are already inactive.
114    fn clear_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
115        let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
116        if idx < self.parent.pruned_chunks {
117            return;
118        }
119        let rel = (loc % Self::CHUNK_BITS) as usize;
120        let chunk = self.chunk_mut(base, idx);
121        chunk[rel / 8] &= !(1 << (rel % 8));
122    }
123
124    /// Get a dirty chunk's bytes, or `None` if unmodified.
125    pub(crate) fn get(&self, idx: usize) -> Option<&[u8; N]> {
126        self.chunks.get(&idx)
127    }
128
129    /// Number of complete chunks.
130    pub(crate) const fn complete_chunks(&self) -> usize {
131        (self.len / Self::CHUNK_BITS) as usize
132    }
133}
134
135/// Bitmap-accelerated floor scan over a layered `BitmapBatch` chain. Fills `out` with up to
136/// `limit` floor-raise candidates in `[floor, tip)`, returning the next `floor`. Skips
137/// locations where the layered bitmap bit is unset (including locations superseded by
138/// uncommitted ancestors), avoiding I/O reads for inactive operations. Produces the same
139/// sequence as repeatedly calling the `next_candidate` test oracle over the chain.
140///
141/// One scan iterator serves the whole batch: overlay chunks resolve lock-free and the
142/// committed base is locked once per untouched chunk, rather than several times per
143/// candidate. The iterator's chunk caching is sound here because bitmap mutators require
144/// `&mut` on the database, which cannot coexist with the `&db` a merkleize holds.
145pub(crate) fn fill_candidates<F: Graftable, const N: usize>(
146    bitmap: &BitmapBatch<N>,
147    floor: Location<F>,
148    tip: u64,
149    limit: usize,
150    out: &mut Vec<Location<F>>,
151) -> Location<F> {
152    Location::new(fill_from(bitmap, *floor, tip, limit, out))
153}
154
155/// Adapter that resolves ops MMR nodes for a batch's `compute_current_layer`.
156///
157/// Tries the batch chain's sync [`Readable`] first (which covers nodes appended or overwritten
158/// by the batch, plus anything still in the in-memory MMR). Falls through to the base's async
159/// [`MerkleStorage`].
160struct BatchStorageAdapter<
161    'a,
162    F: Graftable,
163    D: Digest,
164    R: Readable<Family = F, Digest = D>,
165    S: MerkleStorage<F, Digest = D>,
166> {
167    batch: &'a R,
168    base: &'a S,
169    _phantom: core::marker::PhantomData<(F, D)>,
170}
171
172impl<
173    'a,
174    F: Graftable,
175    D: Digest,
176    R: Readable<Family = F, Digest = D>,
177    S: MerkleStorage<F, Digest = D>,
178> BatchStorageAdapter<'a, F, D, R, S>
179{
180    const fn new(batch: &'a R, base: &'a S) -> Self {
181        Self {
182            batch,
183            base,
184            _phantom: core::marker::PhantomData,
185        }
186    }
187}
188
189impl<F: Graftable, D: Digest, R: Readable<Family = F, Digest = D>, S: MerkleStorage<F, Digest = D>>
190    MerkleStorage<F> for BatchStorageAdapter<'_, F, D, R, S>
191{
192    type Digest = D;
193
194    fn size(&self) -> Position<F> {
195        self.batch.size()
196    }
197    async fn get_node(&self, pos: Position<F>) -> Result<Option<D>, merkle::Error<F>> {
198        if let Some(node) = self.batch.get_node(pos) {
199            return Ok(Some(node));
200        }
201        self.base.get_node(pos).await
202    }
203
204    async fn get_nodes(&self, positions: &[Position<F>]) -> Result<Vec<D>, merkle::Error<F>> {
205        let mut nodes = vec![None; positions.len()];
206        let mut base_positions = Vec::with_capacity(positions.len());
207
208        // Look up nodes already in the batch chain.
209        for (slot, &pos) in nodes.iter_mut().zip(positions) {
210            match self.batch.get_node(pos) {
211                Some(node) => *slot = Some(node),
212                None => base_positions.push(pos),
213            }
214        }
215
216        // Look up remaining nodes from the base.
217        let base_nodes = if base_positions.is_empty() {
218            Vec::new()
219        } else {
220            self.base.get_nodes(&base_positions).await?
221        };
222        let mut base_nodes = base_nodes.into_iter();
223        Ok(nodes
224            .into_iter()
225            .map(|node| node.unwrap_or_else(|| base_nodes.next().expect("one node per base read")))
226            .collect())
227    }
228}
229
230/// Layers a [`GenericMerkleizedBatch`] over a [`Mem`] for node resolution.
231///
232/// [`GenericMerkleizedBatch::get_node`] only covers the batch chain; committed positions
233/// return `None`. This adapter falls through to the committed Mem for those positions.
234struct BatchOverMem<'a, F: Graftable, D: Digest, S: Strategy> {
235    batch: &'a GenericMerkleizedBatch<F, D, S>,
236    mem: &'a Mem<F, D>,
237}
238
239impl<F: Graftable, D: Digest, S: Strategy> Readable for BatchOverMem<'_, F, D, S> {
240    type Family = F;
241    type Digest = D;
242
243    fn size(&self) -> Position<F> {
244        self.batch.size()
245    }
246
247    fn get_node(&self, pos: Position<F>) -> Option<D> {
248        if let Some(d) = self.batch.get_node(pos) {
249            return Some(d);
250        }
251        self.mem.get_node(pos)
252    }
253}
254
255/// A speculative batch of mutations whose root digest has not yet been computed,
256/// in contrast to [`MerkleizedBatch`].
257///
258/// Wraps a [`any::batch::UnmerkleizedBatch`] and adds bitmap and grafted MMR parent state
259/// needed to compute the current layer during [`merkleize`](Self::merkleize).
260pub struct UnmerkleizedBatch<F, H, U, const N: usize, S: Strategy>
261where
262    F: Graftable,
263    U: update::Update,
264    H: Hasher,
265    Operation<F, U>: Codec,
266{
267    /// The inner any-layer batch that handles mutations, journal, and floor raise.
268    inner: any::batch::UnmerkleizedBatch<F, H, U, S>,
269
270    /// Parent's grafted MMR state.
271    grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
272
273    /// Parent's bitmap state (COW, Arc-based).
274    bitmap_parent: BitmapBatch<N>,
275}
276
277/// Staged batch returned by [`UnmerkleizedBatch::stage`].
278pub struct Staged<F, H, U, const N: usize, S: Strategy>
279where
280    F: Graftable,
281    U: update::Update,
282    H: Hasher,
283    Operation<F, U>: Codec,
284{
285    inner: AnyStaged<F, H, U, S>,
286    grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
287    bitmap_parent: BitmapBatch<N>,
288}
289
290/// A speculative batch of operations whose root digest has been computed, in contrast to
291/// [`UnmerkleizedBatch`].
292///
293/// Wraps an [`any::batch::MerkleizedBatch`] and adds the bitmap and grafted MMR state needed to
294/// compute the canonical root.
295///
296/// # Branch validity
297///
298/// A `MerkleizedBatch` is a branch-scoped view rooted at a specific committed prefix of the DB. It
299/// is not an immutable snapshot.
300///
301/// Internally, the batch chain terminates in the DB's committed bitmap via `BitmapBatch::Base`.
302/// That committed bitmap evolves in place as [`Db::apply_batch`](super::db::Db::apply_batch),
303/// [`Db::prune`](super::db::Db::prune), and [`Db::rewind`](super::db::Db::rewind) update the DB.
304///
305/// Reads through this batch's chain, constructing child batches from it, and applying it later are
306/// only semantically correct while its ancestor chain is still the committed prefix of the DB. In
307/// other words, every successful [`apply_batch`](super::db::Db::apply_batch) since this batch was
308/// merkleized must have applied an ancestor of this batch.
309///
310/// Once a non-ancestor batch is applied, this batch and all of its descendants become invalid
311/// objects. The library does not guard against continued use after that point.
312///
313/// Applying an invalid batch is caught by the any-layer authenticated lineage check and returns
314/// [`Error::StaleBatch`] without mutating committed state, so `apply_batch` itself cannot corrupt
315/// the DB.
316///
317/// Rules of thumb:
318/// - Drop any `Arc<MerkleizedBatch>` you no longer intend to apply.
319/// - Extending a batch after `apply_batch` has consumed it (building a child off the just-applied
320///   parent) is safe. The committed bitmap now equals the parent's post-apply state, so child reads
321///   are consistent.
322/// - Extending a batch after a different branch has been applied is not safe. Do not call `get`,
323///   `new_batch`, or `apply_batch` on that branch again.
324pub struct MerkleizedBatch<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
325{
326    /// Inner any-layer batch (ops MMR, diff, floor, commit loc, sizes).
327    pub(crate) inner: Arc<any::batch::MerkleizedBatch<F, D, U, S>>,
328
329    /// Grafted MMR state.
330    pub(crate) grafted: Arc<merkle::batch::MerkleizedBatch<F, D, S>>,
331
332    /// COW bitmap state (for use as a parent in speculative batches).
333    pub(crate) bitmap: BitmapBatch<N>,
334
335    /// The canonical root (ops root + grafted root + partial chunk).
336    pub(crate) canonical_root: D,
337}
338
339impl<F, H, U, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, U, N, S>
340where
341    F: Graftable,
342    U: update::Update,
343    H: Hasher,
344    Operation<F, U>: Codec,
345{
346    pub(super) const fn new(
347        inner: any::batch::UnmerkleizedBatch<F, H, U, S>,
348        grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
349        bitmap_parent: BitmapBatch<N>,
350    ) -> Self {
351        Self {
352            inner,
353            grafted_parent,
354            bitmap_parent,
355        }
356    }
357
358    /// Record a mutation. Use `Some(value)` for update/create, `None` for delete.
359    ///
360    /// If the same key is written multiple times within a batch, the last
361    /// value wins.
362    pub fn write(mut self, key: U::Key, value: Option<U::Value>) -> Self {
363        self.inner = self.inner.write(key, value);
364        self
365    }
366
367    /// Read through: mutations -> ancestor diffs -> committed DB.
368    pub async fn get<E, C, I>(
369        &self,
370        key: &U::Key,
371        db: &super::db::Db<F, E, C, I, H, U, N, S>,
372    ) -> Result<Option<U::Value>, Error<F>>
373    where
374        E: Context,
375        C: Contiguous<Item = Operation<F, U>>,
376        I: UnorderedIndex<Value = Location<F>> + 'static,
377    {
378        self.inner.get(key, &db.any).await
379    }
380
381    /// Batch read multiple keys.
382    ///
383    /// Returns results in the same order as the input keys. Resolved locations are not retained,
384    /// so writing a key read only through `get_many` requires an index re-probe and journal re-read
385    /// during merkleize. Use [`stage`](Self::stage) for keys that may be written. When the writable
386    /// subset is known and much smaller than the full read set, call `get_many` for the read-only
387    /// keys first, then [`stage`](Self::stage) only the writable keys.
388    pub async fn get_many<E, C, I>(
389        &self,
390        keys: &[&U::Key],
391        db: &super::db::Db<F, E, C, I, H, U, N, S>,
392    ) -> Result<Vec<Option<U::Value>>, Error<F>>
393    where
394        E: Context,
395        C: Contiguous<Item = Operation<F, U>>,
396        I: UnorderedIndex<Value = Location<F>> + 'static,
397    {
398        self.inner.get_many(keys, &db.any).await
399    }
400
401    /// Batch read multiple keys and return a staged batch for the same keys.
402    ///
403    /// Returns results in the same order as the input keys. The staged batch records updates by
404    /// read index: the initial keys occupy `0..keys.len()`, and each [`expand`](Staged::expand)
405    /// appends another index range.
406    pub async fn stage<E, C, I>(
407        self,
408        keys: &[&U::Key],
409        db: &super::db::Db<F, E, C, I, H, U, N, S>,
410    ) -> Result<(Vec<Option<U::Value>>, Staged<F, H, U, N, S>), Error<F>>
411    where
412        E: Context,
413        C: Contiguous<Item = Operation<F, U>>,
414        I: UnorderedIndex<Value = Location<F>> + 'static,
415    {
416        let Self {
417            inner,
418            grafted_parent,
419            bitmap_parent,
420        } = self;
421        let (values, inner) = inner.stage(keys, &db.any).await?;
422        Ok((
423            values,
424            Staged {
425                inner,
426                grafted_parent,
427                bitmap_parent,
428            },
429        ))
430    }
431}
432
433impl<F, H, U, const N: usize, S: Strategy> Staged<F, H, U, N, S>
434where
435    F: Graftable,
436    U: update::Update,
437    H: Hasher,
438    Operation<F, U>: Codec,
439{
440    /// Expand this staged batch with more reads.
441    ///
442    /// Existing read indices remain stable. Newly read keys are appended to the staged read set and
443    /// assigned the returned range. The returned values are in the same order as `keys`.
444    ///
445    /// Expansion does not deduplicate against previously staged keys and does not observe values the
446    /// caller has computed for earlier staged slots but not yet passed to
447    /// [`merkleize`](Staged::merkleize).
448    pub async fn expand<E, C, I>(
449        self,
450        keys: &[&U::Key],
451        db: &super::db::Db<F, E, C, I, H, U, N, S>,
452    ) -> Result<(Range<usize>, Vec<Option<U::Value>>, Self), Error<F>>
453    where
454        E: Context,
455        C: Contiguous<Item = Operation<F, U>>,
456        I: UnorderedIndex<Value = Location<F>> + 'static,
457    {
458        let Self {
459            inner,
460            grafted_parent,
461            bitmap_parent,
462        } = self;
463        let (range, values, inner) = inner.expand(keys, &db.any).await?;
464        Ok((
465            range,
466            values,
467            Self {
468                inner,
469                grafted_parent,
470                bitmap_parent,
471            },
472        ))
473    }
474}
475
476impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Unordered<K, V>, N, S>
477where
478    F: Graftable,
479    K: Key,
480    V: ValueEncoding,
481    H: Hasher,
482    Operation<F, update::Unordered<K, V>>: Codec,
483{
484    /// Record updates for staged reads and upserts for unread keys, then merkleize.
485    ///
486    /// Consumes the staged handle and write vectors. Call [`expand`](Staged::expand) before this
487    /// method if more keys must be read into the staged index space.
488    ///
489    /// A `Some` value is an upsert. `None` is a delete. Update indices refer to the staged read
490    /// set: the initial `stage` input followed by any [`expand`](Staged::expand) ranges. `metadata`
491    /// is committed with the returned batch.
492    ///
493    /// # Panics
494    ///
495    /// Panics if any update's `read_index` is out of the staged read range.
496    #[allow(clippy::type_complexity)]
497    #[tracing::instrument(
498        name = "qmdb.current.unordered.batch.merkleize.staged",
499        level = "info",
500        skip_all,
501        fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
502    )]
503    pub async fn merkleize<E, C, I>(
504        self,
505        updates: Vec<(usize, Option<V::Value>)>,
506        upserts: Vec<(K, Option<V::Value>)>,
507        metadata: Option<V::Value>,
508        db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
509    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
510    where
511        E: Context,
512        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
513        I: UnorderedIndex<Value = Location<F>> + 'static,
514    {
515        let Self {
516            inner,
517            grafted_parent,
518            bitmap_parent,
519        } = self;
520
521        // Overlap the update resolution with a committed-prefix candidate prefetch.
522        // Candidates come from the speculative `bitmap_parent` (the same source the floor
523        // raise scans below), clamped to the committed prefix inside the helper.
524        let (inner, staged_updates, prefetched) = inner
525            .resolve_updates_prefetched(updates, upserts, &db.any, |floor, tip, limit, out| {
526                fill_candidates(&bitmap_parent, floor, tip, limit, out)
527            })
528            .await?;
529        let inner = inner
530            .merkleize_with_floor_scan(
531                &db.any,
532                metadata,
533                staged_updates,
534                Some(prefetched),
535                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
536            )
537            .await?;
538        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
539    }
540}
541
542impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Ordered<K, V>, N, S>
543where
544    F: Graftable,
545    K: Key,
546    V: ValueEncoding,
547    H: Hasher,
548    Operation<F, update::Ordered<K, V>>: Codec,
549{
550    /// Record updates for staged reads and upserts for unread keys, then merkleize.
551    ///
552    /// Consumes the staged handle and write vectors. Call [`expand`](Staged::expand) before this
553    /// method if more keys must be read into the staged index space.
554    ///
555    /// A `Some` value is an upsert. `None` is a delete. Update indices refer to the staged read
556    /// set: the initial `stage` input followed by any [`expand`](Staged::expand) ranges. `metadata`
557    /// is committed with the returned batch.
558    ///
559    /// # Panics
560    ///
561    /// Panics if any update's `read_index` is out of the staged read range.
562    #[allow(clippy::type_complexity)]
563    #[tracing::instrument(
564        name = "qmdb.current.ordered.batch.merkleize.staged",
565        level = "info",
566        skip_all,
567        fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
568    )]
569    pub async fn merkleize<E, C, I>(
570        self,
571        updates: Vec<(usize, Option<V::Value>)>,
572        upserts: Vec<(K, Option<V::Value>)>,
573        metadata: Option<V::Value>,
574        db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
575    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
576    where
577        E: Context,
578        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
579        I: crate::index::Ordered<Value = Location<F>> + 'static,
580    {
581        let Self {
582            inner,
583            grafted_parent,
584            bitmap_parent,
585        } = self;
586        let (inner, staged_updates) = inner.resolve_updates(updates, upserts, db.any.strategy());
587        let inner = inner
588            .merkleize_with_floor_scan(
589                &db.any,
590                metadata,
591                staged_updates,
592                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
593            )
594            .await?;
595        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
596    }
597}
598
599// Unordered merkleize.
600impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
601where
602    F: Graftable,
603    K: Key,
604    V: ValueEncoding,
605    H: Hasher,
606    Operation<F, update::Unordered<K, V>>: Codec,
607{
608    /// Resolve mutations into operations, merkleize, and return an `Arc<MerkleizedBatch>`.
609    #[allow(clippy::type_complexity)]
610    #[tracing::instrument(
611        name = "qmdb.current.unordered.batch.merkleize",
612        level = "info",
613        skip_all
614    )]
615    pub async fn merkleize<E, C, I>(
616        self,
617        db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
618        metadata: Option<V::Value>,
619    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
620    where
621        E: Context,
622        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
623        I: UnorderedIndex<Value = Location<F>> + 'static,
624    {
625        let Self {
626            inner,
627            grafted_parent,
628            bitmap_parent,
629        } = self;
630        // Use the speculative parent bitmap rather than the committed `any` bitmap.
631        let inner = inner
632            .merkleize_with_floor_scan(
633                &db.any,
634                metadata,
635                StagedUpdates::<F, update::Unordered<K, V>>::new(),
636                None,
637                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
638            )
639            .await?;
640        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
641    }
642}
643
644// Ordered merkleize.
645impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
646where
647    F: Graftable,
648    K: Key,
649    V: ValueEncoding,
650    H: Hasher,
651    Operation<F, update::Ordered<K, V>>: Codec,
652{
653    /// Resolve mutations into operations, merkleize, and return an `Arc<MerkleizedBatch>`.
654    #[allow(clippy::type_complexity)]
655    #[tracing::instrument(
656        name = "qmdb.current.ordered.batch.merkleize",
657        level = "info",
658        skip_all
659    )]
660    pub async fn merkleize<E, C, I>(
661        self,
662        db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
663        metadata: Option<V::Value>,
664    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
665    where
666        E: Context,
667        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
668        I: crate::index::Ordered<Value = Location<F>> + 'static,
669    {
670        let Self {
671            inner,
672            grafted_parent,
673            bitmap_parent,
674        } = self;
675        // Use the speculative parent bitmap rather than the committed `any` bitmap.
676        let inner = inner
677            .merkleize_with_floor_scan(
678                &db.any,
679                metadata,
680                StagedUpdates::<F, update::Ordered<K, V>>::new(),
681                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
682            )
683            .await?;
684        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
685    }
686}
687
688/// Derive all bitmap mutations (pushes + clears) for this batch in a single pass over the diff and
689/// ancestor diffs. Avoids iterating raw operations.
690///
691/// Pushes: one bit per operation in the batch. All false except active diff entries (whose `loc`
692/// falls in the batch) and the CommitFloor (last op).
693///
694/// Clears: previous CommitFloor, plus the most recent superseded location for each mutated key. We
695/// search back through ancestors to find the most recent active location; if none exists, we clear
696/// the committed DB location (`base_old_loc`).
697#[allow(clippy::type_complexity)]
698fn build_chunk_overlay<F: Graftable, U, B: bitmap::Readable<N>, const N: usize>(
699    base: &B,
700    batch_len: usize,
701    batch_base: u64,
702    diff: &[(U::Key, DiffEntry<F, U::Value>)],
703    ancestor_diffs: &[Arc<Vec<(U::Key, DiffEntry<F, U::Value>)>>],
704) -> ChunkOverlay<N>
705where
706    U: update::Update,
707{
708    let total_bits = base.len() + batch_len as u64;
709    let appended_chunks = (batch_len as u64).div_ceil(ChunkOverlay::<N>::CHUNK_BITS) as usize;
710    let mut overlay = ChunkOverlay::new(base, total_bits, diff.len() + appended_chunks + 1);
711
712    // 1. CommitFloor (last op) is always active.
713    let commit_loc = batch_base + batch_len as u64 - 1;
714    overlay.set_bit(base, commit_loc);
715
716    // 2. Inactivate previous CommitFloor.
717    overlay.clear_bit(base, batch_base - 1);
718
719    // 3. Set active bits + clear superseded locations from the diff. The diff is key-sorted,
720    // so ancestor resolution streams (one cursor per ancestor diff).
721    let mut ancestors = DiffCursors::new(ancestor_diffs.iter().map(|d| d.as_slice()));
722    for (key, entry) in diff {
723        // Set the active bit for this key's final location.
724        if let Some(loc) = entry.loc()
725            && *loc >= batch_base
726            && *loc < batch_base + batch_len as u64
727        {
728            overlay.set_bit(base, *loc);
729        }
730
731        // Clear the most recent superseded location. Older locations were already cleared by the
732        // ancestor batch that superseded them.
733        let mut prev_loc = entry.base_old_loc();
734        if let Some(ancestor_entry) = ancestors.resolve(key) {
735            prev_loc = ancestor_entry.loc();
736        }
737        if let Some(old) = prev_loc {
738            overlay.clear_bit(base, *old);
739        }
740    }
741
742    // Ensure all new complete chunks beyond the parent are materialized, so downstream consumers
743    // don't read from the parent and panic on out-of-range indices. Uses chunk_mut to inherit the
744    // parent's partial chunk data when idx == parent_complete (avoiding loss of existing bits).
745    let parent_complete = overlay.parent.complete_chunks;
746    let new_complete = overlay.complete_chunks();
747    for idx in parent_complete..new_complete {
748        overlay.chunk_mut(base, idx);
749    }
750
751    overlay
752}
753
754/// Merkleize grafted chunk digests while retaining the live ancestor chain.
755async fn merkleize_grafted_batch<F, H, S, const N: usize>(
756    strategy: &S,
757    grafted_parent: Arc<GenericMerkleizedBatch<F, H::Digest, S>>,
758    grafted_tree: &Arc<Mem<F, H::Digest>>,
759    graft_inputs: Vec<(usize, H::Digest, [u8; N])>,
760    grafting_height: u32,
761) -> Arc<GenericMerkleizedBatch<F, H::Digest, S>>
762where
763    F: Graftable,
764    H: Hasher,
765    S: Strategy,
766{
767    let old_grafted_leaves = *grafted_parent.leaves() as usize;
768    let mut grafted_batch = grafted_parent.new_batch();
769    let ancestors = grafted_batch.retain_ancestors();
770    let grafted_tree = Arc::clone(grafted_tree);
771    strategy
772        .clone()
773        .spawn(graft_inputs.len(), move |strategy| {
774            let new_leaves = grafting::graft_chunk_digests::<H, _, N>(&strategy, graft_inputs);
775            for (chunk_idx, digest) in new_leaves {
776                if chunk_idx < old_grafted_leaves {
777                    grafted_batch = grafted_batch
778                        .update_leaf_digest(Location::<F>::new(chunk_idx as u64), digest)
779                        .expect("update_leaf_digest failed");
780                } else {
781                    grafted_batch = grafted_batch.add_leaf_digest(digest);
782                }
783            }
784            let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
785            let merkleized = grafted_batch.merkleize(&grafted_tree, &grafted_hasher);
786            drop(ancestors);
787            merkleized
788        })
789        .await
790}
791
792/// Compute the current layer (bitmap + grafted MMR + canonical root) on top of a merkleized any
793/// batch.
794///
795/// Builds a chunk overlay from the diff, computes grafted MMR leaves from dirty chunks, and
796/// produces the `Arc<MerkleizedBatch>` directly.
797async fn compute_current_layer<F, E, U, C, I, H, const N: usize, S>(
798    inner: Arc<any::batch::MerkleizedBatch<F, H::Digest, U, S>>,
799    current_db: &super::db::Db<F, E, C, I, H, U, N, S>,
800    grafted_parent: &Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
801    bitmap_parent: &BitmapBatch<N>,
802) -> Result<Arc<MerkleizedBatch<F, H::Digest, U, N, S>>, Error<F>>
803where
804    F: Graftable,
805    E: Context,
806    C: Contiguous<Item = Operation<F, U>>,
807    I: UnorderedIndex<Value = Location<F>>,
808    H: Hasher,
809    U: update::Update,
810    S: Strategy,
811    Operation<F, U>: Codec,
812{
813    let batch_len = inner.journal_batch.items().len();
814    let batch_base = *inner.bounds.tip.size - batch_len as u64;
815
816    // Build chunk overlay: materialized bytes for every dirty chunk.
817    let overlay = build_chunk_overlay::<F, U, _, N>(
818        bitmap_parent,
819        batch_len,
820        batch_base,
821        &inner.diff,
822        &inner.ancestor_diffs,
823    );
824
825    let grafting_height = grafting::height::<N>();
826    let ops_tree_adapter =
827        BatchStorageAdapter::new(&inner.journal_batch, &current_db.any.log.merkle);
828
829    // Snapshot ops_leaves for the post-batch state (the canonical root we're about to compute
830    // sees this many ops). Thread it through `graftable_chunks` derivation and root computation.
831    let overlay_ops_leaves = inner.bounds.tip.size;
832
833    // Distinguish three counters:
834    //   - new_complete_chunks: chunks with all bits filled in the post-batch bitmap
835    //   - graftable_overlay:      chunks committed by the grafted tree (have a single h=G ancestor)
836    //   - graftable_parent:       grafted-tree leaf count from the parent (structural source of truth)
837    //
838    // The pending chunk (if any) sits at index `graftable_overlay` and is excluded from the
839    // grafted tree; its digest is hashed directly into the canonical root.
840    let new_complete_chunks = overlay.complete_chunks();
841    let graftable_overlay = grafting::graftable_chunks::<F>(*overlay_ops_leaves, grafting_height)
842        .min(new_complete_chunks as u64) as usize;
843    let graftable_parent = *grafted_parent.leaves() as usize;
844    let pruned_chunks = bitmap_parent.pruned_chunks();
845    assert!(
846        pruned_chunks <= graftable_parent
847            && graftable_parent <= graftable_overlay
848            && graftable_overlay <= new_complete_chunks,
849        "invariant violated: pruned={pruned_chunks} graftable_parent={graftable_parent} graftable_overlay={graftable_overlay} new_complete={new_complete_chunks}"
850    );
851
852    // Build the set of chunk indices whose grafted-leaf needs (re)computing:
853    //   1) Dirty chunks (bits changed in this batch) within the graftable range.
854    //   2) Pending -> graftable transitions: chunks newly graftable because the ops tree built
855    //      their h=G ancestor in this batch. Their bitmap bytes may not be dirty (the chunk
856    //      became graftable via ops growth alone) but they need a grafted-leaf entry now.
857    let mut chunk_indices_to_update: Vec<usize> = overlay
858        .chunks
859        .iter()
860        .filter(|&(&idx, _)| idx < graftable_overlay && idx >= pruned_chunks)
861        .map(|(&idx, _)| idx)
862        .collect();
863    chunk_indices_to_update.extend(graftable_parent..graftable_overlay);
864    chunk_indices_to_update.sort_unstable();
865    chunk_indices_to_update.dedup();
866    let chunks_to_update = chunk_indices_to_update.into_iter().map(|idx| {
867        let chunk = overlay
868            .get(idx)
869            .copied()
870            .unwrap_or_else(|| bitmap_parent.get_chunk(idx));
871        (idx, chunk)
872    });
873
874    // Prefetch each chunk's covering ops-tree node, then run graft hashing and the grafted
875    // MMR build/merkleize as one job through the strategy (against a snapshot of the
876    // committed grafted tree). An empty graft set hashes nothing, so it merkleizes without
877    // submitting a job.
878    let graft_inputs = read_graft_inputs::<F, _, N>(&ops_tree_adapter, chunks_to_update).await?;
879    let grafted_batch = if graft_inputs.is_empty() {
880        let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
881        grafted_parent
882            .new_batch()
883            .merkleize(&current_db.grafted_tree, &grafted_hasher)
884    } else {
885        merkleize_grafted_batch::<F, H, S, N>(
886            &current_db.strategy,
887            Arc::clone(grafted_parent),
888            &current_db.grafted_tree,
889            graft_inputs,
890            grafting_height,
891        )
892        .await
893    };
894
895    // Build the layered bitmap (parent + overlay) before computing the canonical root, so that
896    // compute_db_root sees newly completed chunks. Using bitmap_parent alone would miss chunks
897    // that transitioned from partial to complete in this batch.
898    let bitmap_batch = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
899        parent: bitmap_parent.clone(),
900        overlay: Arc::new(overlay),
901        shared: Arc::clone(bitmap_parent.shared()),
902    }));
903
904    // Compute canonical root. The grafted batch alone cannot resolve committed nodes,
905    // so layer it over the committed grafted MMR.
906    let ops_root = inner.root();
907    let layered = BatchOverMem {
908        batch: &grafted_batch,
909        mem: &current_db.grafted_tree,
910    };
911    let grafted_storage =
912        grafting::Storage::<F, H, _, _>::new(&layered, grafting_height, &ops_tree_adapter);
913    // Compute partial chunk (last incomplete chunk, if any). The partial chunk lives at
914    // index `new_complete_chunks` (the chunk currently being filled with bits) -- distinct
915    // from `graftable_overlay` (the grafted-tree boundary). At gh >= 3, partial and pending can
916    // coexist; this branch only handles partial. The pending chunk (when present) is read
917    // from the bitmap inside `compute_db_root` via `pending_chunk()`.
918    let partial = partial_chunk::<_, N>(&bitmap_batch);
919    let canonical_root = compute_db_root::<F, H, _, _, N>(
920        &bitmap_batch,
921        &grafted_storage,
922        overlay_ops_leaves,
923        partial,
924        inner.bounds.inactivity_floor,
925        &ops_root,
926    )
927    .await?;
928
929    Ok(Arc::new(MerkleizedBatch {
930        inner,
931        grafted: grafted_batch,
932        bitmap: bitmap_batch,
933        canonical_root,
934    }))
935}
936
937/// A view of the committed bitmap plus zero or more speculative overlay `Layer`s.
938///
939/// The chain terminates in a `Base` that references the shared committed bitmap. No validity
940/// check is performed. Callers must ensure they only read through batches whose chains are
941/// still valid prefixes of committed state (see [`Shared`]'s docs).
942#[derive(Clone, Debug)]
943pub(crate) enum BitmapBatch<const N: usize> {
944    /// Chain terminal: shared reference to the committed bitmap.
945    Base(Arc<Shared<N>>),
946    /// Speculative layer on top of a parent batch.
947    Layer(Arc<BitmapBatchLayer<N>>),
948}
949
950/// The data behind a [`BitmapBatch::Layer`].
951#[derive(Debug)]
952pub(crate) struct BitmapBatchLayer<const N: usize> {
953    pub(crate) parent: BitmapBatch<N>,
954    /// Chunk-level overlay: materialized bytes for every chunk that differs from parent.
955    pub(crate) overlay: Arc<ChunkOverlay<N>>,
956    /// Cached terminal [`Shared`] so [`BitmapBatch::shared`] and
957    /// [`BitmapBatch::pruned_chunks`] answer in O(1) instead of walking the chain.
958    pub(crate) shared: Arc<Shared<N>>,
959}
960
961impl<const N: usize> BitmapBatch<N> {
962    const CHUNK_SIZE_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
963
964    /// Return the terminal [`Shared`] at the bottom of the chain.
965    fn shared(&self) -> &Arc<Shared<N>> {
966        match self {
967            Self::Base(s) => s,
968            Self::Layer(layer) => &layer.shared,
969        }
970    }
971
972    /// Return a chain equivalent to `self` with any `Layer` whose overlay is now fully committed
973    /// replaced by a direct reference to the committed bitmap. Since `apply_batch` commits
974    /// contiguous prefixes, committed `Layer`s are always at the bottom of the chain.
975    fn trim_committed(&self) -> Self {
976        let shared = self.shared();
977        let committed = bitmap::Readable::<N>::len(shared.as_ref());
978        let mut kept = Vec::new();
979        let mut current = self;
980        while let Self::Layer(layer) = current {
981            if layer.overlay.len <= committed {
982                break;
983            }
984            kept.push(Arc::clone(&layer.overlay));
985            current = &layer.parent;
986        }
987        let mut result = Self::Base(Arc::clone(shared));
988        for overlay in kept.into_iter().rev() {
989            result = Self::Layer(Arc::new(BitmapBatchLayer {
990                parent: result,
991                overlay,
992                shared: Arc::clone(shared),
993            }));
994        }
995        result
996    }
997}
998
999impl<const N: usize> bitmap::Readable<N> for BitmapBatch<N> {
1000    fn complete_chunks(&self) -> usize {
1001        (self.len() / Self::CHUNK_SIZE_BITS) as usize
1002    }
1003
1004    fn get_chunk(&self, idx: usize) -> [u8; N] {
1005        // Walk the layer chain. Each layer's overlay either holds the chunk (return it) or
1006        // doesn't (descend).
1007        let mut current = self;
1008        loop {
1009            match current {
1010                Self::Base(shared) => return shared.get_chunk(idx),
1011                Self::Layer(layer) => {
1012                    if let Some(&chunk) = layer.overlay.get(idx) {
1013                        return chunk;
1014                    }
1015                    current = &layer.parent;
1016                }
1017            }
1018        }
1019    }
1020
1021    fn last_chunk(&self) -> ([u8; N], u64) {
1022        let total = self.len();
1023        if total == 0 {
1024            return (bitmap::BitMap::<N>::EMPTY_CHUNK, 0);
1025        }
1026        let rem = total % Self::CHUNK_SIZE_BITS;
1027        let bits_in_last = if rem == 0 { Self::CHUNK_SIZE_BITS } else { rem };
1028        let idx = if rem == 0 {
1029            self.complete_chunks().saturating_sub(1)
1030        } else {
1031            self.complete_chunks()
1032        };
1033        (self.get_chunk(idx), bits_in_last)
1034    }
1035
1036    fn pruned_chunks(&self) -> usize {
1037        self.shared().pruned_chunks()
1038    }
1039
1040    fn len(&self) -> u64 {
1041        match self {
1042            Self::Base(shared) => bitmap::Readable::<N>::len(shared.as_ref()),
1043            Self::Layer(layer) => layer.overlay.len,
1044        }
1045    }
1046}
1047
1048impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
1049    MerkleizedBatch<F, D, U, N, S>
1050{
1051    /// Return the canonical root.
1052    pub const fn root(&self) -> D {
1053        self.canonical_root
1054    }
1055
1056    /// Return the QMDB ops-only root.
1057    pub fn ops_root(&self) -> D {
1058        self.inner.root()
1059    }
1060
1061    /// Return the [`Bounds`] of the batch.
1062    pub fn bounds(&self) -> &Bounds<F, D> {
1063        self.inner.bounds()
1064    }
1065
1066    /// Return the operations this batch appends to the ops log and the location of the first.
1067    ///
1068    /// Delegates to the wrapped ops-level batch. The bitmap state contributes to the
1069    /// canonical root but appends no log operations. There are no matching proof or
1070    /// pinned-node methods: an ops-level proof verifies only against [`Self::ops_root`],
1071    /// never the grafted [`Self::root`] that blocks commit to.
1072    pub fn operations(&self) -> (Location<F>, Arc<Vec<Operation<F, U>>>) {
1073        self.inner.operations()
1074    }
1075
1076    /// Return the batch's safe sync boundary.
1077    ///
1078    /// This equals the boundary [`super::db::Db::sync_boundary`] reports once this batch is applied.
1079    pub fn sync_boundary(&self) -> Location<F> {
1080        // Derive from the commit's chunk-aligned inactivity floor, the same quantity the DB uses
1081        // after apply. Deliberately not the physical bitmap pruning boundary, which can lag the
1082        // inactivity floor when pruning has not run.
1083        super::db::sync_boundary::<F, N>(
1084            *self.inner.bounds().inactivity_floor / bitmap::Prunable::<N>::CHUNK_SIZE_BITS,
1085            *self.inner.bounds().tip.size,
1086        )
1087    }
1088}
1089
1090impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
1091    MerkleizedBatch<F, D, U, N, S>
1092where
1093    Operation<F, U>: Codec,
1094{
1095    /// Create a new speculative batch of operations with this batch as its parent.
1096    ///
1097    /// All uncommitted ancestors in the chain must be kept alive until the child (or any
1098    /// descendant) is merkleized. Dropping an uncommitted ancestor causes data
1099    /// loss detected at `apply_batch` time.
1100    ///
1101    /// This is only valid while `self` is still on the winning branch. If a different branch has
1102    /// been applied since `self` was created, `self` is no longer a valid parent and must not be
1103    /// extended.
1104    pub fn new_batch<H>(self: &Arc<Self>) -> UnmerkleizedBatch<F, H, U, N, S>
1105    where
1106        H: Hasher<Digest = D>,
1107    {
1108        UnmerkleizedBatch::new(
1109            self.inner.new_batch::<H>(),
1110            Arc::clone(&self.grafted),
1111            self.bitmap.trim_committed(),
1112        )
1113    }
1114
1115    /// Read through: local diff -> ancestor diffs -> committed DB.
1116    ///
1117    /// This is only valid while `self` remains on the committed prefix. If a non-ancestor batch
1118    /// has been applied since `self` was merkleized, do not read through it.
1119    pub async fn get<E, C, I, H>(
1120        &self,
1121        key: &U::Key,
1122        db: &super::db::Db<F, E, C, I, H, U, N, S>,
1123    ) -> Result<Option<U::Value>, Error<F>>
1124    where
1125        E: Context,
1126        C: Contiguous<Item = Operation<F, U>>,
1127        I: UnorderedIndex<Value = Location<F>> + 'static,
1128        H: Hasher<Digest = D>,
1129    {
1130        self.inner.get(key, &db.any).await
1131    }
1132
1133    /// Batch read multiple keys.
1134    ///
1135    /// Returns results in the same order as the input keys.
1136    pub async fn get_many<E, C, I, H>(
1137        &self,
1138        keys: &[&U::Key],
1139        db: &super::db::Db<F, E, C, I, H, U, N, S>,
1140    ) -> Result<Vec<Option<U::Value>>, Error<F>>
1141    where
1142        E: Context,
1143        C: Contiguous<Item = Operation<F, U>>,
1144        I: UnorderedIndex<Value = Location<F>> + 'static,
1145        H: Hasher<Digest = D>,
1146    {
1147        self.inner.get_many(keys, &db.any).await
1148    }
1149}
1150
1151impl<F, E, C, I, H, U, const N: usize, S> super::db::Db<F, E, C, I, H, U, N, S>
1152where
1153    F: Graftable,
1154    E: Context,
1155    C: Contiguous<Item = Operation<F, U>>,
1156    I: UnorderedIndex<Value = Location<F>>,
1157    H: Hasher,
1158    U: update::Update,
1159    S: Strategy,
1160    Operation<F, U>: Codec,
1161{
1162    /// Create an initial [`MerkleizedBatch`] from the current committed DB state.
1163    ///
1164    /// The returned batch is rooted at the current committed prefix, but it is not a persistent
1165    /// snapshot across later divergent commits. If some other branch is applied afterward, this
1166    /// batch is no longer valid and must not be read through, extended, or applied.
1167    pub fn to_batch(&self) -> Arc<MerkleizedBatch<F, H::Digest, U, N, S>> {
1168        let grafted = self.grafted_snapshot();
1169        Arc::new(MerkleizedBatch {
1170            inner: self.any.to_batch(),
1171            grafted,
1172            bitmap: BitmapBatch::Base(Arc::clone(&self.any.bitmap)),
1173            canonical_root: self.root,
1174        })
1175    }
1176}
1177
1178#[cfg(any(test, feature = "test-traits"))]
1179mod trait_impls {
1180    use super::*;
1181    use crate::{
1182        journal::contiguous::Mutable,
1183        qmdb::any::traits::{
1184            ApplyBatchResult, BatchableDb, MerkleizedBatch as MerkleizedBatchTrait,
1185            UnmerkleizedBatch as UnmerkleizedBatchTrait,
1186        },
1187    };
1188    use std::future::Future;
1189
1190    type CurrentDb<F, E, C, I, H, U, const N: usize, S> =
1191        crate::qmdb::current::db::Db<F, E, C, I, H, U, N, S>;
1192
1193    impl<F, K, V, H, E, C, I, const N: usize, S>
1194        UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>>
1195        for UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
1196    where
1197        F: Graftable,
1198        K: Key,
1199        V: ValueEncoding + 'static,
1200        H: Hasher,
1201        E: Context,
1202        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
1203        I: UnorderedIndex<Value = Location<F>> + 'static,
1204        S: Strategy,
1205        Operation<F, update::Unordered<K, V>>: Codec,
1206    {
1207        type Family = F;
1208        type K = K;
1209        type V = V::Value;
1210        type Metadata = V::Value;
1211        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;
1212
1213        fn write(self, key: K, value: Option<V::Value>) -> Self {
1214            Self::write(self, key, value)
1215        }
1216
1217        async fn merkleize(
1218            self,
1219            db: &CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>,
1220            metadata: Option<V::Value>,
1221        ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
1222            self.merkleize(db, metadata).await
1223        }
1224    }
1225
1226    impl<F, K, V, H, E, C, I, const N: usize, S>
1227        UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>>
1228        for UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
1229    where
1230        F: Graftable,
1231        K: Key,
1232        V: ValueEncoding + 'static,
1233        H: Hasher,
1234        E: Context,
1235        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
1236        I: crate::index::Ordered<Value = Location<F>> + 'static,
1237        S: Strategy,
1238        Operation<F, update::Ordered<K, V>>: Codec,
1239    {
1240        type Family = F;
1241        type K = K;
1242        type V = V::Value;
1243        type Metadata = V::Value;
1244        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;
1245
1246        fn write(self, key: K, value: Option<V::Value>) -> Self {
1247            Self::write(self, key, value)
1248        }
1249
1250        async fn merkleize(
1251            self,
1252            db: &CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>,
1253            metadata: Option<V::Value>,
1254        ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
1255            self.merkleize(db, metadata).await
1256        }
1257    }
1258
1259    impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
1260        MerkleizedBatchTrait for Arc<MerkleizedBatch<F, D, U, N, S>>
1261    where
1262        Operation<F, U>: Codec,
1263    {
1264        type Digest = D;
1265
1266        fn root(&self) -> D {
1267            MerkleizedBatch::root(self)
1268        }
1269    }
1270
1271    impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
1272        for CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>
1273    where
1274        F: Graftable,
1275        E: Context,
1276        K: Key,
1277        V: ValueEncoding + 'static,
1278        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
1279        I: UnorderedIndex<Value = Location<F>> + 'static,
1280        H: Hasher,
1281        S: Strategy,
1282        Operation<F, update::Unordered<K, V>>: Codec,
1283    {
1284        type Family = F;
1285        type K = K;
1286        type V = V::Value;
1287        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;
1288        type Batch = UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>;
1289
1290        fn new_batch(&self) -> Self::Batch {
1291            self.new_batch()
1292        }
1293
1294        fn apply_batch(
1295            self,
1296            batch: Self::Merkleized,
1297        ) -> impl Future<Output = ApplyBatchResult<Self>> {
1298            self.apply_batch(batch)
1299        }
1300    }
1301
1302    impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
1303        for CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>
1304    where
1305        F: Graftable,
1306        E: Context,
1307        K: Key,
1308        V: ValueEncoding + 'static,
1309        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
1310        I: crate::index::Ordered<Value = Location<F>> + 'static,
1311        H: Hasher,
1312        S: Strategy,
1313        Operation<F, update::Ordered<K, V>>: Codec,
1314    {
1315        type Family = F;
1316        type K = K;
1317        type V = V::Value;
1318        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;
1319        type Batch = UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>;
1320
1321        fn new_batch(&self) -> Self::Batch {
1322            self.new_batch()
1323        }
1324
1325        fn apply_batch(
1326            self,
1327            batch: Self::Merkleized,
1328        ) -> impl Future<Output = ApplyBatchResult<Self>> {
1329            self.apply_batch(batch)
1330        }
1331    }
1332}
1333
1334#[cfg(test)]
1335mod tests {
1336    use super::*;
1337    use crate::{mmb, mmr, utils::detached::block_strategy};
1338    use commonware_cryptography::Sha256;
1339    use commonware_macros::test_traced;
1340    use commonware_parallel::{Manual, Rayon};
1341    use commonware_utils::{NZUsize, bitmap::Prunable as BitMap};
1342    use std::{
1343        future::Future as _,
1344        task::Context as TaskContext,
1345        time::{Duration, Instant},
1346    };
1347
1348    // N=4 -> CHUNK_SIZE_BITS = 32
1349    const N: usize = 4;
1350    type Bm = BitMap<N>;
1351    type GraftedBatch =
1352        Arc<GenericMerkleizedBatch<mmb::Family, <Sha256 as Hasher>::Digest, Manual<Rayon>>>;
1353    type Location = mmr::Location;
1354
1355    fn make_bitmap(bits: &[bool]) -> Bm {
1356        let mut bm = Bm::new();
1357        for &b in bits {
1358            bm.push(b);
1359        }
1360        bm
1361    }
1362
1363    fn grafted_chain(
1364        strategy: &Manual<Rayon>,
1365        mem: &Arc<Mem<mmb::Family, <Sha256 as Hasher>::Digest>>,
1366    ) -> (GraftedBatch, GraftedBatch) {
1367        let hasher = grafting::hasher::<mmb::Family, Sha256>(grafting::height::<1>());
1368        let a = mem
1369            .new_batch_with_strategy(strategy.clone())
1370            .add_leaf_digest(Sha256::hash(&[b"a-0"]))
1371            .add_leaf_digest(Sha256::hash(&[b"a-1"]))
1372            .merkleize(mem, &hasher);
1373        let b = a
1374            .new_batch()
1375            .add_leaf_digest(Sha256::hash(&[b"b-0"]))
1376            .merkleize(mem, &hasher);
1377        (a, b)
1378    }
1379
1380    /// A detached grafted-tree merkleization owns the full ancestor chain after cancellation.
1381    #[test_traced]
1382    fn test_grafted_merkleize_retains_ancestors_after_cancellation() {
1383        let strategy = Rayon::new(NZUsize!(2)).unwrap();
1384        let manual = strategy.manual();
1385        let mem = Arc::new(Mem::<mmb::Family, <Sha256 as Hasher>::Digest>::new());
1386        let grafting_height = grafting::height::<1>();
1387        let graft_inputs = || vec![(0, Sha256::hash(&[b"replacement"]), [1u8; 1])];
1388        let waker = futures::task::noop_waker();
1389        let mut context = TaskContext::from_waker(&waker);
1390
1391        // Observe the worker result so a missing grandparent fails the test directly.
1392        let (a, b) = grafted_chain(&manual, &mem);
1393        let ancestor = Arc::downgrade(&a);
1394        let release = block_strategy(&strategy, 2);
1395        let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
1396            &manual,
1397            Arc::clone(&b),
1398            &mem,
1399            graft_inputs(),
1400            grafting_height,
1401        ));
1402        assert!(merkleize.as_mut().poll(&mut context).is_pending());
1403        drop(b);
1404        drop(a);
1405        drop(release);
1406        let _ = futures::executor::block_on(merkleize);
1407        assert!(ancestor.upgrade().is_none());
1408
1409        // Drop the waiter while the worker is queued to prove the guard moved with it.
1410        let (a, b) = grafted_chain(&manual, &mem);
1411        let ancestor = Arc::downgrade(&a);
1412        let release = block_strategy(&strategy, 2);
1413        let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
1414            &manual,
1415            Arc::clone(&b),
1416            &mem,
1417            graft_inputs(),
1418            grafting_height,
1419        ));
1420        assert!(merkleize.as_mut().poll(&mut context).is_pending());
1421        drop(merkleize);
1422        drop(b);
1423        drop(a);
1424        assert!(ancestor.upgrade().is_some());
1425
1426        drop(release);
1427        let deadline = Instant::now() + Duration::from_secs(10);
1428        while ancestor.upgrade().is_some() {
1429            assert!(
1430                Instant::now() < deadline,
1431                "detached grafted merkleization did not release its ancestors"
1432            );
1433            std::thread::yield_now();
1434        }
1435    }
1436
1437    // ---- build_chunk_overlay tests ----
1438
1439    #[test]
1440    fn chunk_overlay_pushes() {
1441        use crate::qmdb::any::value::FixedEncoding;
1442        use commonware_utils::sequence::FixedBytes;
1443
1444        type K = FixedBytes<4>;
1445        type V = FixedEncoding<u64>;
1446        type U = crate::qmdb::any::operation::update::Unordered<K, V>;
1447
1448        let key1 = FixedBytes::from([1, 0, 0, 0]);
1449        let key2 = FixedBytes::from([2, 0, 0, 0]);
1450
1451        // Base: 4 bits, all set (previous commit at loc 3).
1452        // Segment of 4 operations starting at base_size=4.
1453        // Diff: key1 active at loc=4 (in batch), key2 active at loc=99 (not in batch,
1454        // so superseded within this batch).
1455        let base = make_bitmap(&[true; 4]);
1456        let mut diff = vec![
1457            (
1458                key1,
1459                DiffEntry::Active {
1460                    value: 100u64,
1461                    loc: Location::new(4), // offset 0 in batch
1462                    base_old_loc: None,
1463                },
1464            ),
1465            (
1466                key2,
1467                DiffEntry::Active {
1468                    value: 200u64,
1469                    loc: Location::new(99), // not in batch [4,8), so superseded
1470                    base_old_loc: None,
1471                },
1472            ),
1473        ];
1474        diff.sort_by(|a, b| a.0.cmp(&b.0));
1475
1476        let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 4, 4, &diff, &[]);
1477
1478        // Chunk 0 should have: bits 0-3 from base (all set), bit 4 set (key1), bits 5-6 false
1479        // (inactive), bit 7 set (CommitFloor at loc 7). Also bit 3 cleared (previous commit).
1480        let c0 = overlay.get(0).expect("chunk 0 should be dirty");
1481        assert_ne!(c0[0] & (1 << 4), 0); // key1 active
1482        assert_eq!(c0[0] & (1 << 5), 0); // inactive
1483        assert_eq!(c0[0] & (1 << 6), 0); // inactive
1484        assert_ne!(c0[0] & (1 << 7), 0); // CommitFloor
1485        assert_eq!(c0[0] & (1 << 3), 0); // previous commit cleared
1486    }
1487
1488    #[test]
1489    fn chunk_overlay_clears() {
1490        use crate::qmdb::any::value::FixedEncoding;
1491        use commonware_utils::sequence::FixedBytes;
1492
1493        type K = FixedBytes<4>;
1494        type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;
1495
1496        let key1 = FixedBytes::from([1, 0, 0, 0]);
1497        let key2 = FixedBytes::from([2, 0, 0, 0]);
1498        let key3 = FixedBytes::from([3, 0, 0, 0]);
1499
1500        // Base bitmap with 64 bits, all set.
1501        let base = make_bitmap(&[true; 64]);
1502
1503        let mut diff: Vec<(K, DiffEntry<mmr::Family, u64>)> = vec![
1504            (
1505                key1,
1506                DiffEntry::Active {
1507                    value: 100,
1508                    loc: Location::new(70),
1509                    base_old_loc: Some(Location::new(5)),
1510                },
1511            ),
1512            (
1513                key2,
1514                DiffEntry::Deleted {
1515                    base_old_loc: Some(Location::new(10)),
1516                },
1517            ),
1518            (
1519                key3,
1520                DiffEntry::Active {
1521                    value: 300,
1522                    loc: Location::new(71),
1523                    base_old_loc: None,
1524                },
1525            ),
1526        ];
1527        diff.sort_by(|a, b| a.0.cmp(&b.0));
1528
1529        // Segment of 8 ops starting at 64; previous commit at loc 63.
1530        let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 8, 64, &diff, &[]);
1531
1532        // Verify bits 5 and 10 are cleared in chunk 0.
1533        let c0 = overlay.get(0).expect("chunk 0 should be dirty");
1534        assert_eq!(c0[0] & (1 << 5), 0); // bit 5 cleared
1535        assert_eq!(c0[1] & (1 << 2), 0); // bit 10 = byte 1, bit 2 cleared
1536
1537        // Other bits should still be set.
1538        assert_eq!(c0[0] & (1 << 4), 1 << 4); // bit 4 still set
1539        assert_eq!(c0[1] & (1 << 3), 1 << 3); // bit 11 still set
1540    }
1541
1542    /// Regression: when the parent bitmap has a partial last chunk that becomes complete in the
1543    /// child (without any active bits landing in that chunk), the overlay must inherit the parent's
1544    /// partial chunk data, not zero it out.
1545    #[test]
1546    fn chunk_overlay_preserves_partial_parent_chunk() {
1547        use crate::qmdb::any::value::FixedEncoding;
1548        use commonware_utils::sequence::FixedBytes;
1549
1550        type K = FixedBytes<4>;
1551        type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;
1552
1553        // Base: 20 bits set (partial chunk 0, CHUNK_SIZE_BITS=32).
1554        let base = make_bitmap(&[true; 20]);
1555        assert_eq!(base.complete_chunks(), 0); // partial
1556
1557        // Segment of 20 ops starting at loc 20. This pushes total to 40 bits, completing chunk 0
1558        // (32 bits) and starting chunk 1. Diff: only one active key at loc 35 (in chunk 1), plus
1559        // CommitFloor at loc 39. No active bits land in chunk 0's new region (bits 20-31).
1560        let key1 = FixedBytes::from([1, 0, 0, 0]);
1561        let mut diff = vec![(
1562            key1,
1563            DiffEntry::Active {
1564                value: 42u64,
1565                loc: Location::new(35),
1566                base_old_loc: None,
1567            },
1568        )];
1569        diff.sort_by(|a, b| a.0.cmp(&b.0));
1570
1571        let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 20, 20, &diff, &[]);
1572
1573        // Chunk 0 should be materialized and preserve the parent's first 20 bits.
1574        let c0 = overlay.get(0).expect("chunk 0 should be in overlay");
1575        // Bits 0-7 all set -> byte 0 = 0xFF
1576        assert_eq!(c0[0], 0xFF);
1577        // Bits 8-15 all set -> byte 1 = 0xFF
1578        assert_eq!(c0[1], 0xFF);
1579        // Bits 16-18 set, bit 19 cleared (previous commit), 20-23 not set -> byte 2 = 0x07
1580        assert_eq!(c0[2], 0x07);
1581    }
1582
1583    // ---- next_candidate tests ----
1584
1585    /// Single-step oracle for [`fill_candidates`]: return the next floor-raise candidate in
1586    /// `[floor, tip)` over any [`bitmap::Readable`]. `fill_candidates_matches_oracle` proves
1587    /// the production scan produces exactly this sequence over every chain shape.
1588    fn next_candidate<B: bitmap::Readable<N2>, const N2: usize>(
1589        bitmap: &B,
1590        floor: Location,
1591        tip: u64,
1592    ) -> Option<Location> {
1593        let floor = *floor;
1594        let bitmap_len = bitmap.len();
1595        let committed_end = bitmap_len.min(tip);
1596        if floor < committed_end
1597            && let Some(idx) = bitmap.ones_iter_from(floor).next()
1598            && idx < committed_end
1599        {
1600            return Some(Location::new(idx));
1601        }
1602        let candidate = floor.max(bitmap_len);
1603        (candidate < tip).then(|| Location::new(candidate))
1604    }
1605
1606    #[test]
1607    fn bitmap_scan_all_active() {
1608        let bm = make_bitmap(&[true; 8]);
1609        for i in 0..8 {
1610            assert_eq!(
1611                next_candidate(&bm, Location::new(i), 8),
1612                Some(Location::new(i))
1613            );
1614        }
1615        assert_eq!(next_candidate(&bm, Location::new(8), 8), None);
1616    }
1617
1618    #[test]
1619    fn bitmap_scan_all_inactive() {
1620        let bm = make_bitmap(&[false; 8]);
1621        assert_eq!(next_candidate(&bm, Location::new(0), 8), None);
1622    }
1623
1624    #[test]
1625    fn bitmap_scan_skips_inactive() {
1626        // Pattern: inactive, inactive, active, inactive, active
1627        let bm = make_bitmap(&[false, false, true, false, true]);
1628        assert_eq!(
1629            next_candidate(&bm, Location::new(0), 5),
1630            Some(Location::new(2))
1631        );
1632        assert_eq!(
1633            next_candidate(&bm, Location::new(3), 5),
1634            Some(Location::new(4))
1635        );
1636        assert_eq!(next_candidate(&bm, Location::new(5), 5), None);
1637    }
1638
1639    #[test]
1640    fn bitmap_scan_beyond_bitmap_len_returns_candidate() {
1641        // Bitmap has 4 bits, but tip is 8. Locations 4..8 are beyond the bitmap and should be
1642        // returned as candidates.
1643        let bm = make_bitmap(&[false; 4]);
1644        // All bitmap bits are unset, so 0..4 are skipped; loc 4 is beyond bitmap -> candidate.
1645        assert_eq!(
1646            next_candidate(&bm, Location::new(0), 8),
1647            Some(Location::new(4))
1648        );
1649        assert_eq!(
1650            next_candidate(&bm, Location::new(6), 8),
1651            Some(Location::new(6))
1652        );
1653    }
1654
1655    #[test]
1656    fn bitmap_scan_respects_tip() {
1657        let bm = make_bitmap(&[false, false, false, true]);
1658        // Active bit at 3, but tip is 3 so it's excluded.
1659        assert_eq!(next_candidate(&bm, Location::new(0), 3), None);
1660        // With tip=4, bit 3 is included.
1661        assert_eq!(
1662            next_candidate(&bm, Location::new(0), 4),
1663            Some(Location::new(3))
1664        );
1665    }
1666
1667    #[test]
1668    fn bitmap_scan_floor_at_tip() {
1669        let bm = make_bitmap(&[true; 4]);
1670        assert_eq!(next_candidate(&bm, Location::new(4), 4), None);
1671    }
1672
1673    #[test]
1674    fn bitmap_scan_empty_bitmap() {
1675        let bm = Bm::new();
1676        // Empty bitmap, but tip > 0: all locations are beyond bitmap.
1677        assert_eq!(
1678            next_candidate(&bm, Location::new(0), 5),
1679            Some(Location::new(0))
1680        );
1681        // Empty bitmap, tip = 0: no candidates.
1682        assert_eq!(next_candidate(&bm, Location::new(0), 0), None);
1683    }
1684
1685    #[test]
1686    fn fill_candidates_matches_oracle() {
1687        // Sequence parity plus split-resume for one (chain, tip): the scan matches
1688        // single-stepping the oracle over the same chain, and any split point resumes
1689        // seamlessly via the returned continuation.
1690        fn assert_matches(name: &str, chain: &BitmapBatch<N>, tip: u64) {
1691            for floor in 0..=tip {
1692                let mut want = Vec::new();
1693                let mut scan = Location::new(floor);
1694                while let Some(c) = next_candidate(chain, scan, tip) {
1695                    want.push(c);
1696                    scan = c + 1;
1697                }
1698                for split in 0..=want.len() {
1699                    let mut got = Vec::new();
1700                    let next = fill_candidates(chain, Location::new(floor), tip, split, &mut got);
1701                    fill_candidates(chain, next, tip, want.len() + 1, &mut got);
1702                    assert_eq!(got, want, "{name} floor={floor} split={split}");
1703                }
1704            }
1705        }
1706
1707        let bits = [true, false, true, true, false, false, true, false];
1708        let base = make_bitmap(&bits);
1709
1710        // Flat committed base.
1711        let flat = BitmapBatch::Base(Arc::new(Shared::new(make_bitmap(&bits))));
1712
1713        // One layer: clears committed bits 3 and 6, appends 8..12 (only 9 set).
1714        let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1715        let mut overlay = ChunkOverlay::new(&base, 12, 1);
1716        overlay.clear_bit(&base, 3);
1717        overlay.clear_bit(&base, 6);
1718        overlay.set_bit(&base, 9);
1719        let one_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1720            parent: BitmapBatch::Base(Arc::clone(&shared)),
1721            overlay: Arc::new(overlay),
1722            shared,
1723        }));
1724
1725        // Two layers, mirroring `fill_candidates_filters_ancestor_clears`.
1726        let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1727        let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
1728        overlay1.clear_bit(&base, 3);
1729        overlay1.set_bit(&base, 9);
1730        let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1731            parent: BitmapBatch::Base(Arc::clone(&shared)),
1732            overlay: Arc::new(overlay1),
1733            shared: Arc::clone(&shared),
1734        }));
1735        let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
1736        overlay2.clear_bit(&chain1, 6);
1737        overlay2.clear_bit(&chain1, 9);
1738        overlay2.set_bit(&chain1, 13);
1739        let two_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1740            parent: chain1,
1741            overlay: Arc::new(overlay2),
1742            shared,
1743        }));
1744
1745        // Pruned base: 40 bits with chunk 0 pruned (33 and 38 set beyond the pruned
1746        // boundary), plus a layer clearing 38 and appending 40..46 (41 and 44 set).
1747        let make_pruned = || {
1748            let mut bits = [false; 40];
1749            bits[33] = true;
1750            bits[38] = true;
1751            let mut bm = make_bitmap(&bits);
1752            bm.prune_to_bit(32);
1753            bm
1754        };
1755        let pruned_base = make_pruned();
1756        let shared = Arc::new(Shared::new(make_pruned()));
1757        let mut overlay = ChunkOverlay::new(&pruned_base, 46, 1);
1758        overlay.clear_bit(&pruned_base, 38);
1759        overlay.set_bit(&pruned_base, 41);
1760        overlay.set_bit(&pruned_base, 44);
1761        let pruned = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1762            parent: BitmapBatch::Base(Arc::clone(&shared)),
1763            overlay: Arc::new(overlay),
1764            shared,
1765        }));
1766
1767        for (name, chain, committed) in [
1768            ("flat", flat, 8),
1769            ("one-layer", one_layer, 8),
1770            ("two-layer", two_layer, 8),
1771            ("pruned-base", pruned, 40),
1772        ] {
1773            let len = bitmap::Readable::<N>::len(&chain);
1774            for tip in [committed, len, len + 3] {
1775                assert_matches(name, &chain, tip);
1776            }
1777
1778            // Prefetch-then-live handoff: the prefetch is clamped to the committed
1779            // boundary and the live scan resumes from the continuation with the
1780            // post-batch tip. Nothing the raise must revalidate may be lost across the
1781            // handoff (false negatives are forbidden): every set bit in `[floor, len)`
1782            // and every location in `[len, tip)`.
1783            let tip = len + 3;
1784            let cap = tip as usize;
1785            let pruned_bits = bitmap::Readable::<N>::pruned_bits(&chain);
1786            for floor in pruned_bits..=committed {
1787                let mut got = Vec::new();
1788                let next = fill_candidates(&chain, Location::new(floor), committed, cap, &mut got);
1789                fill_candidates(&chain, next, tip, cap, &mut got);
1790                assert!(got.is_sorted_by(|a, b| a < b), "{name} floor={floor}");
1791                for loc in floor..tip {
1792                    let must_emit = loc >= len || bitmap::Readable::<N>::get_bit(&chain, loc);
1793                    assert!(
1794                        !must_emit || got.contains(&Location::new(loc)),
1795                        "{name} floor={floor} lost {loc}"
1796                    );
1797                }
1798            }
1799        }
1800    }
1801
1802    #[test]
1803    fn fill_candidates_filters_ancestor_clears() {
1804        let bits = [true, false, true, true, false, false, true, false];
1805        let base = make_bitmap(&bits);
1806        let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1807
1808        // Layer 1 clears committed bit 3 and appends bits 8..12 (only 9 set).
1809        let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
1810        overlay1.clear_bit(&base, 3);
1811        overlay1.set_bit(&base, 9);
1812        let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1813            parent: BitmapBatch::Base(Arc::clone(&shared)),
1814            overlay: Arc::new(overlay1),
1815            shared: Arc::clone(&shared),
1816        }));
1817
1818        // Layer 2 (materialized against the layer-1 view, as `build_chunk_overlay` does)
1819        // clears bits 6 and 9, and appends bits 12..14 (only 13 set).
1820        let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
1821        overlay2.clear_bit(&chain1, 6);
1822        overlay2.clear_bit(&chain1, 9);
1823        overlay2.set_bit(&chain1, 13);
1824        let chain2 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1825            parent: chain1.clone(),
1826            overlay: Arc::new(overlay2),
1827            shared,
1828        }));
1829
1830        // Bits cleared by any layer are skipped (no wasted log reads), set bits -- committed
1831        // or appended, from whichever layer materialized the chunk last -- are emitted
1832        // ascending, and locations at or beyond the layered length up to `tip` are emitted
1833        // sequentially.
1834        let scan = |chain: &BitmapBatch<N>, tip: u64| {
1835            let mut got = Vec::new();
1836            fill_candidates(chain, Location::new(0), tip, 16, &mut got);
1837            got
1838        };
1839        let want = |locs: &[u64]| locs.iter().copied().map(Location::new).collect::<Vec<_>>();
1840        assert_eq!(scan(&chain1, 12), want(&[0, 2, 6, 9]));
1841        assert_eq!(scan(&chain2, 14), want(&[0, 2, 13]));
1842        assert_eq!(scan(&chain2, 16), want(&[0, 2, 13, 14, 15]));
1843    }
1844
1845    #[test]
1846    fn fill_candidates_mixes_overlay_and_base_chunks() {
1847        // Base spans two chunks (N=4 -> 32-bit chunks): full chunk 0 plus a partial chunk 1.
1848        let mut bits = [false; 40];
1849        for i in [1, 30, 33, 35, 38] {
1850            bits[i] = true;
1851        }
1852        let base = make_bitmap(&bits);
1853        let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1854
1855        // Layer touches only chunk 1: clears committed bit 35 and appends bits 40..44
1856        // (only 41 set). Chunk 0 stays unmaterialized, so the scan must fall through to
1857        // the committed base there.
1858        let mut overlay = ChunkOverlay::new(&base, 44, 1);
1859        overlay.clear_bit(&base, 35);
1860        overlay.set_bit(&base, 41);
1861        let chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1862            parent: BitmapBatch::Base(Arc::clone(&shared)),
1863            overlay: Arc::new(overlay),
1864            shared,
1865        }));
1866
1867        // Chunk 0 bits come from the base, chunk 1 bits from the overlay (35 filtered,
1868        // the appended 41 emitted).
1869        let mut got = Vec::new();
1870        fill_candidates(&chain, Location::new(0), 44, 16, &mut got);
1871        let want: Vec<Location> = [1, 30, 33, 38, 41].into_iter().map(Location::new).collect();
1872        assert_eq!(got, want);
1873    }
1874
1875    // ---- trim_committed tests ----
1876    //
1877    // `trim_committed` is called from `MerkleizedBatch::new_batch` to strip any `Layer`s whose
1878    // overlays have already been absorbed into the shared committed bitmap by a prior apply.
1879    // The implementation is a single loop that collects uncommitted overlays top-down and
1880    // rebuilds a fresh chain rooted at `Base`. These tests cover distinct input shapes directly,
1881    // without going through the full Db/batch machinery, so the function's structural output
1882    // can be asserted.
1883
1884    /// Build a chain `Base(shared) -> Layer(len=L1) -> Layer(len=L2) -> ...` from a list of
1885    /// overlay lengths (bottom to top). Each constructed `Layer` caches `shared` per the
1886    /// struct's invariant.
1887    fn make_chain(shared: &Arc<Shared<N>>, overlay_lens: &[u64]) -> BitmapBatch<N> {
1888        let mut chain = BitmapBatch::Base(Arc::clone(shared));
1889        for &len in overlay_lens {
1890            let overlay = Arc::new(ChunkOverlay::new(&chain, len, 0));
1891            chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1892                parent: chain,
1893                overlay,
1894                shared: Arc::clone(shared),
1895            }));
1896        }
1897        chain
1898    }
1899
1900    /// Walk a chain and return its overlay lengths in bottom-to-top order. Used to assert the
1901    /// structural output of `trim_committed` without touching private fields. Panics if the
1902    /// chain isn't terminated by a single `Base` at the bottom.
1903    fn chain_overlays(batch: &BitmapBatch<N>) -> Vec<u64> {
1904        let mut lens = Vec::new();
1905        let mut current = batch;
1906        while let BitmapBatch::Layer(layer) = current {
1907            lens.push(layer.overlay.len);
1908            current = &layer.parent;
1909        }
1910        assert!(matches!(current, BitmapBatch::Base(_)));
1911        lens.reverse();
1912        lens
1913    }
1914
1915    /// Input is already a bare `Base` with no speculative layers on top -- the loop body never
1916    /// runs, `kept` stays empty, and the result is a freshly constructed `Base` pointing at the
1917    /// same `Shared`. Real-world trigger: `MerkleizedBatch::new_batch` on a batch whose
1918    /// chain was previously trimmed flat (e.g., immediately after an apply collapsed everything).
1919    #[test]
1920    fn trim_committed_already_base() {
1921        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1922        let base = BitmapBatch::Base(Arc::clone(&shared));
1923        let result = base.trim_committed();
1924        // Still `Base`, pointing at the same shared terminal.
1925        match result {
1926            BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
1927            BitmapBatch::Layer(_) => panic!("expected Base"),
1928        }
1929    }
1930
1931    /// Every layer has been absorbed by prior applies -- the loop breaks on the first iteration
1932    /// and `kept` stays empty, so the result is a bare `Base`. This is the steady-state
1933    /// "extend a just-applied batch" flow: after `apply_batch(A)`, `A`'s own layer has
1934    /// `overlay.len == committed` and the next `new_batch` call should start from a clean
1935    /// terminal.
1936    #[test]
1937    fn trim_committed_all_committed() {
1938        // `shared.len() == 64`; the single layer's `overlay.len == 32 (<= 64)`, so it's committed.
1939        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1940        let chain = make_chain(&shared, &[32]);
1941        let result = chain.trim_committed();
1942        // Collapsed to a bare Base, pointing at the original shared.
1943        match result {
1944            BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
1945            BitmapBatch::Layer(_) => panic!("expected Base after full trim"),
1946        }
1947    }
1948
1949    /// Every layer is still speculative -- the loop walks all the way to `Base` without
1950    /// breaking, and `kept` holds every overlay. The rebuilt chain is structurally equivalent
1951    /// to the input (same overlay lens, same shared terminal). Real-world trigger: speculating
1952    /// multiple batches deep (A, then B off A, then C off B) without `apply_batch` in between.
1953    #[test]
1954    fn trim_committed_none_committed() {
1955        // `shared.len() == 32`; both overlays have `len > 32`, so neither is committed.
1956        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 32])));
1957        let chain = make_chain(&shared, &[64, 96]);
1958        let result = chain.trim_committed();
1959        // Structure must be preserved in bottom-to-top order.
1960        assert_eq!(chain_overlays(&result), vec![64, 96]);
1961    }
1962
1963    /// Exactly one layer is uncommitted (the newest) on top of a committed prefix -- the
1964    /// dominant pattern in chained growth. The loop collects the one uncommitted overlay, and
1965    /// the rebuild produces `Layer(Base, overlay_B)`. Also verifies the rebuilt layer carries
1966    /// the cached `shared` reference correctly. Real-world trigger: apply parent A, then B
1967    /// held alive off A, then `B.new_batch()` to build C.
1968    #[test]
1969    fn trim_committed_exactly_one_uncommitted() {
1970        // `shared.len() == 64`; committed layer (`overlay.len == 64`) + uncommitted (`96`).
1971        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1972        let chain = make_chain(&shared, &[64, 96]);
1973        let result = chain.trim_committed();
1974        // The committed layer is gone; only the uncommitted overlay remains.
1975        assert_eq!(chain_overlays(&result), vec![96]);
1976        // And the rebuilt layer's `shared` field still points at the original terminal.
1977        assert!(Arc::ptr_eq(result.shared(), &shared));
1978    }
1979
1980    /// Two or more uncommitted layers on top of a committed prefix -- exercises the loop's
1981    /// iterated `kept.push` and the rebuild's iterated `Arc::new(BitmapBatchLayer)`, including
1982    /// the cached `shared` wire-through on every reconstructed layer. Real-world trigger:
1983    /// build A, then B off A, then C off B; apply only A; then call `C.new_batch()`.
1984    #[test]
1985    fn trim_committed_multiple_uncommitted() {
1986        // `shared.len() == 64`; committed layer (64), then two uncommitted (96, 128).
1987        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1988        let chain = make_chain(&shared, &[64, 96, 128]);
1989        let result = chain.trim_committed();
1990        // Committed layer dropped; uncommitted pair preserved in order.
1991        assert_eq!(chain_overlays(&result), vec![96, 128]);
1992        // Every reconstructed layer must still cache the original shared terminal.
1993        assert!(Arc::ptr_eq(result.shared(), &shared));
1994    }
1995}