//! Batch mutation API for Current QMDBs.
//!
//! Wraps the [`any::batch`] API.

use crate::{
    Context,
    index::Unordered as UnorderedIndex,
    journal::contiguous::{Contiguous, Mutable},
    merkle::{
        self, Graftable, Location, Position, Readable,
        batch::MerkleizedBatch as GenericMerkleizedBatch, mem::Mem,
        storage::Storage as MerkleStorage,
    },
    qmdb::{
        Error,
        any::{
            self, ValueEncoding,
            batch::{DiffCursors, DiffEntry, Staged as AnyStaged, StagedUpdates},
            operation::{Operation, update},
        },
        batch_chain::Bounds,
        bitmap::{Shared, fill_from},
        current::{
            db::{compute_db_root, partial_chunk, read_graft_inputs},
            grafting,
        },
        operation::Key,
    },
};
use ahash::AHashMap;
use commonware_codec::Codec;
use commonware_cryptography::{Digest, Hasher};
use commonware_parallel::Strategy;
use commonware_utils::bitmap::{self, Readable as _};
use core::ops::Range;
use std::sync::Arc;

/// Speculative chunk-level bitmap overlay.
///
/// Instead of tracking individual pushed bits and cleared locations, maintains materialized chunk
/// bytes for every chunk that differs from the parent bitmap. This directly produces the chunk data
/// needed for grafted MMR leaf computation.
#[derive(Clone, Debug, Default)]
pub(crate) struct ChunkOverlay<const N: usize> {
    /// Dirty chunks: chunk_idx -> materialized chunk bytes.
    ///
    /// Iteration order is not observed by any consumer.
    pub(crate) chunks: AHashMap<usize, [u8; N]>,
    /// Total number of bits (parent + new operations).
    pub(crate) len: u64,
    /// The parent bitmap's dimensions, captured at construction.
    parent: Dimensions,
}

/// Parent-bitmap dimensions captured once per overlay. `chunk_mut` needs them on every newly
/// materialized chunk, and reading each through a `Base` chain costs a lock acquisition on
/// the shared committed bitmap, so they are read once instead of per touched chunk.
#[derive(Clone, Copy, Debug, Default)]
struct Dimensions {
    len: u64,
    complete_chunks: usize,
    pruned_chunks: usize,
}

impl Dimensions {
    fn of<B: bitmap::Readable<N>, const N: usize>(base: &B) -> Self {
        Self {
            len: base.len(),
            complete_chunks: base.complete_chunks(),
            pruned_chunks: base.pruned_chunks(),
        }
    }
}

impl<const N: usize> ChunkOverlay<N> {
    const CHUNK_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;

    /// Create an overlay of `len` total bits on top of `base`. The `base` handed to later
    /// `set_bit` / `clear_bit` / `chunk_mut` calls must be the bitmap given here.
    fn new<B: bitmap::Readable<N>>(base: &B, len: u64, capacity: usize) -> Self {
        Self {
            chunks: AHashMap::with_capacity(capacity),
            len,
            parent: Dimensions::of(base),
        }
    }

    /// Load-or-create a chunk: returns a mutable reference to the materialized chunk bytes. On
    /// first access for an existing chunk, reads from `base`.
    fn chunk_mut<B: bitmap::Readable<N>>(&mut self, base: &B, idx: usize) -> &mut [u8; N] {
        let parent = self.parent;
        self.chunks.entry(idx).or_insert_with(|| {
            let base_has_partial = !parent.len.is_multiple_of(Self::CHUNK_BITS);
            if idx < parent.complete_chunks {
                base.get_chunk(idx)
            } else if idx == parent.complete_chunks && base_has_partial {
                base.last_chunk().0
            } else {
                bitmap::BitMap::<N>::EMPTY_CHUNK
            }
        })
    }

    /// Set a single bit (used for pushes and active operations).
    fn set_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
        let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
        let rel = (loc % Self::CHUNK_BITS) as usize;
        let chunk = self.chunk_mut(base, idx);
        chunk[rel / 8] |= 1 << (rel % 8);
    }

    /// Clear a single bit (used for superseded locations). Skips locations in pruned chunks
    /// since those bits are already inactive.
    fn clear_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
        let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
        if idx < self.parent.pruned_chunks {
            return;
        }
        let rel = (loc % Self::CHUNK_BITS) as usize;
        let chunk = self.chunk_mut(base, idx);
        chunk[rel / 8] &= !(1 << (rel % 8));
    }

    /// Get a dirty chunk's bytes, or `None` if unmodified.
    pub(crate) fn get(&self, idx: usize) -> Option<&[u8; N]> {
        self.chunks.get(&idx)
    }

    /// Number of complete chunks.
    pub(crate) const fn complete_chunks(&self) -> usize {
        (self.len / Self::CHUNK_BITS) as usize
    }
}

/// Bitmap-accelerated floor scan over a layered `BitmapBatch` chain. Fills `out` with up to
/// `limit` floor-raise candidates in `[floor, tip)`, returning the next `floor`. Skips
/// locations where the layered bitmap bit is unset (including locations superseded by
/// uncommitted ancestors), avoiding I/O reads for inactive operations. Produces the same
/// sequence as repeatedly calling the `next_candidate` test oracle over the chain.
///
/// One scan iterator serves the whole batch: overlay chunks resolve lock-free and the
/// committed base is locked once per untouched chunk, rather than several times per
/// candidate. The iterator's chunk caching is sound here because bitmap mutators require
/// `&mut` on the database, which cannot coexist with the `&db` a merkleize holds.
pub(crate) fn fill_candidates<F: Graftable, const N: usize>(
    bitmap: &BitmapBatch<N>,
    floor: Location<F>,
    tip: u64,
    limit: usize,
    out: &mut Vec<Location<F>>,
) -> Location<F> {
    Location::new(fill_from(bitmap, *floor, tip, limit, out))
}

/// Adapter that resolves ops MMR nodes for a batch's `compute_current_layer`.
///
/// Tries the batch chain's sync [`Readable`] first (which covers nodes appended or overwritten
/// by the batch, plus anything still in the in-memory MMR). Falls through to the base's async
/// [`MerkleStorage`].
struct BatchStorageAdapter<
    'a,
    F: Graftable,
    D: Digest,
    R: Readable<Family = F, Digest = D>,
    S: MerkleStorage<F, Digest = D>,
> {
    batch: &'a R,
    base: &'a S,
    _phantom: core::marker::PhantomData<(F, D)>,
}

impl<
    'a,
    F: Graftable,
    D: Digest,
    R: Readable<Family = F, Digest = D>,
    S: MerkleStorage<F, Digest = D>,
> BatchStorageAdapter<'a, F, D, R, S>
{
    const fn new(batch: &'a R, base: &'a S) -> Self {
        Self {
            batch,
            base,
            _phantom: core::marker::PhantomData,
        }
    }
}

impl<F: Graftable, D: Digest, R: Readable<Family = F, Digest = D>, S: MerkleStorage<F, Digest = D>>
    MerkleStorage<F> for BatchStorageAdapter<'_, F, D, R, S>
{
    type Digest = D;

    fn size(&self) -> Position<F> {
        self.batch.size()
    }
    async fn get_node(&self, pos: Position<F>) -> Result<Option<D>, merkle::Error<F>> {
        if let Some(node) = self.batch.get_node(pos) {
            return Ok(Some(node));
        }
        self.base.get_node(pos).await
    }

    async fn get_nodes(&self, positions: &[Position<F>]) -> Result<Vec<D>, merkle::Error<F>> {
        let mut nodes = vec![None; positions.len()];
        let mut base_positions = Vec::with_capacity(positions.len());

        // Look up nodes already in the batch chain.
        for (slot, &pos) in nodes.iter_mut().zip(positions) {
            match self.batch.get_node(pos) {
                Some(node) => *slot = Some(node),
                None => base_positions.push(pos),
            }
        }

        // Look up remaining nodes from the base.
        let base_nodes = if base_positions.is_empty() {
            Vec::new()
        } else {
            self.base.get_nodes(&base_positions).await?
        };
        let mut base_nodes = base_nodes.into_iter();
        Ok(nodes
            .into_iter()
            .map(|node| node.unwrap_or_else(|| base_nodes.next().expect("one node per base read")))
            .collect())
    }
}

/// Layers a [`GenericMerkleizedBatch`] over a [`Mem`] for node resolution.
///
/// [`GenericMerkleizedBatch::get_node`] only covers the batch chain; committed positions
/// return `None`. This adapter falls through to the committed Mem for those positions.
struct BatchOverMem<'a, F: Graftable, D: Digest, S: Strategy> {
    batch: &'a GenericMerkleizedBatch<F, D, S>,
    mem: &'a Mem<F, D>,
}

impl<F: Graftable, D: Digest, S: Strategy> Readable for BatchOverMem<'_, F, D, S> {
    type Family = F;
    type Digest = D;

    fn size(&self) -> Position<F> {
        self.batch.size()
    }

    fn get_node(&self, pos: Position<F>) -> Option<D> {
        if let Some(d) = self.batch.get_node(pos) {
            return Some(d);
        }
        self.mem.get_node(pos)
    }
}

/// A speculative batch of mutations whose root digest has not yet been computed,
/// in contrast to [`MerkleizedBatch`].
///
/// Wraps a [`any::batch::UnmerkleizedBatch`] and adds bitmap and grafted MMR parent state
/// needed to compute the current layer during [`merkleize`](Self::merkleize).
pub struct UnmerkleizedBatch<F, H, U, const N: usize, S: Strategy>
where
    F: Graftable,
    U: update::Update,
    H: Hasher,
    Operation<F, U>: Codec,
{
    /// The inner any-layer batch that handles mutations, journal, and floor raise.
    inner: any::batch::UnmerkleizedBatch<F, H, U, S>,

    /// Parent's grafted MMR state.
    grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,

    /// Parent's bitmap state (COW, Arc-based).
    bitmap_parent: BitmapBatch<N>,
}

/// Staged batch returned by [`UnmerkleizedBatch::stage`].
pub struct Staged<F, H, U, const N: usize, S: Strategy>
where
    F: Graftable,
    U: update::Update,
    H: Hasher,
    Operation<F, U>: Codec,
{
    inner: AnyStaged<F, H, U, S>,
    grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
    bitmap_parent: BitmapBatch<N>,
}

/// A speculative batch of operations whose root digest has been computed, in contrast to
/// [`UnmerkleizedBatch`].
///
/// Wraps an [`any::batch::MerkleizedBatch`] and adds the bitmap and grafted MMR state needed to
/// compute the canonical root.
///
/// # Branch validity
///
/// A `MerkleizedBatch` is a branch-scoped view rooted at a specific committed prefix of the DB. It
/// is not an immutable snapshot.
///
/// Internally, the batch chain terminates in the DB's committed bitmap via `BitmapBatch::Base`.
/// That committed bitmap evolves in place as [`Db::apply_batch`](super::db::Db::apply_batch),
/// [`Db::prune`](super::db::Db::prune), and [`Db::rewind`](super::db::Db::rewind) update the DB.
///
/// Reads through this batch's chain, constructing child batches from it, and applying it later are
/// only semantically correct while its ancestor chain is still the committed prefix of the DB. In
/// other words, every successful [`apply_batch`](super::db::Db::apply_batch) since this batch was
/// merkleized must have applied an ancestor of this batch.
///
/// Once a non-ancestor batch is applied, this batch and all of its descendants become invalid
/// objects. The library does not guard against continued use after that point.
///
/// Applying an invalid batch is caught by the any-layer authenticated lineage check and returns
/// [`Error::StaleBatch`] without mutating committed state, so `apply_batch` itself cannot corrupt
/// the DB.
///
/// Rules of thumb:
/// - Drop any `Arc<MerkleizedBatch>` you no longer intend to apply.
/// - Extending a batch after `apply_batch` has consumed it (building a child off the just-applied
///   parent) is safe. The committed bitmap now equals the parent's post-apply state, so child reads
///   are consistent.
/// - Extending a batch after a different branch has been applied is not safe. Do not call `get`,
///   `new_batch`, or `apply_batch` on that branch again.
pub struct MerkleizedBatch<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
{
    /// Inner any-layer batch (ops MMR, diff, floor, commit loc, sizes).
    pub(crate) inner: Arc<any::batch::MerkleizedBatch<F, D, U, S>>,

    /// Grafted MMR state.
    pub(crate) grafted: Arc<merkle::batch::MerkleizedBatch<F, D, S>>,

    /// COW bitmap state (for use as a parent in speculative batches).
    pub(crate) bitmap: BitmapBatch<N>,

    /// The canonical root (ops root + grafted root + partial chunk).
    pub(crate) canonical_root: D,
}

impl<F, H, U, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, U, N, S>
where
    F: Graftable,
    U: update::Update,
    H: Hasher,
    Operation<F, U>: Codec,
{
    pub(super) const fn new(
        inner: any::batch::UnmerkleizedBatch<F, H, U, S>,
        grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
        bitmap_parent: BitmapBatch<N>,
    ) -> Self {
        Self {
            inner,
            grafted_parent,
            bitmap_parent,
        }
    }

    /// Record a mutation. Use `Some(value)` for update/create, `None` for delete.
    ///
    /// If the same key is written multiple times within a batch, the last
    /// value wins.
    pub fn write(mut self, key: U::Key, value: Option<U::Value>) -> Self {
        self.inner = self.inner.write(key, value);
        self
    }

    /// Read through: mutations -> ancestor diffs -> committed DB.
    pub async fn get<E, C, I>(
        &self,
        key: &U::Key,
        db: &super::db::Db<F, E, C, I, H, U, N, S>,
    ) -> Result<Option<U::Value>, Error<F>>
    where
        E: Context,
        C: Contiguous<Item = Operation<F, U>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
    {
        self.inner.get(key, &db.any).await
    }

    /// Batch read multiple keys.
    ///
    /// Returns results in the same order as the input keys. Resolved locations are not retained,
    /// so writing a key read only through `get_many` requires an index re-probe and journal re-read
    /// during merkleize. Use [`stage`](Self::stage) for keys that may be written. When the writable
    /// subset is known and much smaller than the full read set, call `get_many` for the read-only
    /// keys first, then [`stage`](Self::stage) only the writable keys.
    pub async fn get_many<E, C, I>(
        &self,
        keys: &[&U::Key],
        db: &super::db::Db<F, E, C, I, H, U, N, S>,
    ) -> Result<Vec<Option<U::Value>>, Error<F>>
    where
        E: Context,
        C: Contiguous<Item = Operation<F, U>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
    {
        self.inner.get_many(keys, &db.any).await
    }

    /// Batch read multiple keys and return a staged batch for the same keys.
    ///
    /// Returns results in the same order as the input keys. The staged batch records updates by
    /// read index: the initial keys occupy `0..keys.len()`, and each [`expand`](Staged::expand)
    /// appends another index range.
    pub async fn stage<E, C, I>(
        self,
        keys: &[&U::Key],
        db: &super::db::Db<F, E, C, I, H, U, N, S>,
    ) -> Result<(Vec<Option<U::Value>>, Staged<F, H, U, N, S>), Error<F>>
    where
        E: Context,
        C: Contiguous<Item = Operation<F, U>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
    {
        let Self {
            inner,
            grafted_parent,
            bitmap_parent,
        } = self;
        let (values, inner) = inner.stage(keys, &db.any).await?;
        Ok((
            values,
            Staged {
                inner,
                grafted_parent,
                bitmap_parent,
            },
        ))
    }
}

impl<F, H, U, const N: usize, S: Strategy> Staged<F, H, U, N, S>
where
    F: Graftable,
    U: update::Update,
    H: Hasher,
    Operation<F, U>: Codec,
{
    /// Expand this staged batch with more reads.
    ///
    /// Existing read indices remain stable. Newly read keys are appended to the staged read set and
    /// assigned the returned range. The returned values are in the same order as `keys`.
    ///
    /// Expansion does not deduplicate against previously staged keys and does not observe values the
    /// caller has computed for earlier staged slots but not yet passed to
    /// [`merkleize`](Staged::merkleize).
    pub async fn expand<E, C, I>(
        self,
        keys: &[&U::Key],
        db: &super::db::Db<F, E, C, I, H, U, N, S>,
    ) -> Result<(Range<usize>, Vec<Option<U::Value>>, Self), Error<F>>
    where
        E: Context,
        C: Contiguous<Item = Operation<F, U>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
    {
        let Self {
            inner,
            grafted_parent,
            bitmap_parent,
        } = self;
        let (range, values, inner) = inner.expand(keys, &db.any).await?;
        Ok((
            range,
            values,
            Self {
                inner,
                grafted_parent,
                bitmap_parent,
            },
        ))
    }
}

impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Unordered<K, V>, N, S>
where
    F: Graftable,
    K: Key,
    V: ValueEncoding,
    H: Hasher,
    Operation<F, update::Unordered<K, V>>: Codec,
{
    /// Record updates for staged reads and upserts for unread keys, then merkleize.
    ///
    /// Consumes the staged handle and write vectors. Call [`expand`](Staged::expand) before this
    /// method if more keys must be read into the staged index space.
    ///
    /// A `Some` value is an upsert. `None` is a delete. Update indices refer to the staged read
    /// set: the initial `stage` input followed by any [`expand`](Staged::expand) ranges. `metadata`
    /// is committed with the returned batch.
    ///
    /// # Panics
    ///
    /// Panics if any update's `read_index` is out of the staged read range.
    #[allow(clippy::type_complexity)]
    #[tracing::instrument(
        name = "qmdb.current.unordered.batch.merkleize.staged",
        level = "info",
        skip_all,
        fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
    )]
    pub async fn merkleize<E, C, I>(
        self,
        updates: Vec<(usize, Option<V::Value>)>,
        upserts: Vec<(K, Option<V::Value>)>,
        metadata: Option<V::Value>,
        db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
    {
        let Self {
            inner,
            grafted_parent,
            bitmap_parent,
        } = self;

        // Overlap the update resolution with a committed-prefix candidate prefetch.
        // Candidates come from the speculative `bitmap_parent` (the same source the floor
        // raise scans below), clamped to the committed prefix inside the helper.
        let (inner, staged_updates, prefetched) = inner
            .resolve_updates_prefetched(updates, upserts, &db.any, |floor, tip, limit, out| {
                fill_candidates(&bitmap_parent, floor, tip, limit, out)
            })
            .await?;
        let inner = inner
            .merkleize_with_floor_scan(
                &db.any,
                metadata,
                staged_updates,
                Some(prefetched),
                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
            )
            .await?;
        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
    }
}

impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Ordered<K, V>, N, S>
where
    F: Graftable,
    K: Key,
    V: ValueEncoding,
    H: Hasher,
    Operation<F, update::Ordered<K, V>>: Codec,
{
    /// Record updates for staged reads and upserts for unread keys, then merkleize.
    ///
    /// Consumes the staged handle and write vectors. Call [`expand`](Staged::expand) before this
    /// method if more keys must be read into the staged index space.
    ///
    /// A `Some` value is an upsert. `None` is a delete. Update indices refer to the staged read
    /// set: the initial `stage` input followed by any [`expand`](Staged::expand) ranges. `metadata`
    /// is committed with the returned batch.
    ///
    /// # Panics
    ///
    /// Panics if any update's `read_index` is out of the staged read range.
    #[allow(clippy::type_complexity)]
    #[tracing::instrument(
        name = "qmdb.current.ordered.batch.merkleize.staged",
        level = "info",
        skip_all,
        fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
    )]
    pub async fn merkleize<E, C, I>(
        self,
        updates: Vec<(usize, Option<V::Value>)>,
        upserts: Vec<(K, Option<V::Value>)>,
        metadata: Option<V::Value>,
        db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
        I: crate::index::Ordered<Value = Location<F>> + 'static,
    {
        let Self {
            inner,
            grafted_parent,
            bitmap_parent,
        } = self;
        let (inner, staged_updates) = inner.resolve_updates(updates, upserts, db.any.strategy());
        let inner = inner
            .merkleize_with_floor_scan(
                &db.any,
                metadata,
                staged_updates,
                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
            )
            .await?;
        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
    }
}

// Unordered merkleize.
impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
where
    F: Graftable,
    K: Key,
    V: ValueEncoding,
    H: Hasher,
    Operation<F, update::Unordered<K, V>>: Codec,
{
    /// Resolve mutations into operations, merkleize, and return an `Arc<MerkleizedBatch>`.
    #[allow(clippy::type_complexity)]
    #[tracing::instrument(
        name = "qmdb.current.unordered.batch.merkleize",
        level = "info",
        skip_all
    )]
    pub async fn merkleize<E, C, I>(
        self,
        db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
        metadata: Option<V::Value>,
    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
    {
        let Self {
            inner,
            grafted_parent,
            bitmap_parent,
        } = self;
        // Use the speculative parent bitmap rather than the committed `any` bitmap.
        let inner = inner
            .merkleize_with_floor_scan(
                &db.any,
                metadata,
                StagedUpdates::<F, update::Unordered<K, V>>::new(),
                None,
                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
            )
            .await?;
        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
    }
}

// Ordered merkleize.
impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
where
    F: Graftable,
    K: Key,
    V: ValueEncoding,
    H: Hasher,
    Operation<F, update::Ordered<K, V>>: Codec,
{
    /// Resolve mutations into operations, merkleize, and return an `Arc<MerkleizedBatch>`.
    #[allow(clippy::type_complexity)]
    #[tracing::instrument(
        name = "qmdb.current.ordered.batch.merkleize",
        level = "info",
        skip_all
    )]
    pub async fn merkleize<E, C, I>(
        self,
        db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
        metadata: Option<V::Value>,
    ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
        I: crate::index::Ordered<Value = Location<F>> + 'static,
    {
        let Self {
            inner,
            grafted_parent,
            bitmap_parent,
        } = self;
        // Use the speculative parent bitmap rather than the committed `any` bitmap.
        let inner = inner
            .merkleize_with_floor_scan(
                &db.any,
                metadata,
                StagedUpdates::<F, update::Ordered<K, V>>::new(),
                |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
            )
            .await?;
        compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
    }
}

/// Derive all bitmap mutations (pushes + clears) for this batch in a single pass over the diff and
/// ancestor diffs. Avoids iterating raw operations.
///
/// Pushes: one bit per operation in the batch. All false except active diff entries (whose `loc`
/// falls in the batch) and the CommitFloor (last op).
///
/// Clears: previous CommitFloor, plus the most recent superseded location for each mutated key. We
/// search back through ancestors to find the most recent active location; if none exists, we clear
/// the committed DB location (`base_old_loc`).
#[allow(clippy::type_complexity)]
fn build_chunk_overlay<F: Graftable, U, B: bitmap::Readable<N>, const N: usize>(
    base: &B,
    batch_len: usize,
    batch_base: u64,
    diff: &[(U::Key, DiffEntry<F, U::Value>)],
    ancestor_diffs: &[Arc<Vec<(U::Key, DiffEntry<F, U::Value>)>>],
) -> ChunkOverlay<N>
where
    U: update::Update,
{
    let total_bits = base.len() + batch_len as u64;
    let appended_chunks = (batch_len as u64).div_ceil(ChunkOverlay::<N>::CHUNK_BITS) as usize;
    let mut overlay = ChunkOverlay::new(base, total_bits, diff.len() + appended_chunks + 1);

    // 1. CommitFloor (last op) is always active.
    let commit_loc = batch_base + batch_len as u64 - 1;
    overlay.set_bit(base, commit_loc);

    // 2. Inactivate previous CommitFloor.
    overlay.clear_bit(base, batch_base - 1);

    // 3. Set active bits + clear superseded locations from the diff. The diff is key-sorted,
    // so ancestor resolution streams (one cursor per ancestor diff).
    let mut ancestors = DiffCursors::new(ancestor_diffs.iter().map(|d| d.as_slice()));
    for (key, entry) in diff {
        // Set the active bit for this key's final location.
        if let Some(loc) = entry.loc()
            && *loc >= batch_base
            && *loc < batch_base + batch_len as u64
        {
            overlay.set_bit(base, *loc);
        }

        // Clear the most recent superseded location. Older locations were already cleared by the
        // ancestor batch that superseded them.
        let mut prev_loc = entry.base_old_loc();
        if let Some(ancestor_entry) = ancestors.resolve(key) {
            prev_loc = ancestor_entry.loc();
        }
        if let Some(old) = prev_loc {
            overlay.clear_bit(base, *old);
        }
    }

    // Ensure all new complete chunks beyond the parent are materialized, so downstream consumers
    // don't read from the parent and panic on out-of-range indices. Uses chunk_mut to inherit the
    // parent's partial chunk data when idx == parent_complete (avoiding loss of existing bits).
    let parent_complete = overlay.parent.complete_chunks;
    let new_complete = overlay.complete_chunks();
    for idx in parent_complete..new_complete {
        overlay.chunk_mut(base, idx);
    }

    overlay
}

/// Merkleize grafted chunk digests while retaining the live ancestor chain.
async fn merkleize_grafted_batch<F, H, S, const N: usize>(
    strategy: &S,
    grafted_parent: Arc<GenericMerkleizedBatch<F, H::Digest, S>>,
    grafted_tree: &Arc<Mem<F, H::Digest>>,
    graft_inputs: Vec<(usize, H::Digest, [u8; N])>,
    grafting_height: u32,
) -> Arc<GenericMerkleizedBatch<F, H::Digest, S>>
where
    F: Graftable,
    H: Hasher,
    S: Strategy,
{
    let old_grafted_leaves = *grafted_parent.leaves() as usize;
    let mut grafted_batch = grafted_parent.new_batch();
    let ancestors = grafted_batch.retain_ancestors();
    let grafted_tree = Arc::clone(grafted_tree);
    strategy
        .clone()
        .spawn(graft_inputs.len(), move |strategy| {
            let new_leaves = grafting::graft_chunk_digests::<H, _, N>(&strategy, graft_inputs);
            for (chunk_idx, digest) in new_leaves {
                if chunk_idx < old_grafted_leaves {
                    grafted_batch = grafted_batch
                        .update_leaf_digest(Location::<F>::new(chunk_idx as u64), digest)
                        .expect("update_leaf_digest failed");
                } else {
                    grafted_batch = grafted_batch.add_leaf_digest(digest);
                }
            }
            let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
            let merkleized = grafted_batch.merkleize(&grafted_tree, &grafted_hasher);
            drop(ancestors);
            merkleized
        })
        .await
}

/// Compute the current layer (bitmap + grafted MMR + canonical root) on top of a merkleized any
/// batch.
///
/// Builds a chunk overlay from the diff, computes grafted MMR leaves from dirty chunks, and
/// produces the `Arc<MerkleizedBatch>` directly.
async fn compute_current_layer<F, E, U, C, I, H, const N: usize, S>(
    inner: Arc<any::batch::MerkleizedBatch<F, H::Digest, U, S>>,
    current_db: &super::db::Db<F, E, C, I, H, U, N, S>,
    grafted_parent: &Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
    bitmap_parent: &BitmapBatch<N>,
) -> Result<Arc<MerkleizedBatch<F, H::Digest, U, N, S>>, Error<F>>
where
    F: Graftable,
    E: Context,
    C: Contiguous<Item = Operation<F, U>>,
    I: UnorderedIndex<Value = Location<F>>,
    H: Hasher,
    U: update::Update,
    S: Strategy,
    Operation<F, U>: Codec,
{
    let batch_len = inner.journal_batch.items().len();
    let batch_base = *inner.bounds.tip.size - batch_len as u64;

    // Build chunk overlay: materialized bytes for every dirty chunk.
    let overlay = build_chunk_overlay::<F, U, _, N>(
        bitmap_parent,
        batch_len,
        batch_base,
        &inner.diff,
        &inner.ancestor_diffs,
    );

    let grafting_height = grafting::height::<N>();
    let ops_tree_adapter =
        BatchStorageAdapter::new(&inner.journal_batch, &current_db.any.log.merkle);

    // Snapshot ops_leaves for the post-batch state (the canonical root we're about to compute
    // sees this many ops). Thread it through `graftable_chunks` derivation and root computation.
    let overlay_ops_leaves = inner.bounds.tip.size;

    // Distinguish three counters:
    //   - new_complete_chunks: chunks with all bits filled in the post-batch bitmap
    //   - graftable_overlay:      chunks committed by the grafted tree (have a single h=G ancestor)
    //   - graftable_parent:       grafted-tree leaf count from the parent (structural source of truth)
    //
    // The pending chunk (if any) sits at index `graftable_overlay` and is excluded from the
    // grafted tree; its digest is hashed directly into the canonical root.
    let new_complete_chunks = overlay.complete_chunks();
    let graftable_overlay = grafting::graftable_chunks::<F>(*overlay_ops_leaves, grafting_height)
        .min(new_complete_chunks as u64) as usize;
    let graftable_parent = *grafted_parent.leaves() as usize;
    let pruned_chunks = bitmap_parent.pruned_chunks();
    assert!(
        pruned_chunks <= graftable_parent
            && graftable_parent <= graftable_overlay
            && graftable_overlay <= new_complete_chunks,
        "invariant violated: pruned={pruned_chunks} graftable_parent={graftable_parent} graftable_overlay={graftable_overlay} new_complete={new_complete_chunks}"
    );

    // Build the set of chunk indices whose grafted-leaf needs (re)computing:
    //   1) Dirty chunks (bits changed in this batch) within the graftable range.
    //   2) Pending -> graftable transitions: chunks newly graftable because the ops tree built
    //      their h=G ancestor in this batch. Their bitmap bytes may not be dirty (the chunk
    //      became graftable via ops growth alone) but they need a grafted-leaf entry now.
    let mut chunk_indices_to_update: Vec<usize> = overlay
        .chunks
        .iter()
        .filter(|&(&idx, _)| idx < graftable_overlay && idx >= pruned_chunks)
        .map(|(&idx, _)| idx)
        .collect();
    chunk_indices_to_update.extend(graftable_parent..graftable_overlay);
    chunk_indices_to_update.sort_unstable();
    chunk_indices_to_update.dedup();
    let chunks_to_update = chunk_indices_to_update.into_iter().map(|idx| {
        let chunk = overlay
            .get(idx)
            .copied()
            .unwrap_or_else(|| bitmap_parent.get_chunk(idx));
        (idx, chunk)
    });

    // Prefetch each chunk's covering ops-tree node, then run graft hashing and the grafted
    // MMR build/merkleize as one job through the strategy (against a snapshot of the
    // committed grafted tree). An empty graft set hashes nothing, so it merkleizes without
    // submitting a job.
    let graft_inputs = read_graft_inputs::<F, _, N>(&ops_tree_adapter, chunks_to_update).await?;
    let grafted_batch = if graft_inputs.is_empty() {
        let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
        grafted_parent
            .new_batch()
            .merkleize(&current_db.grafted_tree, &grafted_hasher)
    } else {
        merkleize_grafted_batch::<F, H, S, N>(
            &current_db.strategy,
            Arc::clone(grafted_parent),
            &current_db.grafted_tree,
            graft_inputs,
            grafting_height,
        )
        .await
    };

    // Build the layered bitmap (parent + overlay) before computing the canonical root, so that
    // compute_db_root sees newly completed chunks. Using bitmap_parent alone would miss chunks
    // that transitioned from partial to complete in this batch.
    let bitmap_batch = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
        parent: bitmap_parent.clone(),
        overlay: Arc::new(overlay),
        shared: Arc::clone(bitmap_parent.shared()),
    }));

    // Compute canonical root. The grafted batch alone cannot resolve committed nodes,
    // so layer it over the committed grafted MMR.
    let ops_root = inner.root();
    let layered = BatchOverMem {
        batch: &grafted_batch,
        mem: &current_db.grafted_tree,
    };
    let grafted_storage =
        grafting::Storage::<F, H, _, _>::new(&layered, grafting_height, &ops_tree_adapter);
    // Compute partial chunk (last incomplete chunk, if any). The partial chunk lives at
    // index `new_complete_chunks` (the chunk currently being filled with bits) -- distinct
    // from `graftable_overlay` (the grafted-tree boundary). At gh >= 3, partial and pending can
    // coexist; this branch only handles partial. The pending chunk (when present) is read
    // from the bitmap inside `compute_db_root` via `pending_chunk()`.
    let partial = partial_chunk::<_, N>(&bitmap_batch);
    let canonical_root = compute_db_root::<F, H, _, _, N>(
        &bitmap_batch,
        &grafted_storage,
        overlay_ops_leaves,
        partial,
        inner.bounds.inactivity_floor,
        &ops_root,
    )
    .await?;

    Ok(Arc::new(MerkleizedBatch {
        inner,
        grafted: grafted_batch,
        bitmap: bitmap_batch,
        canonical_root,
    }))
}

/// A view of the committed bitmap plus zero or more speculative overlay `Layer`s.
///
/// The chain terminates in a `Base` that references the shared committed bitmap. No validity
/// check is performed. Callers must ensure they only read through batches whose chains are
/// still valid prefixes of committed state (see [`Shared`]'s docs).
#[derive(Clone, Debug)]
pub(crate) enum BitmapBatch<const N: usize> {
    /// Chain terminal: shared reference to the committed bitmap.
    Base(Arc<Shared<N>>),
    /// Speculative layer on top of a parent batch.
    Layer(Arc<BitmapBatchLayer<N>>),
}

/// The data behind a [`BitmapBatch::Layer`].
#[derive(Debug)]
pub(crate) struct BitmapBatchLayer<const N: usize> {
    pub(crate) parent: BitmapBatch<N>,
    /// Chunk-level overlay: materialized bytes for every chunk that differs from parent.
    pub(crate) overlay: Arc<ChunkOverlay<N>>,
    /// Cached terminal [`Shared`] so [`BitmapBatch::shared`] and
    /// [`BitmapBatch::pruned_chunks`] answer in O(1) instead of walking the chain.
    pub(crate) shared: Arc<Shared<N>>,
}

impl<const N: usize> BitmapBatch<N> {
    const CHUNK_SIZE_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;

    /// Return the terminal [`Shared`] at the bottom of the chain.
    fn shared(&self) -> &Arc<Shared<N>> {
        match self {
            Self::Base(s) => s,
            Self::Layer(layer) => &layer.shared,
        }
    }

    /// Return a chain equivalent to `self` with any `Layer` whose overlay is now fully committed
    /// replaced by a direct reference to the committed bitmap. Since `apply_batch` commits
    /// contiguous prefixes, committed `Layer`s are always at the bottom of the chain.
    fn trim_committed(&self) -> Self {
        let shared = self.shared();
        let committed = bitmap::Readable::<N>::len(shared.as_ref());
        let mut kept = Vec::new();
        let mut current = self;
        while let Self::Layer(layer) = current {
            if layer.overlay.len <= committed {
                break;
            }
            kept.push(Arc::clone(&layer.overlay));
            current = &layer.parent;
        }
        let mut result = Self::Base(Arc::clone(shared));
        for overlay in kept.into_iter().rev() {
            result = Self::Layer(Arc::new(BitmapBatchLayer {
                parent: result,
                overlay,
                shared: Arc::clone(shared),
            }));
        }
        result
    }
}

impl<const N: usize> bitmap::Readable<N> for BitmapBatch<N> {
    fn complete_chunks(&self) -> usize {
        (self.len() / Self::CHUNK_SIZE_BITS) as usize
    }

    fn get_chunk(&self, idx: usize) -> [u8; N] {
        // Walk the layer chain. Each layer's overlay either holds the chunk (return it) or
        // doesn't (descend).
        let mut current = self;
        loop {
            match current {
                Self::Base(shared) => return shared.get_chunk(idx),
                Self::Layer(layer) => {
                    if let Some(&chunk) = layer.overlay.get(idx) {
                        return chunk;
                    }
                    current = &layer.parent;
                }
            }
        }
    }

    fn last_chunk(&self) -> ([u8; N], u64) {
        let total = self.len();
        if total == 0 {
            return (bitmap::BitMap::<N>::EMPTY_CHUNK, 0);
        }
        let rem = total % Self::CHUNK_SIZE_BITS;
        let bits_in_last = if rem == 0 { Self::CHUNK_SIZE_BITS } else { rem };
        let idx = if rem == 0 {
            self.complete_chunks().saturating_sub(1)
        } else {
            self.complete_chunks()
        };
        (self.get_chunk(idx), bits_in_last)
    }

    fn pruned_chunks(&self) -> usize {
        self.shared().pruned_chunks()
    }

    fn len(&self) -> u64 {
        match self {
            Self::Base(shared) => bitmap::Readable::<N>::len(shared.as_ref()),
            Self::Layer(layer) => layer.overlay.len,
        }
    }
}

impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
    MerkleizedBatch<F, D, U, N, S>
{
    /// Return the canonical root.
    pub const fn root(&self) -> D {
        self.canonical_root
    }

    /// Return the QMDB ops-only root.
    pub fn ops_root(&self) -> D {
        self.inner.root()
    }

    /// Return the [`Bounds`] of the batch.
    pub fn bounds(&self) -> &Bounds<F, D> {
        self.inner.bounds()
    }

    /// Return the operations this batch appends to the ops log and the location of the first.
    ///
    /// Delegates to the wrapped ops-level batch. The bitmap state contributes to the
    /// canonical root but appends no log operations. There are no matching proof or
    /// pinned-node methods: an ops-level proof verifies only against [`Self::ops_root`],
    /// never the grafted [`Self::root`] that blocks commit to.
    pub fn operations(&self) -> (Location<F>, Arc<Vec<Operation<F, U>>>) {
        self.inner.operations()
    }

    /// Return the batch's safe sync boundary.
    ///
    /// This equals the boundary [`super::db::Db::sync_boundary`] reports once this batch is applied.
    pub fn sync_boundary(&self) -> Location<F> {
        // Derive from the commit's chunk-aligned inactivity floor, the same quantity the DB uses
        // after apply. Deliberately not the physical bitmap pruning boundary, which can lag the
        // inactivity floor when pruning has not run.
        super::db::sync_boundary::<F, N>(
            *self.inner.bounds().inactivity_floor / bitmap::Prunable::<N>::CHUNK_SIZE_BITS,
            *self.inner.bounds().tip.size,
        )
    }
}

impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
    MerkleizedBatch<F, D, U, N, S>
where
    Operation<F, U>: Codec,
{
    /// Create a new speculative batch of operations with this batch as its parent.
    ///
    /// All uncommitted ancestors in the chain must be kept alive until the child (or any
    /// descendant) is merkleized. Dropping an uncommitted ancestor causes data
    /// loss detected at `apply_batch` time.
    ///
    /// This is only valid while `self` is still on the winning branch. If a different branch has
    /// been applied since `self` was created, `self` is no longer a valid parent and must not be
    /// extended.
    pub fn new_batch<H>(self: &Arc<Self>) -> UnmerkleizedBatch<F, H, U, N, S>
    where
        H: Hasher<Digest = D>,
    {
        UnmerkleizedBatch::new(
            self.inner.new_batch::<H>(),
            Arc::clone(&self.grafted),
            self.bitmap.trim_committed(),
        )
    }

    /// Read through: local diff -> ancestor diffs -> committed DB.
    ///
    /// This is only valid while `self` remains on the committed prefix. If a non-ancestor batch
    /// has been applied since `self` was merkleized, do not read through it.
    pub async fn get<E, C, I, H>(
        &self,
        key: &U::Key,
        db: &super::db::Db<F, E, C, I, H, U, N, S>,
    ) -> Result<Option<U::Value>, Error<F>>
    where
        E: Context,
        C: Contiguous<Item = Operation<F, U>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
        H: Hasher<Digest = D>,
    {
        self.inner.get(key, &db.any).await
    }

    /// Batch read multiple keys.
    ///
    /// Returns results in the same order as the input keys.
    pub async fn get_many<E, C, I, H>(
        &self,
        keys: &[&U::Key],
        db: &super::db::Db<F, E, C, I, H, U, N, S>,
    ) -> Result<Vec<Option<U::Value>>, Error<F>>
    where
        E: Context,
        C: Contiguous<Item = Operation<F, U>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
        H: Hasher<Digest = D>,
    {
        self.inner.get_many(keys, &db.any).await
    }
}

impl<F, E, C, I, H, U, const N: usize, S> super::db::Db<F, E, C, I, H, U, N, S>
where
    F: Graftable,
    E: Context,
    C: Contiguous<Item = Operation<F, U>>,
    I: UnorderedIndex<Value = Location<F>>,
    H: Hasher,
    U: update::Update,
    S: Strategy,
    Operation<F, U>: Codec,
{
    /// Create an initial [`MerkleizedBatch`] from the current committed DB state.
    ///
    /// The returned batch is rooted at the current committed prefix, but it is not a persistent
    /// snapshot across later divergent commits. If some other branch is applied afterward, this
    /// batch is no longer valid and must not be read through, extended, or applied.
    pub fn to_batch(&self) -> Arc<MerkleizedBatch<F, H::Digest, U, N, S>> {
        let grafted = self.grafted_snapshot();
        Arc::new(MerkleizedBatch {
            inner: self.any.to_batch(),
            grafted,
            bitmap: BitmapBatch::Base(Arc::clone(&self.any.bitmap)),
            canonical_root: self.root,
        })
    }
}

#[cfg(any(test, feature = "test-traits"))]
mod trait_impls {
    use super::*;
    use crate::{
        journal::contiguous::Mutable,
        qmdb::any::traits::{
            ApplyBatchResult, BatchableDb, MerkleizedBatch as MerkleizedBatchTrait,
            UnmerkleizedBatch as UnmerkleizedBatchTrait,
        },
    };
    use std::future::Future;

    type CurrentDb<F, E, C, I, H, U, const N: usize, S> =
        crate::qmdb::current::db::Db<F, E, C, I, H, U, N, S>;

    impl<F, K, V, H, E, C, I, const N: usize, S>
        UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>>
        for UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
    where
        F: Graftable,
        K: Key,
        V: ValueEncoding + 'static,
        H: Hasher,
        E: Context,
        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
        S: Strategy,
        Operation<F, update::Unordered<K, V>>: Codec,
    {
        type Family = F;
        type K = K;
        type V = V::Value;
        type Metadata = V::Value;
        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;

        fn write(self, key: K, value: Option<V::Value>) -> Self {
            Self::write(self, key, value)
        }

        async fn merkleize(
            self,
            db: &CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>,
            metadata: Option<V::Value>,
        ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
            self.merkleize(db, metadata).await
        }
    }

    impl<F, K, V, H, E, C, I, const N: usize, S>
        UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>>
        for UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
    where
        F: Graftable,
        K: Key,
        V: ValueEncoding + 'static,
        H: Hasher,
        E: Context,
        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
        I: crate::index::Ordered<Value = Location<F>> + 'static,
        S: Strategy,
        Operation<F, update::Ordered<K, V>>: Codec,
    {
        type Family = F;
        type K = K;
        type V = V::Value;
        type Metadata = V::Value;
        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;

        fn write(self, key: K, value: Option<V::Value>) -> Self {
            Self::write(self, key, value)
        }

        async fn merkleize(
            self,
            db: &CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>,
            metadata: Option<V::Value>,
        ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
            self.merkleize(db, metadata).await
        }
    }

    impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
        MerkleizedBatchTrait for Arc<MerkleizedBatch<F, D, U, N, S>>
    where
        Operation<F, U>: Codec,
    {
        type Digest = D;

        fn root(&self) -> D {
            MerkleizedBatch::root(self)
        }
    }

    impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
        for CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>
    where
        F: Graftable,
        E: Context,
        K: Key,
        V: ValueEncoding + 'static,
        C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
        I: UnorderedIndex<Value = Location<F>> + 'static,
        H: Hasher,
        S: Strategy,
        Operation<F, update::Unordered<K, V>>: Codec,
    {
        type Family = F;
        type K = K;
        type V = V::Value;
        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;
        type Batch = UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>;

        fn new_batch(&self) -> Self::Batch {
            self.new_batch()
        }

        fn apply_batch(
            self,
            batch: Self::Merkleized,
        ) -> impl Future<Output = ApplyBatchResult<Self>> {
            self.apply_batch(batch)
        }
    }

    impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
        for CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>
    where
        F: Graftable,
        E: Context,
        K: Key,
        V: ValueEncoding + 'static,
        C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
        I: crate::index::Ordered<Value = Location<F>> + 'static,
        H: Hasher,
        S: Strategy,
        Operation<F, update::Ordered<K, V>>: Codec,
    {
        type Family = F;
        type K = K;
        type V = V::Value;
        type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;
        type Batch = UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>;

        fn new_batch(&self) -> Self::Batch {
            self.new_batch()
        }

        fn apply_batch(
            self,
            batch: Self::Merkleized,
        ) -> impl Future<Output = ApplyBatchResult<Self>> {
            self.apply_batch(batch)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{mmb, mmr, utils::detached::block_strategy};
    use commonware_cryptography::Sha256;
    use commonware_macros::test_traced;
    use commonware_parallel::{Manual, Rayon};
    use commonware_utils::{NZUsize, bitmap::Prunable as BitMap};
    use std::{
        future::Future as _,
        task::Context as TaskContext,
        time::{Duration, Instant},
    };

    // N=4 -> CHUNK_SIZE_BITS = 32
    const N: usize = 4;
    type Bm = BitMap<N>;
    type GraftedBatch =
        Arc<GenericMerkleizedBatch<mmb::Family, <Sha256 as Hasher>::Digest, Manual<Rayon>>>;
    type Location = mmr::Location;

    fn make_bitmap(bits: &[bool]) -> Bm {
        let mut bm = Bm::new();
        for &b in bits {
            bm.push(b);
        }
        bm
    }

    fn grafted_chain(
        strategy: &Manual<Rayon>,
        mem: &Arc<Mem<mmb::Family, <Sha256 as Hasher>::Digest>>,
    ) -> (GraftedBatch, GraftedBatch) {
        let hasher = grafting::hasher::<mmb::Family, Sha256>(grafting::height::<1>());
        let a = mem
            .new_batch_with_strategy(strategy.clone())
            .add_leaf_digest(Sha256::hash(&[b"a-0"]))
            .add_leaf_digest(Sha256::hash(&[b"a-1"]))
            .merkleize(mem, &hasher);
        let b = a
            .new_batch()
            .add_leaf_digest(Sha256::hash(&[b"b-0"]))
            .merkleize(mem, &hasher);
        (a, b)
    }

    /// A detached grafted-tree merkleization owns the full ancestor chain after cancellation.
    #[test_traced]
    fn test_grafted_merkleize_retains_ancestors_after_cancellation() {
        let strategy = Rayon::new(NZUsize!(2)).unwrap();
        let manual = strategy.manual();
        let mem = Arc::new(Mem::<mmb::Family, <Sha256 as Hasher>::Digest>::new());
        let grafting_height = grafting::height::<1>();
        let graft_inputs = || vec![(0, Sha256::hash(&[b"replacement"]), [1u8; 1])];
        let waker = futures::task::noop_waker();
        let mut context = TaskContext::from_waker(&waker);

        // Observe the worker result so a missing grandparent fails the test directly.
        let (a, b) = grafted_chain(&manual, &mem);
        let ancestor = Arc::downgrade(&a);
        let release = block_strategy(&strategy, 2);
        let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
            &manual,
            Arc::clone(&b),
            &mem,
            graft_inputs(),
            grafting_height,
        ));
        assert!(merkleize.as_mut().poll(&mut context).is_pending());
        drop(b);
        drop(a);
        drop(release);
        let _ = futures::executor::block_on(merkleize);
        assert!(ancestor.upgrade().is_none());

        // Drop the waiter while the worker is queued to prove the guard moved with it.
        let (a, b) = grafted_chain(&manual, &mem);
        let ancestor = Arc::downgrade(&a);
        let release = block_strategy(&strategy, 2);
        let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
            &manual,
            Arc::clone(&b),
            &mem,
            graft_inputs(),
            grafting_height,
        ));
        assert!(merkleize.as_mut().poll(&mut context).is_pending());
        drop(merkleize);
        drop(b);
        drop(a);
        assert!(ancestor.upgrade().is_some());

        drop(release);
        let deadline = Instant::now() + Duration::from_secs(10);
        while ancestor.upgrade().is_some() {
            assert!(
                Instant::now() < deadline,
                "detached grafted merkleization did not release its ancestors"
            );
            std::thread::yield_now();
        }
    }

    // ---- build_chunk_overlay tests ----

    #[test]
    fn chunk_overlay_pushes() {
        use crate::qmdb::any::value::FixedEncoding;
        use commonware_utils::sequence::FixedBytes;

        type K = FixedBytes<4>;
        type V = FixedEncoding<u64>;
        type U = crate::qmdb::any::operation::update::Unordered<K, V>;

        let key1 = FixedBytes::from([1, 0, 0, 0]);
        let key2 = FixedBytes::from([2, 0, 0, 0]);

        // Base: 4 bits, all set (previous commit at loc 3).
        // Segment of 4 operations starting at base_size=4.
        // Diff: key1 active at loc=4 (in batch), key2 active at loc=99 (not in batch,
        // so superseded within this batch).
        let base = make_bitmap(&[true; 4]);
        let mut diff = vec![
            (
                key1,
                DiffEntry::Active {
                    value: 100u64,
                    loc: Location::new(4), // offset 0 in batch
                    base_old_loc: None,
                },
            ),
            (
                key2,
                DiffEntry::Active {
                    value: 200u64,
                    loc: Location::new(99), // not in batch [4,8), so superseded
                    base_old_loc: None,
                },
            ),
        ];
        diff.sort_by(|a, b| a.0.cmp(&b.0));

        let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 4, 4, &diff, &[]);

        // Chunk 0 should have: bits 0-3 from base (all set), bit 4 set (key1), bits 5-6 false
        // (inactive), bit 7 set (CommitFloor at loc 7). Also bit 3 cleared (previous commit).
        let c0 = overlay.get(0).expect("chunk 0 should be dirty");
        assert_ne!(c0[0] & (1 << 4), 0); // key1 active
        assert_eq!(c0[0] & (1 << 5), 0); // inactive
        assert_eq!(c0[0] & (1 << 6), 0); // inactive
        assert_ne!(c0[0] & (1 << 7), 0); // CommitFloor
        assert_eq!(c0[0] & (1 << 3), 0); // previous commit cleared
    }

    #[test]
    fn chunk_overlay_clears() {
        use crate::qmdb::any::value::FixedEncoding;
        use commonware_utils::sequence::FixedBytes;

        type K = FixedBytes<4>;
        type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;

        let key1 = FixedBytes::from([1, 0, 0, 0]);
        let key2 = FixedBytes::from([2, 0, 0, 0]);
        let key3 = FixedBytes::from([3, 0, 0, 0]);

        // Base bitmap with 64 bits, all set.
        let base = make_bitmap(&[true; 64]);

        let mut diff: Vec<(K, DiffEntry<mmr::Family, u64>)> = vec![
            (
                key1,
                DiffEntry::Active {
                    value: 100,
                    loc: Location::new(70),
                    base_old_loc: Some(Location::new(5)),
                },
            ),
            (
                key2,
                DiffEntry::Deleted {
                    base_old_loc: Some(Location::new(10)),
                },
            ),
            (
                key3,
                DiffEntry::Active {
                    value: 300,
                    loc: Location::new(71),
                    base_old_loc: None,
                },
            ),
        ];
        diff.sort_by(|a, b| a.0.cmp(&b.0));

        // Segment of 8 ops starting at 64; previous commit at loc 63.
        let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 8, 64, &diff, &[]);

        // Verify bits 5 and 10 are cleared in chunk 0.
        let c0 = overlay.get(0).expect("chunk 0 should be dirty");
        assert_eq!(c0[0] & (1 << 5), 0); // bit 5 cleared
        assert_eq!(c0[1] & (1 << 2), 0); // bit 10 = byte 1, bit 2 cleared

        // Other bits should still be set.
        assert_eq!(c0[0] & (1 << 4), 1 << 4); // bit 4 still set
        assert_eq!(c0[1] & (1 << 3), 1 << 3); // bit 11 still set
    }

    /// Regression: when the parent bitmap has a partial last chunk that becomes complete in the
    /// child (without any active bits landing in that chunk), the overlay must inherit the parent's
    /// partial chunk data, not zero it out.
    #[test]
    fn chunk_overlay_preserves_partial_parent_chunk() {
        use crate::qmdb::any::value::FixedEncoding;
        use commonware_utils::sequence::FixedBytes;

        type K = FixedBytes<4>;
        type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;

        // Base: 20 bits set (partial chunk 0, CHUNK_SIZE_BITS=32).
        let base = make_bitmap(&[true; 20]);
        assert_eq!(base.complete_chunks(), 0); // partial

        // Segment of 20 ops starting at loc 20. This pushes total to 40 bits, completing chunk 0
        // (32 bits) and starting chunk 1. Diff: only one active key at loc 35 (in chunk 1), plus
        // CommitFloor at loc 39. No active bits land in chunk 0's new region (bits 20-31).
        let key1 = FixedBytes::from([1, 0, 0, 0]);
        let mut diff = vec![(
            key1,
            DiffEntry::Active {
                value: 42u64,
                loc: Location::new(35),
                base_old_loc: None,
            },
        )];
        diff.sort_by(|a, b| a.0.cmp(&b.0));

        let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 20, 20, &diff, &[]);

        // Chunk 0 should be materialized and preserve the parent's first 20 bits.
        let c0 = overlay.get(0).expect("chunk 0 should be in overlay");
        // Bits 0-7 all set -> byte 0 = 0xFF
        assert_eq!(c0[0], 0xFF);
        // Bits 8-15 all set -> byte 1 = 0xFF
        assert_eq!(c0[1], 0xFF);
        // Bits 16-18 set, bit 19 cleared (previous commit), 20-23 not set -> byte 2 = 0x07
        assert_eq!(c0[2], 0x07);
    }

    // ---- next_candidate tests ----

    /// Single-step oracle for [`fill_candidates`]: return the next floor-raise candidate in
    /// `[floor, tip)` over any [`bitmap::Readable`]. `fill_candidates_matches_oracle` proves
    /// the production scan produces exactly this sequence over every chain shape.
    fn next_candidate<B: bitmap::Readable<N2>, const N2: usize>(
        bitmap: &B,
        floor: Location,
        tip: u64,
    ) -> Option<Location> {
        let floor = *floor;
        let bitmap_len = bitmap.len();
        let committed_end = bitmap_len.min(tip);
        if floor < committed_end
            && let Some(idx) = bitmap.ones_iter_from(floor).next()
            && idx < committed_end
        {
            return Some(Location::new(idx));
        }
        let candidate = floor.max(bitmap_len);
        (candidate < tip).then(|| Location::new(candidate))
    }

    #[test]
    fn bitmap_scan_all_active() {
        let bm = make_bitmap(&[true; 8]);
        for i in 0..8 {
            assert_eq!(
                next_candidate(&bm, Location::new(i), 8),
                Some(Location::new(i))
            );
        }
        assert_eq!(next_candidate(&bm, Location::new(8), 8), None);
    }

    #[test]
    fn bitmap_scan_all_inactive() {
        let bm = make_bitmap(&[false; 8]);
        assert_eq!(next_candidate(&bm, Location::new(0), 8), None);
    }

    #[test]
    fn bitmap_scan_skips_inactive() {
        // Pattern: inactive, inactive, active, inactive, active
        let bm = make_bitmap(&[false, false, true, false, true]);
        assert_eq!(
            next_candidate(&bm, Location::new(0), 5),
            Some(Location::new(2))
        );
        assert_eq!(
            next_candidate(&bm, Location::new(3), 5),
            Some(Location::new(4))
        );
        assert_eq!(next_candidate(&bm, Location::new(5), 5), None);
    }

    #[test]
    fn bitmap_scan_beyond_bitmap_len_returns_candidate() {
        // Bitmap has 4 bits, but tip is 8. Locations 4..8 are beyond the bitmap and should be
        // returned as candidates.
        let bm = make_bitmap(&[false; 4]);
        // All bitmap bits are unset, so 0..4 are skipped; loc 4 is beyond bitmap -> candidate.
        assert_eq!(
            next_candidate(&bm, Location::new(0), 8),
            Some(Location::new(4))
        );
        assert_eq!(
            next_candidate(&bm, Location::new(6), 8),
            Some(Location::new(6))
        );
    }

    #[test]
    fn bitmap_scan_respects_tip() {
        let bm = make_bitmap(&[false, false, false, true]);
        // Active bit at 3, but tip is 3 so it's excluded.
        assert_eq!(next_candidate(&bm, Location::new(0), 3), None);
        // With tip=4, bit 3 is included.
        assert_eq!(
            next_candidate(&bm, Location::new(0), 4),
            Some(Location::new(3))
        );
    }

    #[test]
    fn bitmap_scan_floor_at_tip() {
        let bm = make_bitmap(&[true; 4]);
        assert_eq!(next_candidate(&bm, Location::new(4), 4), None);
    }

    #[test]
    fn bitmap_scan_empty_bitmap() {
        let bm = Bm::new();
        // Empty bitmap, but tip > 0: all locations are beyond bitmap.
        assert_eq!(
            next_candidate(&bm, Location::new(0), 5),
            Some(Location::new(0))
        );
        // Empty bitmap, tip = 0: no candidates.
        assert_eq!(next_candidate(&bm, Location::new(0), 0), None);
    }

    #[test]
    fn fill_candidates_matches_oracle() {
        // Sequence parity plus split-resume for one (chain, tip): the scan matches
        // single-stepping the oracle over the same chain, and any split point resumes
        // seamlessly via the returned continuation.
        fn assert_matches(name: &str, chain: &BitmapBatch<N>, tip: u64) {
            for floor in 0..=tip {
                let mut want = Vec::new();
                let mut scan = Location::new(floor);
                while let Some(c) = next_candidate(chain, scan, tip) {
                    want.push(c);
                    scan = c + 1;
                }
                for split in 0..=want.len() {
                    let mut got = Vec::new();
                    let next = fill_candidates(chain, Location::new(floor), tip, split, &mut got);
                    fill_candidates(chain, next, tip, want.len() + 1, &mut got);
                    assert_eq!(got, want, "{name} floor={floor} split={split}");
                }
            }
        }

        let bits = [true, false, true, true, false, false, true, false];
        let base = make_bitmap(&bits);

        // Flat committed base.
        let flat = BitmapBatch::Base(Arc::new(Shared::new(make_bitmap(&bits))));

        // One layer: clears committed bits 3 and 6, appends 8..12 (only 9 set).
        let shared = Arc::new(Shared::new(make_bitmap(&bits)));
        let mut overlay = ChunkOverlay::new(&base, 12, 1);
        overlay.clear_bit(&base, 3);
        overlay.clear_bit(&base, 6);
        overlay.set_bit(&base, 9);
        let one_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
            parent: BitmapBatch::Base(Arc::clone(&shared)),
            overlay: Arc::new(overlay),
            shared,
        }));

        // Two layers, mirroring `fill_candidates_filters_ancestor_clears`.
        let shared = Arc::new(Shared::new(make_bitmap(&bits)));
        let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
        overlay1.clear_bit(&base, 3);
        overlay1.set_bit(&base, 9);
        let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
            parent: BitmapBatch::Base(Arc::clone(&shared)),
            overlay: Arc::new(overlay1),
            shared: Arc::clone(&shared),
        }));
        let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
        overlay2.clear_bit(&chain1, 6);
        overlay2.clear_bit(&chain1, 9);
        overlay2.set_bit(&chain1, 13);
        let two_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
            parent: chain1,
            overlay: Arc::new(overlay2),
            shared,
        }));

        // Pruned base: 40 bits with chunk 0 pruned (33 and 38 set beyond the pruned
        // boundary), plus a layer clearing 38 and appending 40..46 (41 and 44 set).
        let make_pruned = || {
            let mut bits = [false; 40];
            bits[33] = true;
            bits[38] = true;
            let mut bm = make_bitmap(&bits);
            bm.prune_to_bit(32);
            bm
        };
        let pruned_base = make_pruned();
        let shared = Arc::new(Shared::new(make_pruned()));
        let mut overlay = ChunkOverlay::new(&pruned_base, 46, 1);
        overlay.clear_bit(&pruned_base, 38);
        overlay.set_bit(&pruned_base, 41);
        overlay.set_bit(&pruned_base, 44);
        let pruned = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
            parent: BitmapBatch::Base(Arc::clone(&shared)),
            overlay: Arc::new(overlay),
            shared,
        }));

        for (name, chain, committed) in [
            ("flat", flat, 8),
            ("one-layer", one_layer, 8),
            ("two-layer", two_layer, 8),
            ("pruned-base", pruned, 40),
        ] {
            let len = bitmap::Readable::<N>::len(&chain);
            for tip in [committed, len, len + 3] {
                assert_matches(name, &chain, tip);
            }

            // Prefetch-then-live handoff: the prefetch is clamped to the committed
            // boundary and the live scan resumes from the continuation with the
            // post-batch tip. Nothing the raise must revalidate may be lost across the
            // handoff (false negatives are forbidden): every set bit in `[floor, len)`
            // and every location in `[len, tip)`.
            let tip = len + 3;
            let cap = tip as usize;
            let pruned_bits = bitmap::Readable::<N>::pruned_bits(&chain);
            for floor in pruned_bits..=committed {
                let mut got = Vec::new();
                let next = fill_candidates(&chain, Location::new(floor), committed, cap, &mut got);
                fill_candidates(&chain, next, tip, cap, &mut got);
                assert!(got.is_sorted_by(|a, b| a < b), "{name} floor={floor}");
                for loc in floor..tip {
                    let must_emit = loc >= len || bitmap::Readable::<N>::get_bit(&chain, loc);
                    assert!(
                        !must_emit || got.contains(&Location::new(loc)),
                        "{name} floor={floor} lost {loc}"
                    );
                }
            }
        }
    }

    #[test]
    fn fill_candidates_filters_ancestor_clears() {
        let bits = [true, false, true, true, false, false, true, false];
        let base = make_bitmap(&bits);
        let shared = Arc::new(Shared::new(make_bitmap(&bits)));

        // Layer 1 clears committed bit 3 and appends bits 8..12 (only 9 set).
        let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
        overlay1.clear_bit(&base, 3);
        overlay1.set_bit(&base, 9);
        let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
            parent: BitmapBatch::Base(Arc::clone(&shared)),
            overlay: Arc::new(overlay1),
            shared: Arc::clone(&shared),
        }));

        // Layer 2 (materialized against the layer-1 view, as `build_chunk_overlay` does)
        // clears bits 6 and 9, and appends bits 12..14 (only 13 set).
        let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
        overlay2.clear_bit(&chain1, 6);
        overlay2.clear_bit(&chain1, 9);
        overlay2.set_bit(&chain1, 13);
        let chain2 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
            parent: chain1.clone(),
            overlay: Arc::new(overlay2),
            shared,
        }));

        // Bits cleared by any layer are skipped (no wasted log reads), set bits -- committed
        // or appended, from whichever layer materialized the chunk last -- are emitted
        // ascending, and locations at or beyond the layered length up to `tip` are emitted
        // sequentially.
        let scan = |chain: &BitmapBatch<N>, tip: u64| {
            let mut got = Vec::new();
            fill_candidates(chain, Location::new(0), tip, 16, &mut got);
            got
        };
        let want = |locs: &[u64]| locs.iter().copied().map(Location::new).collect::<Vec<_>>();
        assert_eq!(scan(&chain1, 12), want(&[0, 2, 6, 9]));
        assert_eq!(scan(&chain2, 14), want(&[0, 2, 13]));
        assert_eq!(scan(&chain2, 16), want(&[0, 2, 13, 14, 15]));
    }

    #[test]
    fn fill_candidates_mixes_overlay_and_base_chunks() {
        // Base spans two chunks (N=4 -> 32-bit chunks): full chunk 0 plus a partial chunk 1.
        let mut bits = [false; 40];
        for i in [1, 30, 33, 35, 38] {
            bits[i] = true;
        }
        let base = make_bitmap(&bits);
        let shared = Arc::new(Shared::new(make_bitmap(&bits)));

        // Layer touches only chunk 1: clears committed bit 35 and appends bits 40..44
        // (only 41 set). Chunk 0 stays unmaterialized, so the scan must fall through to
        // the committed base there.
        let mut overlay = ChunkOverlay::new(&base, 44, 1);
        overlay.clear_bit(&base, 35);
        overlay.set_bit(&base, 41);
        let chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
            parent: BitmapBatch::Base(Arc::clone(&shared)),
            overlay: Arc::new(overlay),
            shared,
        }));

        // Chunk 0 bits come from the base, chunk 1 bits from the overlay (35 filtered,
        // the appended 41 emitted).
        let mut got = Vec::new();
        fill_candidates(&chain, Location::new(0), 44, 16, &mut got);
        let want: Vec<Location> = [1, 30, 33, 38, 41].into_iter().map(Location::new).collect();
        assert_eq!(got, want);
    }

    // ---- trim_committed tests ----
    //
    // `trim_committed` is called from `MerkleizedBatch::new_batch` to strip any `Layer`s whose
    // overlays have already been absorbed into the shared committed bitmap by a prior apply.
    // The implementation is a single loop that collects uncommitted overlays top-down and
    // rebuilds a fresh chain rooted at `Base`. These tests cover distinct input shapes directly,
    // without going through the full Db/batch machinery, so the function's structural output
    // can be asserted.

    /// Build a chain `Base(shared) -> Layer(len=L1) -> Layer(len=L2) -> ...` from a list of
    /// overlay lengths (bottom to top). Each constructed `Layer` caches `shared` per the
    /// struct's invariant.
    fn make_chain(shared: &Arc<Shared<N>>, overlay_lens: &[u64]) -> BitmapBatch<N> {
        let mut chain = BitmapBatch::Base(Arc::clone(shared));
        for &len in overlay_lens {
            let overlay = Arc::new(ChunkOverlay::new(&chain, len, 0));
            chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
                parent: chain,
                overlay,
                shared: Arc::clone(shared),
            }));
        }
        chain
    }

    /// Walk a chain and return its overlay lengths in bottom-to-top order. Used to assert the
    /// structural output of `trim_committed` without touching private fields. Panics if the
    /// chain isn't terminated by a single `Base` at the bottom.
    fn chain_overlays(batch: &BitmapBatch<N>) -> Vec<u64> {
        let mut lens = Vec::new();
        let mut current = batch;
        while let BitmapBatch::Layer(layer) = current {
            lens.push(layer.overlay.len);
            current = &layer.parent;
        }
        assert!(matches!(current, BitmapBatch::Base(_)));
        lens.reverse();
        lens
    }

    /// Input is already a bare `Base` with no speculative layers on top -- the loop body never
    /// runs, `kept` stays empty, and the result is a freshly constructed `Base` pointing at the
    /// same `Shared`. Real-world trigger: `MerkleizedBatch::new_batch` on a batch whose
    /// chain was previously trimmed flat (e.g., immediately after an apply collapsed everything).
    #[test]
    fn trim_committed_already_base() {
        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
        let base = BitmapBatch::Base(Arc::clone(&shared));
        let result = base.trim_committed();
        // Still `Base`, pointing at the same shared terminal.
        match result {
            BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
            BitmapBatch::Layer(_) => panic!("expected Base"),
        }
    }

    /// Every layer has been absorbed by prior applies -- the loop breaks on the first iteration
    /// and `kept` stays empty, so the result is a bare `Base`. This is the steady-state
    /// "extend a just-applied batch" flow: after `apply_batch(A)`, `A`'s own layer has
    /// `overlay.len == committed` and the next `new_batch` call should start from a clean
    /// terminal.
    #[test]
    fn trim_committed_all_committed() {
        // `shared.len() == 64`; the single layer's `overlay.len == 32 (<= 64)`, so it's committed.
        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
        let chain = make_chain(&shared, &[32]);
        let result = chain.trim_committed();
        // Collapsed to a bare Base, pointing at the original shared.
        match result {
            BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
            BitmapBatch::Layer(_) => panic!("expected Base after full trim"),
        }
    }

    /// Every layer is still speculative -- the loop walks all the way to `Base` without
    /// breaking, and `kept` holds every overlay. The rebuilt chain is structurally equivalent
    /// to the input (same overlay lens, same shared terminal). Real-world trigger: speculating
    /// multiple batches deep (A, then B off A, then C off B) without `apply_batch` in between.
    #[test]
    fn trim_committed_none_committed() {
        // `shared.len() == 32`; both overlays have `len > 32`, so neither is committed.
        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 32])));
        let chain = make_chain(&shared, &[64, 96]);
        let result = chain.trim_committed();
        // Structure must be preserved in bottom-to-top order.
        assert_eq!(chain_overlays(&result), vec![64, 96]);
    }

    /// Exactly one layer is uncommitted (the newest) on top of a committed prefix -- the
    /// dominant pattern in chained growth. The loop collects the one uncommitted overlay, and
    /// the rebuild produces `Layer(Base, overlay_B)`. Also verifies the rebuilt layer carries
    /// the cached `shared` reference correctly. Real-world trigger: apply parent A, then B
    /// held alive off A, then `B.new_batch()` to build C.
    #[test]
    fn trim_committed_exactly_one_uncommitted() {
        // `shared.len() == 64`; committed layer (`overlay.len == 64`) + uncommitted (`96`).
        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
        let chain = make_chain(&shared, &[64, 96]);
        let result = chain.trim_committed();
        // The committed layer is gone; only the uncommitted overlay remains.
        assert_eq!(chain_overlays(&result), vec![96]);
        // And the rebuilt layer's `shared` field still points at the original terminal.
        assert!(Arc::ptr_eq(result.shared(), &shared));
    }

    /// Two or more uncommitted layers on top of a committed prefix -- exercises the loop's
    /// iterated `kept.push` and the rebuild's iterated `Arc::new(BitmapBatchLayer)`, including
    /// the cached `shared` wire-through on every reconstructed layer. Real-world trigger:
    /// build A, then B off A, then C off B; apply only A; then call `C.new_batch()`.
    #[test]
    fn trim_committed_multiple_uncommitted() {
        // `shared.len() == 64`; committed layer (64), then two uncommitted (96, 128).
        let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
        let chain = make_chain(&shared, &[64, 96, 128]);
        let result = chain.trim_committed();
        // Committed layer dropped; uncommitted pair preserved in order.
        assert_eq!(chain_overlays(&result), vec![96, 128]);
        // Every reconstructed layer must still cache the original shared terminal.
        assert!(Arc::ptr_eq(result.shared(), &shared));
    }
}