commonware-storage 2026.10.0

//! Batch mutation API for Immutable QMDBs.

use super::Immutable;
use crate::{
    Context,
    journal::{authenticated, contiguous::Mutable},
    merkle::{Family, Location, Proof},
    qmdb::{
        Error,
        any::{ValueEncoding, batch::lookup_sorted},
        chain::{self, Bounds, Commitment, Compatible},
        immutable::operation::Operation,
        operation::Key,
    },
    translator::Translator,
};
use commonware_codec::EncodeShared;
use commonware_cryptography::{Digest, Hasher};
use commonware_parallel::Strategy;
use commonware_utils::iter::zip_eq;
use std::{
    collections::BTreeMap,
    iter,
    sync::{Arc, Weak},
};

type DiffVec<K, F, V> = Vec<(K, DiffEntry<F, V>)>;

/// What happened to a key in this batch.
#[derive(Clone)]
pub(crate) struct DiffEntry<F: Family, V> {
    pub(crate) value: V,
    pub(crate) loc: Location<F>,
}

/// Look up `key` in `diff`, ignoring an entry below the inactivity `floor`.
///
/// A reopened database indexes only operations at or above its floor, so batch reads skip
/// entries below the floor to answer alike on every database at the same commitment.
fn lookup_active<'a, K: Ord, F: Family, V>(
    diff: &'a [(K, DiffEntry<F, V>)],
    key: &K,
    floor: Location<F>,
) -> Option<&'a DiffEntry<F, V>> {
    lookup_sorted(diff, key).filter(|entry| entry.loc >= floor)
}

/// Result of merkleizing a batch.
type MerkleizeResult<F, D, K, V, S> = Result<Arc<MerkleizedBatch<F, D, K, V, S>>, Error<F>>;

/// A speculative batch of operations whose root digest has not yet been computed, in contrast
/// to [`MerkleizedBatch`].
///
/// Consuming [`UnmerkleizedBatch::merkleize`] produces a merkleized batch.
/// Methods that need the DB (e.g. [`get`](Self::get)) accept it as a parameter.
#[allow(clippy::type_complexity)]
pub struct UnmerkleizedBatch<F, H, K, V, S: Strategy>
where
    F: Family,
    K: Key,
    V: ValueEncoding,
    H: Hasher,
{
    /// Authenticated journal batch for computing the speculative Merkle root.
    journal_batch: authenticated::UnmerkleizedBatch<F, H, Operation<F, K, V>, S>,

    /// Pending mutations.
    mutations: BTreeMap<K, V::Value>,

    /// Parent batch in the chain. `None` for batches created directly from the DB.
    parent: Option<Arc<MerkleizedBatch<F, H::Digest, K, V, S>>>,

    /// The state immediately before this batch's operations.
    /// This batch's i-th operation lands at location `base.size + i`.
    base: Commitment<F, H::Digest>,
}

/// Merkleized authenticated-journal batch wrapping an [`Operation`] payload.
type JournalBatch<F, D, K, V, S> = Arc<authenticated::MerkleizedBatch<F, D, Operation<F, K, V>, S>>;

/// A speculative batch of operations whose root digest has been computed, in contrast to
/// [`UnmerkleizedBatch`]. Reads through it return [`crate::qmdb::Error::StaleRead`] once the
/// database is at no state of its chain, for example after a sibling or one of its descendants
/// is applied (see [`crate::qmdb::chain`]).
#[derive(Clone)]
pub struct MerkleizedBatch<F: Family, D: Digest, K: Key, V: ValueEncoding, S: Strategy> {
    /// Authenticated journal batch (Merkle state + local items).
    pub(super) journal_batch: JournalBatch<F, D, K, V, S>,

    /// This batch's local key-level changes only (not accumulated from ancestors).
    /// Sorted by key with no duplicates; queried via `lookup_sorted` (binary search).
    pub(super) diff: Arc<DiffVec<K, F, V::Value>>,

    /// The parent batch in the chain, if any.
    pub(super) parent: Option<Weak<Self>>,

    /// Arc refs to the diffs of the reachable ancestors, collected during `merkleize()`. Reads
    /// consult them (see `read_diffs`), and `apply_batch` applies the unapplied ones. 1:1 with
    /// `bounds.ancestors` (same length, same ordering).
    pub(super) ancestor_diffs: Vec<Arc<DiffVec<K, F, V::Value>>>,

    /// Position and floor bounds for this batch chain.
    pub(super) bounds: chain::Bounds<F, D>,
}

impl<F, H, K, V, S: Strategy> UnmerkleizedBatch<F, H, K, V, S>
where
    F: Family,
    K: Key,
    V: ValueEncoding,
    H: Hasher,
    Operation<F, K, V>: EncodeShared,
{
    /// Create a batch from the DB (no parent chain).
    pub(super) fn new<E, C, T>(
        immutable: &Immutable<F, E, K, V, C, H, T, S>,
        base: Commitment<F, H::Digest>,
    ) -> Self
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        T: Translator,
    {
        Self {
            journal_batch: immutable.journal.new_batch(),
            mutations: BTreeMap::new(),
            parent: None,
            base,
        }
    }

    /// The database boundary for this batch chain.
    ///
    /// A batch created from the database uses its base. A child inherits its parent's `db`.
    fn db(&self) -> Commitment<F, H::Digest> {
        self.parent
            .as_ref()
            .map_or(self.base, |parent| parent.bounds.db)
    }

    /// The inactivity floor this batch builds on: its parent's, or `db`'s for a batch with no
    /// parent.
    fn floor<E, C, T>(&self, db: &Immutable<F, E, K, V, C, H, T, S>) -> Location<F>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        T: Translator,
    {
        self.parent.as_ref().map_or_else(
            || db.inactivity_floor_loc(),
            |parent| parent.bounds.inactivity_floor,
        )
    }

    /// Check that `db` is at a state of this batch's chain and return the [`Compatible`]
    /// witness that database reads require.
    #[allow(clippy::type_complexity)]
    fn compatible<'a, E, C, T>(
        &self,
        db: &'a Immutable<F, E, K, V, C, H, T, S>,
    ) -> Result<Compatible<'a, Immutable<F, E, K, V, C, H, T, S>>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        T: Translator,
    {
        self.parent.as_ref().map_or_else(
            || self.base.compatible(db, db.commitment()),
            |parent| parent.bounds.compatible(db, db.commitment()),
        )
    }

    /// Set a key to a value.
    ///
    /// The key must not already be set in the database history or in any ancestor batch.
    pub fn set(mut self, key: K, value: V::Value) -> Self {
        self.mutations.insert(key, value);
        self
    }

    /// Read through: mutations -> ancestor diffs -> database.
    ///
    /// Only operations at or above the inactivity floor this batch builds on are read, so every
    /// database at the same commitment answers alike, however far it has pruned.
    ///
    /// # Errors
    ///
    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
    pub async fn get<E, C, T>(
        &self,
        key: &K,
        db: &Immutable<F, E, K, V, C, H, T, S>,
    ) -> Result<Option<V::Value>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        T: Translator,
    {
        let db = self.compatible(db)?;
        // Check this batch's pending mutations.
        if let Some(value) = self.mutations.get(key) {
            return Ok(Some(value.clone()));
        }
        let floor = self.floor(&*db);
        if let Some(parent) = self.parent.as_ref() {
            for diff in parent.read_diffs(db.bounds().end) {
                if let Some(entry) = lookup_active(diff, key, floor) {
                    return Ok(Some(entry.value.clone()));
                }
            }
        }
        // Fall through to base DB.
        db.get_from(key, floor).await
    }

    /// Batch read multiple keys.
    ///
    /// Returns results in the same order as the input keys, matching [`Self::get`] for each key.
    ///
    /// # Errors
    ///
    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
    pub async fn get_many<E, C, T>(
        &self,
        keys: &[&K],
        db: &Immutable<F, E, K, V, C, H, T, S>,
    ) -> Result<Vec<Option<V::Value>>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        T: Translator,
    {
        let db = self.compatible(db)?;
        if keys.is_empty() {
            return Ok(Vec::new());
        }
        let floor = self.floor(&*db);

        let mut results: Vec<Option<V::Value>> = Vec::with_capacity(keys.len());
        let mut db_indices = Vec::new();
        let mut db_keys = Vec::new();

        for (i, key) in keys.iter().enumerate() {
            // Check local mutations.
            if let Some(value) = self.mutations.get(*key) {
                results.push(Some(value.clone()));
                continue;
            }

            // Check the parent's retained diff chain.
            if let Some(entry) = self.parent.as_ref().and_then(|parent| {
                parent
                    .read_diffs(db.bounds().end)
                    .find_map(|diff| lookup_active(diff, *key, floor))
            }) {
                results.push(Some(entry.value.clone()));
                continue;
            }

            // Need DB fallthrough.
            db_indices.push(i);
            db_keys.push(*key);
            results.push(None);
        }

        if !db_keys.is_empty() {
            let db_results = db.get_many_from(&db_keys, floor).await?;
            for (slot, value) in zip_eq(db_indices, db_results) {
                results[slot] = value;
            }
        }

        Ok(results)
    }

    /// Resolve mutations into operations and return a merkleized batch.
    ///
    /// `inactivity_floor` declares that all operations before this location are inactive. It must
    /// be at least the floor this batch builds on (its parent's, or the database's) and at most
    /// this batch's own commit location.
    ///
    /// # Errors
    ///
    /// - Returns [`Error::StaleBatch`] if `db` is not at a state of this batch's chain or an
    ///   unapplied ancestor was dropped.
    /// - Returns [`Error::FloorRegressed`] if `inactivity_floor` is below the floor this batch
    ///   builds on, and [`Error::FloorBeyondSize`] if it is past the commit location.
    #[tracing::instrument(name = "qmdb.immutable.batch.merkleize", level = "info", skip_all)]
    pub async fn merkleize<E, C, T>(
        self,
        db: &Immutable<F, E, K, V, C, H, T, S>,
        metadata: Option<V::Value>,
        inactivity_floor: Location<F>,
    ) -> MerkleizeResult<F, H::Digest, K, V, S>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        T: Translator,
    {
        let base = self.base.size;

        let live_ancestors: Vec<_> =
            chain::parent_and_ancestors(self.parent.as_ref(), |parent| parent.ancestors())
                .collect();
        let boundary = chain::effective_boundary(
            self.db(),
            live_ancestors.last().map(|oldest| oldest.bounds.base),
        );

        // Compute the batch chain bounds.
        let mut ancestor_diffs = Vec::new();
        let mut ancestors = Vec::new();
        for batch in &live_ancestors {
            ancestor_diffs.push(Arc::clone(&batch.diff));
            ancestors.push(batch.commitment());
        }
        let db = chain::merkleizable(db, db.commitment(), boundary, ancestors.iter().copied())?;
        let start_floor = self.floor(&*db);

        // Build operations: one Set per key, then Commit. `self.mutations` is a BTreeMap, so
        // iteration yields keys in sorted order, which `diff` relies on for binary search.
        let mut ops: Vec<Operation<F, K, V>> = Vec::with_capacity(self.mutations.len() + 1);
        let mut diff: DiffVec<K, F, V::Value> = Vec::with_capacity(self.mutations.len());

        for (key, value) in self.mutations {
            let loc = base + ops.len() as u64;
            ops.push(Operation::Set(key.clone(), value.clone()));
            diff.push((key, DiffEntry { value, loc }));
        }

        ops.push(Operation::Commit(metadata, inactivity_floor));

        let total_size = base + ops.len() as u64;
        chain::validate_merkleize_floor::<F>(start_floor, inactivity_floor, total_size - 1)?;
        let inactive_peaks = F::inactive_peaks(total_size, inactivity_floor);

        // Leaf and node hashing dominate merkleization, so run them as one job through the
        // strategy (see `Journal::merkleize`).
        let (journal, root) = db
            .journal
            .merkleize(self.journal_batch, ops, inactive_peaks)
            .await?;

        // Hold ancestor batches until the journal has captured their operations and nodes.
        drop(live_ancestors);

        Ok(Arc::new(MerkleizedBatch {
            journal_batch: journal,
            diff: Arc::new(diff),
            parent: self.parent.as_ref().map(Arc::downgrade),
            ancestor_diffs,
            bounds: chain::Bounds {
                base: self.base,
                db: boundary,
                tip: Commitment::new(total_size, root),
                ancestors,
                inactivity_floor,
            },
        }))
    }
}

impl<F: Family, D: Digest, K: Key, V: ValueEncoding, S: Strategy> MerkleizedBatch<F, D, K, V, S>
where
    Operation<F, K, V>: EncodeShared,
{
    /// Return the speculative root.
    pub const fn root(&self) -> D {
        self.bounds.tip.root
    }

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

    /// Return the operations this batch appends to the log and the location of the first.
    #[allow(clippy::type_complexity)]
    pub fn operations(&self) -> (Location<F>, Arc<Vec<Operation<F, K, V>>>) {
        (
            self.bounds.base.size,
            Arc::clone(self.journal_batch.items()),
        )
    }

    /// Inclusion proof for the operations returned by [`Self::operations`], anchored at
    /// this batch's tip. The pair verifies against [`Self::root`] via
    /// [`crate::qmdb::verify_proof`]. Together with [`Self::pinned_nodes`] they verify via
    /// [`crate::qmdb::verify_proof_and_pinned_nodes`].
    ///
    /// Nodes of unapplied ancestors are read through the chain, so those ancestors must still be
    /// held. Nodes below the chain are read from `db`'s [Merkle store][crate::merkle::mem::Mem],
    /// which retains them at least until this batch's changes are flushed (by a commit or sync
    /// after apply).
    ///
    /// # Errors
    ///
    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain,
    /// [`crate::merkle::Error::ElementPruned`] if a required node has been pruned or belongs to a
    /// dropped unapplied ancestor, and [`crate::merkle::Error::Empty`] if the batch has no
    /// operations (such as a batch from [`Immutable::to_batch`]).
    pub fn proof<E, C, H, T>(
        &self,
        db: &Immutable<F, E, K, V, C, H, T, S>,
    ) -> Result<Proof<F, D>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        H: Hasher<Digest = D>,
        T: Translator,
    {
        let db = self.bounds.compatible(db, db.commitment())?;
        let inactive_peaks = F::inactive_peaks(self.bounds.tip.size, self.bounds.inactivity_floor);
        db.journal
            .speculative_proof(&self.journal_batch, inactive_peaks)
            .map_err(Into::into)
    }

    /// The Merkle frontier at the first operation returned by [`Self::operations`]
    /// ([`Family::nodes_to_pin`]), which lets a consumer holding only this batch's base rebuild
    /// compact state and replay the operations. The operations, [`Self::proof`], and pinned
    /// nodes verify against [`Self::root`] via [`crate::qmdb::verify_proof_and_pinned_nodes`].
    ///
    /// Nodes of unapplied ancestors are read through the chain, so those ancestors must still be
    /// held. Nodes below the chain are read from `db`'s [Merkle store][crate::merkle::mem::Mem],
    /// which retains them at least until this batch's changes are flushed (by a commit or sync
    /// after apply).
    ///
    /// # Errors
    ///
    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain, and
    /// [`crate::merkle::Error::ElementPruned`] if a required node has been pruned or belongs to a
    /// dropped unapplied ancestor.
    pub fn pinned_nodes<E, C, H, T>(
        &self,
        db: &Immutable<F, E, K, V, C, H, T, S>,
    ) -> Result<Vec<D>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        H: Hasher<Digest = D>,
        T: Translator,
    {
        let db = self.bounds.compatible(db, db.commitment())?;
        db.journal
            .speculative_pinned_nodes(&self.journal_batch)
            .map_err(Into::into)
    }

    /// Yield this batch's diff, then the ancestor diffs a read consults before a compatible
    /// database at `db_size` operations: each reachable or unapplied ancestor's, newest first.
    ///
    /// The walk ends at the first ancestor that is both unreachable and applied. The database
    /// holds that ancestor's writes and every older layer's. A reachable applied ancestor stays in
    /// the walk so its writes resolve from memory. The batch keeps every retained diff, because a
    /// bounded initialization can make those layers unapplied again.
    fn read_diffs(
        &self,
        db_size: Location<F>,
    ) -> impl Iterator<Item = &[(K, DiffEntry<F, V::Value>)]> {
        let mut reachable = iter::once_with(|| self.ancestors()).flatten();
        iter::once(self.diff.as_slice()).chain(
            self.ancestor_diffs
                .iter()
                .zip(&self.bounds.ancestors)
                .take_while(move |(_, state)| reachable.next().is_some() || state.size > db_size)
                .map(|(diff, _)| diff.as_slice()),
        )
    }

    /// Iterate over ancestor batches (parent first, then grandparent, etc.).
    pub(super) fn ancestors(&self) -> impl Iterator<Item = Arc<Self>> + use<F, D, K, V, S> {
        chain::ancestors(self.parent.clone(), |batch| batch.parent.as_ref())
    }

    /// The [`Commitment`] this batch commits to.
    pub(super) const fn commitment(&self) -> Commitment<F, D> {
        self.bounds.tip
    }

    /// Read through: local diff -> ancestor diffs -> database.
    ///
    /// Only operations at or above this batch's inactivity floor are read, so every database at
    /// the same commitment answers alike, however far it has pruned.
    ///
    /// # Errors
    ///
    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
    pub async fn get<E, C, H, T>(
        &self,
        key: &K,
        db: &Immutable<F, E, K, V, C, H, T, S>,
    ) -> Result<Option<V::Value>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        H: Hasher<Digest = D>,
        T: Translator,
    {
        let db = self.bounds.compatible(db, db.commitment())?;
        let floor = self.bounds.inactivity_floor;
        for diff in self.read_diffs(db.bounds().end) {
            if let Some(entry) = lookup_active(diff, key, floor) {
                return Ok(Some(entry.value.clone()));
            }
        }
        db.get_from(key, floor).await
    }

    /// Batch read multiple keys.
    ///
    /// Returns results in the same order as the input keys, matching [`Self::get`] for each key.
    ///
    /// # Errors
    ///
    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
    pub async fn get_many<E, C, H, T>(
        &self,
        keys: &[&K],
        db: &Immutable<F, E, K, V, C, H, T, S>,
    ) -> Result<Vec<Option<V::Value>>, Error<F>>
    where
        E: Context,
        C: Mutable<Item = Operation<F, K, V>>,
        C::Item: EncodeShared,
        H: Hasher<Digest = D>,
        T: Translator,
    {
        let db = self.bounds.compatible(db, db.commitment())?;
        if keys.is_empty() {
            return Ok(Vec::new());
        }
        let floor = self.bounds.inactivity_floor;

        let mut results: Vec<Option<V::Value>> = Vec::with_capacity(keys.len());
        let mut db_indices = Vec::new();
        let mut db_keys = Vec::new();

        for (i, key) in keys.iter().enumerate() {
            if let Some(entry) = self
                .read_diffs(db.bounds().end)
                .find_map(|diff| lookup_active(diff, *key, floor))
            {
                results.push(Some(entry.value.clone()));
                continue;
            }

            // Need DB fallthrough.
            db_indices.push(i);
            db_keys.push(*key);
            results.push(None);
        }

        if !db_keys.is_empty() {
            let db_results = db.get_many_from(&db_keys, floor).await?;
            for (slot, value) in zip_eq(db_indices, db_results) {
                results[slot] = value;
            }
        }

        Ok(results)
    }

    /// Create a new speculative batch of operations with this batch as its parent.
    ///
    /// Every unapplied ancestor in the chain must be held until the child (or any descendant) is
    /// merkleized. Otherwise, `merkleize` returns [`Error::StaleBatch`].
    ///
    /// A child can be created from a stale parent. While the database is at none of the chain's
    /// states, the child's reads return [`Error::StaleRead`], and merkleizing it returns
    /// [`Error::StaleBatch`].
    pub fn new_batch<H>(self: &Arc<Self>) -> UnmerkleizedBatch<F, H, K, V, S>
    where
        H: Hasher<Digest = D>,
    {
        UnmerkleizedBatch {
            journal_batch: self.journal_batch.new_batch::<H>(),
            mutations: BTreeMap::new(),
            parent: Some(Arc::clone(self)),
            base: self.commitment(),
        }
    }
}

impl<F, E, K, V, C, H, T, S> Immutable<F, E, K, V, C, H, T, S>
where
    F: Family,
    E: Context,
    K: Key,
    V: ValueEncoding,
    C: Mutable<Item = Operation<F, K, V>>,
    C::Item: EncodeShared,
    H: Hasher,
    T: Translator,
    S: Strategy,
{
    /// Create an initial [`MerkleizedBatch`] from the DB's current state.
    pub fn to_batch(&self) -> Arc<MerkleizedBatch<F, H::Digest, K, V, S>> {
        Arc::new(MerkleizedBatch {
            journal_batch: self.journal.to_merkleized_batch(),
            diff: Arc::new(Vec::new()),
            parent: None,
            ancestor_diffs: Vec::new(),
            bounds: chain::Bounds::from_db(self.commitment(), self.inactivity_floor_loc),
        })
    }
}