Skip to main content

commonware_storage/qmdb/current/
db.rs

1//! A shared, generic implementation of the _Current_ QMDB.
2//!
3//! The impl blocks in this file define shared functionality across all Current QMDB variants.
4
5use crate::{
6    index::Unordered as UnorderedIndex,
7    journal::{
8        contiguous::{Contiguous, Mutable},
9        Error as JournalError,
10    },
11    merkle::{
12        self, hasher::Hasher as _, mem::Mem, storage::Storage as MerkleStorage, Graftable,
13        Location, Position,
14    },
15    metadata::{Config as MConfig, Metadata},
16    qmdb::{
17        self,
18        any::{
19            self,
20            operation::{update::Update, Operation},
21        },
22        current::{
23            batch::BitmapBatch,
24            grafting,
25            proof::{OperationProof, OpsRootWitness, RangeProof, RangeProofSpec},
26        },
27        operation::Operation as _,
28        Error,
29    },
30    Context,
31};
32use commonware_codec::{Codec, CodecShared, DecodeExt};
33use commonware_cryptography::{Digest, DigestOf, Hasher};
34use commonware_macros::boxed;
35use commonware_parallel::Strategy;
36use commonware_runtime::telemetry::metrics::{
37    histogram::{ScopedTimer, Timed},
38    Counter, Gauge, GaugeExt as _, MetricsExt as _,
39};
40use commonware_utils::{
41    bitmap::{self, Readable as _},
42    sequence::prefixed_u64::U64,
43};
44use core::{num::NonZeroU64, ops::Range};
45use futures::future::try_join_all;
46use std::{collections::BTreeMap, sync::Arc};
47use tracing::{error, warn};
48
49/// Prefix used for the metadata key for grafted tree pinned nodes.
50const NODE_PREFIX: u8 = 0;
51
52/// Prefix used for the metadata key for the number of pruned bitmap chunks.
53const PRUNED_CHUNKS_PREFIX: u8 = 1;
54
55/// Metrics for the Current layer.
56pub(crate) struct Metrics<E: Context> {
57    /// Pruned bitmap chunks.
58    pruned_chunks: Gauge,
59    /// Most recent safe sync/prune boundary location.
60    sync_boundary: Gauge,
61    /// Current-layer apply-batch calls.
62    pub apply_batch_calls: Counter,
63    /// Duration of Current-layer apply-batch calls.
64    apply_batch_duration: Timed,
65    /// Current-layer sync calls.
66    pub sync_calls: Counter,
67    /// Duration of Current-layer sync calls.
68    sync_duration: Timed,
69    /// Current-layer prune calls.
70    pub prune_calls: Counter,
71    /// Duration of Current-layer prune calls.
72    prune_duration: Timed,
73    /// Clock used by the duration timers.
74    clock: Arc<E>,
75}
76
77impl<E: Context> Metrics<E> {
78    /// Register the full metric set under `context`, retaining it as the timers' clock.
79    pub fn new(context: E) -> Self {
80        Self {
81            pruned_chunks: context.gauge("pruned_chunks", "Number of pruned bitmap chunks"),
82            sync_boundary: context
83                .gauge("sync_boundary", "Most recent safe sync boundary location"),
84            apply_batch_calls: context.counter("apply_batch_calls", "Number of apply-batch calls"),
85            apply_batch_duration: Timed::register(
86                &context,
87                "apply_batch_duration",
88                "Duration of apply-batch calls",
89            ),
90            sync_calls: context.counter("sync_calls", "Number of sync calls"),
91            sync_duration: Timed::register(&context, "sync_duration", "Duration of sync calls"),
92            prune_calls: context.counter("prune_calls", "Number of prune calls"),
93            prune_duration: Timed::register(&context, "prune_duration", "Duration of prune calls"),
94            clock: Arc::new(context),
95        }
96    }
97
98    pub fn apply_batch_timer(&self) -> ScopedTimer<E> {
99        self.apply_batch_duration.scoped(&self.clock)
100    }
101
102    pub fn sync_timer(&self) -> ScopedTimer<E> {
103        self.sync_duration.scoped(&self.clock)
104    }
105
106    pub fn prune_timer(&self) -> ScopedTimer<E> {
107        self.prune_duration.scoped(&self.clock)
108    }
109
110    /// Update Current-specific state gauges.
111    pub fn update(&self, pruned_chunks: u64, sync_boundary: u64) {
112        let _ = self.pruned_chunks.try_set(pruned_chunks);
113        let _ = self.sync_boundary.try_set(sync_boundary);
114    }
115}
116
117/// A Current QMDB implementation generic over ordered/unordered keys and variable/fixed values.
118pub struct Db<
119    F: merkle::Graftable,
120    E: Context,
121    C: Contiguous<Item: CodecShared>,
122    I: UnorderedIndex<Value = Location<F>>,
123    H: Hasher,
124    U: Send + Sync,
125    const N: usize,
126    S: Strategy,
127> {
128    /// An authenticated database that provides the ability to prove whether a key ever had a
129    /// specific value. Owns the activity-status bitmap (`any.bitmap`) that this layer reads to
130    /// install grafted-tree updates and serve proofs.
131    pub(super) any: any::db::Db<F, E, C, I, H, U, N, S>,
132
133    /// Each leaf corresponds to a complete bitmap chunk at the grafting height.
134    /// See the [grafted leaf formula](super) in the module documentation.
135    ///
136    /// Internal nodes are hashed using their position in the ops tree rather than their
137    /// grafted position.
138    ///
139    /// Held in an [`Arc`] so merkleize can hand a zero-copy, immutable snapshot to the
140    /// grafted-layer hashing job running off the calling task. Mutations go through
141    /// [`Arc::make_mut`]: they are in-place while no snapshot is alive and copy-on-write
142    /// otherwise, so a snapshot never observes later mutations.
143    pub(super) grafted_tree: Arc<Mem<F, H::Digest>>,
144
145    /// Persists:
146    /// - The number of pruned bitmap chunks at key [PRUNED_CHUNKS_PREFIX]
147    /// - The grafted tree pinned nodes at key [NODE_PREFIX]
148    pub(super) metadata: Metadata<E, U64, Vec<u8>>,
149
150    /// Strategy used to parallelize batch operations across the ops tree, the grafted tree,
151    /// and grafted leaf computation.
152    pub(super) strategy: S,
153
154    /// The cached canonical root.
155    /// See the [Root structure](super) section in the module documentation.
156    pub(super) root: DigestOf<H>,
157
158    /// Metrics for the Current layer.
159    pub(super) metrics: Metrics<E>,
160
161    /// Test-only: park [Self::prune] after the pruning-metadata sync, before the log prune,
162    /// so tests can drop the pending future at that exact point.
163    #[cfg(test)]
164    pub(super) halt_before_prune_log: bool,
165}
166
167// Shared read-only functionality.
168impl<F, E, C, I, H, U, const N: usize, S> Db<F, E, C, I, H, U, N, S>
169where
170    F: merkle::Graftable,
171    E: Context,
172    U: Update,
173    C: Contiguous<Item = Operation<F, U>>,
174    I: UnorderedIndex<Value = Location<F>>,
175    H: Hasher,
176    S: Strategy,
177    Operation<F, U>: Codec,
178{
179    /// Return the inactivity floor location. This is the location before which all operations are
180    /// known to be inactive.
181    #[cfg(any(test, feature = "test-traits"))]
182    pub(crate) const fn inactivity_floor_loc(&self) -> Location<F> {
183        self.any.inactivity_floor_loc()
184    }
185
186    /// Whether the snapshot currently has no active keys.
187    pub const fn is_empty(&self) -> bool {
188        self.any.is_empty()
189    }
190
191    /// Get the metadata associated with the last commit.
192    pub async fn get_metadata(&self) -> Result<Option<U::Value>, Error<F>> {
193        self.any.get_metadata().await
194    }
195
196    /// Return [start, end) where `start` and `end - 1` are the Locations of the oldest and newest
197    /// retained operations respectively.
198    pub fn bounds(&self) -> std::ops::Range<Location<F>> {
199        self.any.bounds()
200    }
201
202    /// Return true if the given sequence of `ops` were applied starting at location `start_loc`
203    /// in the log with the provided `root`, having the activity status described by `chunks`.
204    pub fn verify_range_proof(
205        proof: &RangeProof<F, H::Digest>,
206        start_loc: Location<F>,
207        ops: &[Operation<F, U>],
208        chunks: &[[u8; N]],
209        root: &H::Digest,
210    ) -> bool {
211        proof.verify::<H, _, N>(start_loc, ops, chunks, root)
212    }
213}
214
215// Functionality requiring non-mutable journal.
216impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
217where
218    F: merkle::Graftable,
219    E: Context,
220    U: Update,
221    C: Contiguous<Item = Operation<F, U>>,
222    I: UnorderedIndex<Value = Location<F>>,
223    H: Hasher,
224    S: Strategy,
225    Operation<F, U>: Codec,
226{
227    /// Returns a virtual [grafting::Storage] over the grafted tree and ops tree. For positions at
228    /// or above the grafting height, returns the grafted node. For positions below the grafting
229    /// height, the ops tree is used.
230    fn grafted_storage(&self) -> impl MerkleStorage<F, Digest = H::Digest> + '_ {
231        grafting::Storage::<F, H, _, _>::new(
232            &self.grafted_tree,
233            grafting::height::<N>(),
234            &self.any.log.merkle,
235        )
236    }
237
238    /// Returns the canonical root.
239    /// See the [Root structure](super) section in the module documentation.
240    pub const fn root(&self) -> H::Digest {
241        self.root
242    }
243
244    /// Return a reference to the merkleization strategy.
245    pub const fn strategy(&self) -> &S {
246        &self.strategy
247    }
248
249    /// Returns the ops tree root.
250    ///
251    /// This is the root of the raw operations log, without the activity bitmap. It is used as the
252    /// sync target because the sync engine verifies batches against the ops root, not the canonical
253    /// root.
254    ///
255    /// External consumers that receive a trusted canonical `current` root should use
256    /// [`Self::ops_root_witness`] to authenticate this ops root against it.
257    ///
258    /// See the [Root structure](super) section in the module documentation.
259    pub const fn ops_root(&self) -> H::Digest {
260        self.any.root()
261    }
262
263    /// Returns a witness that this database's canonical root commits to its ops root.
264    ///
265    /// This can be used to authenticate an ops root against a trusted canonical `current` root.
266    pub async fn ops_root_witness(&self) -> Result<OpsRootWitness<F, H::Digest>, Error<F>> {
267        let storage = self.grafted_storage();
268        let ops_size = storage.size();
269        let ops_leaves = Location::<F>::try_from(ops_size)?;
270        let grafted_root = compute_grafted_root::<F, H, _, _, N>(
271            self.any.bitmap.as_ref(),
272            &storage,
273            ops_leaves,
274            self.any.inactivity_floor_loc,
275        )
276        .await?;
277        let hasher = qmdb::hasher::<H>();
278        let partial_chunk = partial_chunk::<_, N>(self.any.bitmap.as_ref())
279            .map(|(chunk, next_bit)| (next_bit, hasher.digest(chunk.as_slice())));
280        let pending_chunk_digest: F::PendingChunk<H::Digest> = pending_chunk::<F, _, N>(
281            self.any.bitmap.as_ref(),
282            ops_leaves,
283            grafting::height::<N>(),
284        )?
285        .map(|chunk| hasher.digest(chunk.as_slice()))
286        .try_into()
287        .expect("pending_chunk must be consistent with family");
288        Ok(OpsRootWitness {
289            grafted_root,
290            pending_chunk_digest,
291            partial_chunk,
292        })
293    }
294
295    /// Snapshot of the grafted tree for use in batch chains.
296    pub(super) fn grafted_snapshot(&self) -> Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>> {
297        merkle::batch::MerkleizedBatch::from_mem_with_strategy(
298            &self.grafted_tree,
299            self.strategy.clone(),
300        )
301    }
302
303    /// Create a new speculative batch of operations with this database as its parent.
304    pub fn new_batch(&self) -> super::batch::UnmerkleizedBatch<F, H, U, N, S> {
305        super::batch::UnmerkleizedBatch::new(
306            self.any.new_batch(),
307            self.grafted_snapshot(),
308            BitmapBatch::Base(Arc::clone(&self.any.bitmap)),
309        )
310    }
311
312    /// Returns a proof for the operation at `loc`.
313    pub(super) async fn operation_proof(
314        &self,
315        loc: Location<F>,
316    ) -> Result<OperationProof<F, H::Digest, N>, Error<F>> {
317        let storage = self.grafted_storage();
318        let ops_root = self.any.root();
319        OperationProof::new::<H, _>(
320            self.any.bitmap.as_ref(),
321            &storage,
322            self.any.inactivity_floor_loc,
323            loc,
324            ops_root,
325        )
326        .await
327    }
328
329    /// Returns a proof that the specified range of operations are part of the database, along with
330    /// the operations from the range. A truncated range (from hitting the max) can be detected by
331    /// looking at the length of the returned operations vector. Also returns the bitmap chunks
332    /// required to verify the proof.
333    ///
334    /// # Errors
335    ///
336    /// Returns [Error::OperationPruned] if `start_loc` falls in a pruned bitmap chunk. Returns
337    /// [`crate::merkle::Error::LocationOverflow`] if `start_loc` >
338    /// [`crate::merkle::Family::MAX_LEAVES`]. Returns [`crate::merkle::Error::RangeOutOfBounds`] if
339    /// `start_loc` >= number of leaves in the tree.
340    #[allow(clippy::type_complexity)]
341    #[tracing::instrument(
342        name = "qmdb.current.db.range_proof",
343        level = "info",
344        skip_all,
345        fields(
346            start_loc = *start_loc,
347            max_ops = max_ops.get(),
348        ),
349    )]
350    pub async fn range_proof(
351        &self,
352        start_loc: Location<F>,
353        max_ops: NonZeroU64,
354    ) -> Result<(RangeProof<F, H::Digest>, Vec<Operation<F, U>>, Vec<[u8; N]>), Error<F>> {
355        let storage = self.grafted_storage();
356        let ops_root = self.any.root();
357        RangeProof::new_with_ops::<H, _, _, N>(
358            self.any.bitmap.as_ref(),
359            &storage,
360            &self.any.log,
361            RangeProofSpec {
362                start_loc,
363                max_ops,
364                inactivity_floor: self.any.inactivity_floor_loc,
365                ops_root,
366            },
367        )
368        .await
369    }
370}
371
372// Functionality requiring mutable journal.
373impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
374where
375    F: merkle::Graftable,
376    E: Context,
377    U: Update,
378    C: Mutable<Item = Operation<F, U>>,
379    I: UnorderedIndex<Value = Location<F>>,
380    H: Hasher,
381    S: Strategy,
382    Operation<F, U>: Codec,
383{
384    /// Returns an ops-level historical proof for the specified range.
385    ///
386    /// Unlike [`range_proof`](Self::range_proof) which returns grafted proofs incorporating the
387    /// activity bitmap, this returns ops-tree Merkle proofs suitable for state sync. Direct
388    /// verifiers should use [`crate::qmdb::verify_proof`].
389    pub async fn ops_historical_proof(
390        &self,
391        historical_size: Location<F>,
392        start_loc: Location<F>,
393        max_ops: NonZeroU64,
394    ) -> Result<(merkle::Proof<F, H::Digest>, Vec<Operation<F, U>>), Error<F>> {
395        self.any
396            .historical_proof(historical_size, start_loc, max_ops)
397            .await
398    }
399
400    /// Return the pinned nodes for a lower operation boundary of `loc`.
401    pub async fn pinned_nodes_at(&self, loc: Location<F>) -> Result<Vec<H::Digest>, Error<F>> {
402        self.any.pinned_nodes_at(loc).await
403    }
404
405    /// Returns the most recent location from which this database can safely be synced, and the
406    /// upper bound on [`Self::prune`]'s `prune_loc`.
407    ///
408    /// Callers constructing a sync [`Target`](crate::qmdb::sync::Target) may use this value, or
409    /// any earlier retained location, as `range.start`. Values *above* this boundary are unsafe:
410    /// the receiver's grafted-pin derivation requires absorption-settled state for every fully
411    /// pruned chunk, which this value guarantees.
412    ///
413    /// # Computation
414    ///
415    /// Starts from the inactivity floor (the most chunks we could possibly prune) and walks
416    /// backward until two conditions hold for the youngest chunk that would be pruned:
417    ///
418    /// 1. **Settled**: the chunk's ops subtree root at height `gh` has been born in the ops
419    ///    tree (its `peak_birth_size <= ops_leaves`).
420    ///
421    /// 2. **Absorbed**: the chunk-pair parent at height `gh+1` has been born. This guarantees
422    ///    that the ops tree has no individual height-`gh` peaks for pruned chunks, so
423    ///    `compute_grafted_root` never queries a discarded grafted leaf.
424    ///
425    /// Because older chunk-pairs have strictly earlier birth times, checking only the youngest
426    /// pair is sufficient: if the youngest pair's parent is born, all older pairs' parents are
427    /// too. In the worst case the loop decrements twice (once past the unsettled chunk, once
428    /// to land on the older pair boundary).
429    ///
430    /// For families without delayed merges (e.g. MMR), `peak_birth_size` at height `gh` equals
431    /// the chunk's last leaf, so condition (1) always holds and the function returns the
432    /// inactivity floor rounded down to the nearest chunk boundary.
433    pub fn sync_boundary(&self) -> Location<F> {
434        sync_boundary::<F, N>(
435            *self.any.inactivity_floor_loc / bitmap::Prunable::<N>::CHUNK_SIZE_BITS,
436            *self.any.last_commit_loc + 1,
437        )
438    }
439
440    /// Update Current-specific state gauges.
441    pub(super) fn update_metrics(&self) {
442        self.metrics.update(
443            self.any.bitmap.pruned_chunks() as u64,
444            *self.sync_boundary(),
445        );
446    }
447
448    /// Returns the minimum rewind target that keeps delayed-merge grafting queries valid
449    /// for the current bitmap pruning boundary.
450    ///
451    /// This is the same absorption threshold used by [`Self::sync_boundary`]: the
452    /// `peak_birth_size` of the youngest pruned chunk-pair's height-(gh+1) parent.
453    /// Rewinding below this size would put the ops tree in a state where the parent has not
454    /// been born, re-exposing individual height-`gh` ops peaks for pruned chunks whose
455    /// grafted leaves are no longer available.
456    ///
457    /// Returns `None` for families without delayed merges.
458    fn delayed_merge_rewind_floor(&self) -> Option<u64> {
459        pair_absorption_threshold::<F, N>(self.any.bitmap.pruned_chunks() as u64)
460    }
461
462    /// Prune the grafted tree to match the committed bitmap's pruned chunks.
463    fn prune_grafted_tree_to_bitmap(&mut self) -> Result<(), Error<F>> {
464        let pruned_chunks = self.any.bitmap.pruned_chunks() as u64;
465        if pruned_chunks == 0 {
466            return Ok(());
467        }
468
469        let prune_loc = Location::<F>::new(pruned_chunks);
470        if prune_loc <= self.grafted_tree.bounds().start {
471            return Ok(());
472        }
473
474        let prune_pos = Position::try_from(prune_loc)
475            .map_err(|_| Error::<F>::DataCorrupted("prune location overflow"))?;
476        let size = self.grafted_tree.size();
477
478        let mut pinned = BTreeMap::new();
479        for pos in F::nodes_to_pin(prune_loc) {
480            let digest = self
481                .grafted_tree
482                .get_node(pos)
483                .ok_or(Error::<F>::DataCorrupted("missing grafted pinned node"))?;
484            pinned.insert(pos, digest);
485        }
486
487        let mut retained = Vec::with_capacity((*size - *prune_pos) as usize);
488        for p in *prune_pos..*size {
489            let digest = self
490                .grafted_tree
491                .get_node(Position::new(p))
492                .ok_or(Error::<F>::DataCorrupted("missing retained grafted node"))?;
493            retained.push(digest);
494        }
495
496        self.grafted_tree = Arc::new(Mem::from_pruned_with_retained(prune_pos, pinned, retained));
497        Ok(())
498    }
499
500    /// Prunes historical operations prior to `prune_loc`. This does not affect the db's root or
501    /// snapshot.
502    ///
503    /// `prune` requires no prior commit. After a crash, the database remains recoverable;
504    /// uncommitted operations are not guaranteed to survive.
505    ///
506    /// `prune_loc` must be at most [`Self::sync_boundary`]: the ops log's lower bound must not
507    /// advance past the point where the grafting overlay has been pruned. The bitmap and grafted
508    /// tree advance to the sync boundary regardless of `prune_loc`.
509    ///
510    /// # Errors
511    ///
512    /// - Returns [Error::PruneBeyondMinRequired] if `prune_loc` > [`Self::sync_boundary`].
513    /// - Returns [`crate::merkle::Error::LocationOverflow`] if `prune_loc` >
514    ///   [crate::merkle::Family::MAX_LEAVES].
515    /// - Returns [Error::DataCorrupted] if internal grafted-tree state is inconsistent (a pinned
516    ///   or retained node is missing, or the prune location overflows a [Position]).
517    #[tracing::instrument(name = "qmdb.current.db.prune", level = "info", skip_all)]
518    pub async fn prune(&mut self, prune_loc: Location<F>) -> Result<(), Error<F>> {
519        let _timer = self.metrics.prune_timer();
520        self.metrics.prune_calls.inc();
521        let sync_boundary = self.sync_boundary();
522        if prune_loc > sync_boundary {
523            return Err(Error::PruneBeyondMinRequired(prune_loc, sync_boundary));
524        }
525
526        // The sync boundary may be advanced by applied-but-uncommitted operations, and the
527        // pruning metadata persisted below durably records it. Commit the log first so
528        // recovery can replay to that boundary: otherwise a crash before the log prune
529        // recovers the older durable floor alongside newer pruning metadata and fails to
530        // initialize the bitmap.
531        self.any.log.commit().await?;
532
533        // Prune the bitmap to the sync boundary (most aggressive safe location).
534        self.any.prune_bitmap(sync_boundary);
535        self.prune_grafted_tree_to_bitmap()?;
536
537        // Persist grafted tree pruning state before pruning the ops log. If the subsequent
538        // `any.prune_log` fails, the metadata is ahead of the log, which is safe: on recovery,
539        // `build_grafted_tree` will recompute from the (un-pruned) log and the metadata
540        // simply records peaks that haven't been pruned yet. The reverse order would be unsafe:
541        // a pruned log with stale metadata would lose peak digests permanently.
542        self.sync_metadata().await?;
543
544        #[cfg(test)]
545        if self.halt_before_prune_log {
546            std::future::pending::<()>().await;
547        }
548
549        self.any.prune_log(prune_loc).await?;
550        self.any.update_metrics();
551        self.update_metrics();
552        Ok(())
553    }
554
555    /// Rewind the database to `size` operations, where `size` is the location of the next append.
556    ///
557    /// This rewinds the underlying Any database and rebuilds the Current overlay state (bitmap,
558    /// grafted tree, and canonical root) for the rewound size.
559    ///
560    /// # Errors
561    ///
562    /// Returns an error when:
563    /// - `size` is not a valid rewind target
564    /// - the target's required logical range is not fully retained (for Current, this includes the
565    ///   underlying Any inactivity-floor boundary and bitmap pruning boundary)
566    /// - `size - 1` is not a commit operation
567    /// - `size` is below the bitmap pruning boundary
568    ///
569    /// Any error from this method is fatal for this handle. Rewind may mutate state in the
570    /// underlying Any database before this Current overlay finishes rebuilding. Callers must drop
571    /// this database handle after any `Err` from `rewind` and reopen from storage.
572    ///
573    /// A successful rewind is not restart-stable until a subsequent [`Db::commit`] or
574    /// [`Db::sync`].
575    #[tracing::instrument(name = "qmdb.current.db.rewind", level = "info", skip_all)]
576    pub async fn rewind(&mut self, size: Location<F>) -> Result<(), Error<F>> {
577        let rewind_size = *size;
578        let current_size = *self.any.last_commit_loc + 1;
579        // No-op short-circuit. Avoids the post-rewind grafted-tree rebuild and the validation
580        // and journal-read overhead below. Validation runs after this on the non-no-op path.
581        if rewind_size == current_size {
582            return Ok(());
583        }
584        // Reject zero / out-of-range up front: lines below compute `rewind_size - 1`, which
585        // underflows when `rewind_size == 0`. `any::Db::rewind` would catch these, but it isn't
586        // called until after those subtractions.
587        if rewind_size == 0 || rewind_size > current_size {
588            return Err(Error::Journal(JournalError::InvalidRewind(rewind_size)));
589        }
590
591        let pruned_chunks = self.any.bitmap.pruned_chunks();
592        let pruned_bits = (pruned_chunks as u64)
593            .checked_mul(bitmap::Prunable::<N>::CHUNK_SIZE_BITS)
594            .ok_or_else(|| Error::DataCorrupted("pruned ops leaves overflow"))?;
595        if rewind_size < pruned_bits {
596            return Err(Error::Journal(JournalError::ItemPruned(rewind_size - 1)));
597        }
598        if let Some(rewind_floor) = self.delayed_merge_rewind_floor() {
599            if rewind_size < rewind_floor {
600                return Err(Error::Journal(JournalError::ItemPruned(rewind_size - 1)));
601            }
602        }
603
604        // Ensure the target commit's logical range is fully representable with the current
605        // bitmap pruning boundary. Even if the ops log still retains older entries, rewinding
606        // to a commit with floor below `pruned_bits` would require bitmap chunks we've already
607        // discarded.
608        {
609            let rewind_last_loc = Location::<F>::new(rewind_size - 1);
610            let rewind_last_op = self.any.log.read(*rewind_last_loc).await?;
611            let Some(rewind_floor) = rewind_last_op.has_floor() else {
612                return Err(Error::<F>::UnexpectedData(rewind_last_loc));
613            };
614            if *rewind_floor < pruned_bits {
615                return Err(Error::<F>::Journal(JournalError::ItemPruned(*rewind_floor)));
616            }
617        }
618
619        // Extract pinned nodes for the existing pruning boundary from the in-memory grafted tree.
620        let pinned_nodes = if pruned_chunks > 0 {
621            let grafted_leaves = Location::<F>::new(pruned_chunks as u64);
622            let mut pinned_nodes = Vec::new();
623            for pos in F::nodes_to_pin(grafted_leaves) {
624                let digest = self
625                    .grafted_tree
626                    .get_node(pos)
627                    .ok_or(Error::<F>::DataCorrupted("missing grafted pinned node"))?;
628                pinned_nodes.push(digest);
629            }
630            pinned_nodes
631        } else {
632            Vec::new()
633        };
634
635        // `any.rewind` rewinds the log and patches the shared bitmap (truncate + restore active
636        // bits + set the rewound tail's CommitFloor). Live pre-rewind batches must be dropped by
637        // the caller; reads through them now return inconsistent data.
638        self.any.rewind(size).await?;
639
640        let ops_size = self.any.log.merkle.size();
641        let ops_leaves = Location::<F>::try_from(ops_size)?;
642        let grafted_tree = build_grafted_tree::<F, H, S, N>(
643            self.any.bitmap.as_ref(),
644            &pinned_nodes,
645            &self.any.log.merkle,
646            ops_leaves,
647            &self.strategy,
648        )
649        .await?;
650        let storage = grafting::Storage::<F, H, _, _>::new(
651            &grafted_tree,
652            grafting::height::<N>(),
653            &self.any.log.merkle,
654        );
655        let partial_chunk = partial_chunk(self.any.bitmap.as_ref());
656        let ops_root = self.any.root();
657        let root = compute_db_root::<F, H, _, _, N>(
658            self.any.bitmap.as_ref(),
659            &storage,
660            ops_leaves,
661            partial_chunk,
662            self.any.inactivity_floor_loc,
663            &ops_root,
664        )
665        .await?;
666
667        self.grafted_tree = Arc::new(grafted_tree);
668        self.root = root;
669        self.update_metrics();
670
671        Ok(())
672    }
673
674    /// Sync the metadata to disk.
675    pub(crate) async fn sync_metadata(&mut self) -> Result<(), Error<F>> {
676        self.metadata.clear();
677
678        // Snapshot the pruning boundary under the read lock; the guard drops before any await.
679        let pruned_chunks_u64 = self.any.bitmap.pruned_chunks() as u64;
680
681        // Write the number of pruned chunks.
682        let key = U64::new(PRUNED_CHUNKS_PREFIX, 0);
683        self.metadata
684            .put(key, pruned_chunks_u64.to_be_bytes().to_vec());
685
686        // Write the pinned nodes of the grafted tree.
687        let pruned_chunks = Location::<F>::new(pruned_chunks_u64);
688        for (i, grafted_pos) in F::nodes_to_pin(pruned_chunks).enumerate() {
689            let digest = self
690                .grafted_tree
691                .get_node(grafted_pos)
692                .ok_or(Error::<F>::DataCorrupted("missing grafted pinned node"))?;
693            let key = U64::new(NODE_PREFIX, i as u64);
694            self.metadata.put(key, digest.to_vec());
695        }
696
697        self.metadata.sync().await?;
698
699        Ok(())
700    }
701}
702
703/// Compute the safe sync boundary from the chunk-aligned inactivity floor and the current
704/// ops-tree size.
705///
706/// `floor_chunks` is the inactivity floor expressed in bitmap chunks (`floor / CHUNK_SIZE_BITS`),
707/// not the number of physically pruned chunks. Shared by the live DB and speculative batch
708/// wrappers, which both derive it from the inactivity floor so they report the same range start.
709pub(crate) fn sync_boundary<F: Graftable, const N: usize>(
710    mut floor_chunks: u64,
711    ops_leaves: u64,
712) -> Location<F> {
713    let chunk_bits = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
714    let grafting_height = grafting::height::<N>();
715
716    while floor_chunks > 0 {
717        let required_ops = pair_absorption_threshold::<F, N>(floor_chunks).unwrap_or_else(|| {
718            let youngest_start = (floor_chunks - 1) * chunk_bits;
719            let pos = F::subtree_root_position(Location::<F>::new(youngest_start), grafting_height);
720            F::peak_birth_size(pos, grafting_height)
721        });
722
723        if ops_leaves >= required_ops {
724            break;
725        }
726        floor_chunks -= 1;
727    }
728
729    Location::new(floor_chunks * chunk_bits)
730}
731
732/// For the youngest of `chunk_count` chunks, return the `peak_birth_size` of its
733/// chunk-pair parent at height `gh+1`. Returns `None` for families without delayed merges
734/// (where `peak_birth_size` at height `gh` equals the chunk boundary).
735fn pair_absorption_threshold<F: Graftable, const N: usize>(chunk_count: u64) -> Option<u64> {
736    if chunk_count == 0 {
737        return None;
738    }
739
740    let grafting_height = grafting::height::<N>();
741    let youngest = chunk_count - 1;
742    let youngest_start = youngest << grafting_height;
743    let youngest_end = (youngest + 1) << grafting_height;
744    let youngest_pos =
745        F::subtree_root_position(Location::<F>::new(youngest_start), grafting_height);
746
747    if F::peak_birth_size(youngest_pos, grafting_height) <= youngest_end {
748        return None;
749    }
750
751    let pair_chunk = youngest & !1;
752    let pair_start = pair_chunk << grafting_height;
753    let pair_pos = F::subtree_root_position(Location::<F>::new(pair_start), grafting_height + 1);
754    Some(F::peak_birth_size(pair_pos, grafting_height + 1))
755}
756
757// Functionality requiring mutable + persistable journal.
758impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
759where
760    F: merkle::Graftable,
761    E: Context,
762    U: Update,
763    C: Mutable<Item = Operation<F, U>>,
764    I: UnorderedIndex<Value = Location<F>>,
765    H: Hasher,
766    S: Strategy,
767    Operation<F, U>: Codec,
768{
769    /// Durably commit the journal state published by prior [`Db::apply_batch`]
770    /// calls.
771    #[tracing::instrument(name = "qmdb.current.db.commit", level = "info", skip_all)]
772    pub async fn commit(&mut self) -> Result<(), Error<F>> {
773        self.any.commit().await
774    }
775
776    /// Sync all database state to disk.
777    #[tracing::instrument(name = "qmdb.current.db.sync", level = "info", skip_all)]
778    pub async fn sync(&mut self) -> Result<(), Error<F>> {
779        let _timer = self.metrics.sync_timer();
780        self.metrics.sync_calls.inc();
781        self.any.sync().await?;
782
783        // Write the bitmap pruning boundary to disk so that next startup doesn't have to
784        // re-Merkleize the inactive portion up to the inactivity floor.
785        self.sync_metadata().await?;
786        self.update_metrics();
787        Ok(())
788    }
789
790    /// Destroy the db, removing all data from disk.
791    #[boxed]
792    pub async fn destroy(self) -> Result<(), Error<F>> {
793        // Destructure before the await boundary to avoid stack growth from
794        // retaining the entire `self` in the future.
795        let Self { any, metadata, .. } = self;
796        metadata.destroy().await?;
797        any.destroy().await
798    }
799}
800
801impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
802where
803    F: merkle::Graftable,
804    E: Context,
805    U: Update + 'static,
806    C: Mutable<Item = Operation<F, U>>,
807    I: UnorderedIndex<Value = Location<F>>,
808    H: Hasher,
809    S: Strategy,
810    Operation<F, U>: Codec,
811{
812    /// Apply a batch to the database, returning the range of written operations.
813    ///
814    /// A batch is valid only if every batch applied to the database since this batch's
815    /// ancestor chain was created is an ancestor of this batch. Applying a batch from a
816    /// different fork returns [`Error::StaleBatch`] (see [`crate::qmdb::batch_chain`] for
817    /// more details).
818    ///
819    /// This publishes the batch to the in-memory Current view and appends it to the journal,
820    /// but does not durably persist it. Call [`Db::commit`] or [`Db::sync`] to guarantee
821    /// durability.
822    #[tracing::instrument(name = "qmdb.current.db.apply_batch", level = "info", skip_all)]
823    pub async fn apply_batch(
824        &mut self,
825        batch: Arc<super::batch::MerkleizedBatch<F, H::Digest, U, N, S>>,
826    ) -> Result<Range<Location<F>>, Error<F>> {
827        let _timer = self.metrics.apply_batch_timer();
828        self.metrics.apply_batch_calls.inc();
829        let range = self.any.apply_batch(Arc::clone(&batch.inner)).await?;
830        Arc::make_mut(&mut self.grafted_tree).apply_batch(&batch.grafted)?;
831        self.root = batch.canonical_root;
832        self.update_metrics();
833        Ok(range)
834    }
835}
836
837/// Returns `Some((last_chunk, next_bit))` if the bitmap has an incomplete trailing chunk, or
838/// `None` if all bits fall on complete chunk boundaries.
839pub(super) fn partial_chunk<B: bitmap::Readable<N>, const N: usize>(
840    bitmap: &B,
841) -> Option<([u8; N], u64)> {
842    let (last_chunk, next_bit) = bitmap.last_chunk();
843    if next_bit == bitmap::Prunable::<N>::CHUNK_SIZE_BITS {
844        None
845    } else {
846        Some((last_chunk, next_bit))
847    }
848}
849
850/// Return complete and graftable chunk counts, enforcing the pending and pruning invariants.
851///
852/// Returns [`Error::DataCorrupted`] if `bitmap` and `ops_leaves` imply more than one
853/// pending chunk, or if pruning has advanced past the graftable chunk boundary.
854fn graftable_chunk_window<F: merkle::Graftable, B: bitmap::Readable<N>, const N: usize>(
855    bitmap: &B,
856    ops_leaves: Location<F>,
857    grafting_height: u32,
858) -> Result<(u64, u64), Error<F>> {
859    let complete = bitmap.complete_chunks() as u64;
860    let graftable = grafting::graftable_chunks::<F>(*ops_leaves, grafting_height).min(complete);
861    let pending = complete - graftable;
862    if pending > 1 {
863        return Err(Error::DataCorrupted("multiple pending bitmap chunks"));
864    }
865
866    let pruned = bitmap.pruned_chunks() as u64;
867    if pruned > graftable {
868        return Err(Error::DataCorrupted(
869            "pruned chunks exceed graftable chunks",
870        ));
871    }
872
873    Ok((complete, graftable))
874}
875
876/// Returns the bytes of the "pending" chunk if the bitmap currently has one, else `None`.
877///
878/// A chunk is pending when its bits are fully written to the bitmap but its h=G ancestor
879/// has not yet been born in the ops tree. At most one chunk is ever in this state (the most
880/// recently completed one); see [`super::grafting::graftable_chunks`] for the structural
881/// argument.
882///
883/// The caller must pass a consistent snapshot of `ops_leaves` (the ops tree's leaf count)
884/// and the bitmap state. Both inputs are used to derive `graftable_chunks`; deriving them from
885/// independent snapshots can violate the pending-window or pruning invariants.
886///
887/// Returns [`Error::DataCorrupted`] when those invariants are violated.
888pub(super) fn pending_chunk<F: merkle::Graftable, B: bitmap::Readable<N>, const N: usize>(
889    bitmap: &B,
890    ops_leaves: Location<F>,
891    grafting_height: u32,
892) -> Result<Option<[u8; N]>, Error<F>> {
893    let (complete, graftable) =
894        graftable_chunk_window::<F, B, N>(bitmap, ops_leaves, grafting_height)?;
895    if complete - graftable != 1 {
896        return Ok(None);
897    }
898    Ok(Some(bitmap.get_chunk(graftable as usize)))
899}
900
901/// Compute the canonical root from the ops root, grafted tree root, and optional pending /
902/// partial chunk digests.
903///
904/// See [Canonical root structure](super::proof#canonical-root-structure) for the full layout.
905/// The pending and partial inputs are independent: either, both, or neither may be set, and
906/// pending precedes partial in hash order when both are present.
907///
908/// # Collision resistance
909///
910/// `pending` contributes `D` bytes when present; `partial` contributes `D + 8` bytes (`D` =
911/// digest size). Different fixed lengths, so the two cannot produce the same input bytes,
912/// even when their digests are identical. Collisions reduce to H.
913pub(super) fn combine_roots<H: Hasher>(
914    ops_root: &H::Digest,
915    grafted_root: &H::Digest,
916    pending: Option<&H::Digest>,
917    partial: Option<(u64, &H::Digest)>,
918) -> H::Digest {
919    let hasher = qmdb::hasher::<H>();
920    match (pending, partial) {
921        (None, None) => hasher.hash([ops_root.as_ref(), grafted_root.as_ref()]),
922        (Some(pe), None) => hasher.hash([ops_root.as_ref(), grafted_root.as_ref(), pe.as_ref()]),
923        (None, Some((nb, p))) => {
924            let nb_bytes = nb.to_be_bytes();
925            hasher.hash([
926                ops_root.as_ref(),
927                grafted_root.as_ref(),
928                nb_bytes.as_slice(),
929                p.as_ref(),
930            ])
931        }
932        (Some(pe), Some((nb, p))) => {
933            let nb_bytes = nb.to_be_bytes();
934            hasher.hash([
935                ops_root.as_ref(),
936                grafted_root.as_ref(),
937                pe.as_ref(),
938                nb_bytes.as_slice(),
939                p.as_ref(),
940            ])
941        }
942    }
943}
944
945/// Compute the canonical root digest of a [Db].
946///
947/// See the [Root structure](super) section in the module documentation.
948///
949/// `ops_leaves` must be a single consistent snapshot of the ops tree's leaf count, taken
950/// in the same logical instant as the bitmap state passed via `status`. Both the pending
951/// chunk derivation and `compute_grafted_root` use this value to compute `graftable_chunks`;
952/// deriving them from independent snapshots risks the inconsistent state where a chunk is
953/// counted in one path but not the other.
954#[allow(clippy::too_many_arguments)]
955pub(super) async fn compute_db_root<
956    F: merkle::Graftable,
957    H: Hasher,
958    B: bitmap::Readable<N>,
959    S: MerkleStorage<F, Digest = H::Digest>,
960    const N: usize,
961>(
962    status: &B,
963    storage: &S,
964    ops_leaves: Location<F>,
965    partial_chunk: Option<([u8; N], u64)>,
966    inactivity_floor: Location<F>,
967    ops_root: &H::Digest,
968) -> Result<H::Digest, Error<F>> {
969    let grafted_root =
970        compute_grafted_root::<F, H, B, S, N>(status, storage, ops_leaves, inactivity_floor)
971            .await?;
972    let hasher = qmdb::hasher::<H>();
973    let pending = pending_chunk::<F, B, N>(status, ops_leaves, grafting::height::<N>())?
974        .map(|chunk| hasher.digest(&chunk));
975    let partial = partial_chunk.map(|(chunk, next_bit)| {
976        let digest = hasher.digest(&chunk);
977        (next_bit, digest)
978    });
979    Ok(combine_roots::<H>(
980        ops_root,
981        &grafted_root,
982        pending.as_ref(),
983        partial.as_ref().map(|(nb, d)| (*nb, d)),
984    ))
985}
986
987/// Compute the root of the grafted structure represented by `storage`.
988///
989/// Only **graftable** chunks (those whose h=G ancestor has been born in the ops tree) are
990/// committed by the grafted tree. The most recently completed but ungraftable chunk, if
991/// any, is hashed into the canonical root directly by [`combine_roots`] as the pending
992/// chunk, not by this function.
993///
994/// `ops_leaves` must come from the same single snapshot as `status` to preserve the
995/// `pruned_chunks <= graftable_chunks <= complete_chunks` invariant.
996pub(super) async fn compute_grafted_root<
997    F: merkle::Graftable,
998    H: Hasher,
999    B: bitmap::Readable<N>,
1000    S: MerkleStorage<F, Digest = H::Digest>,
1001    const N: usize,
1002>(
1003    status: &B,
1004    storage: &S,
1005    ops_leaves: Location<F>,
1006    inactivity_floor: Location<F>,
1007) -> Result<H::Digest, Error<F>> {
1008    let size = storage.size();
1009    let leaves = Location::try_from(size)?;
1010
1011    // Collect peak digests of the grafted structure.
1012    let mut peaks: Vec<H::Digest> = Vec::new();
1013    for (peak_pos, _) in F::peaks(size) {
1014        let digest = storage
1015            .get_node(peak_pos)
1016            .await?
1017            .ok_or_else(|| merkle::Error::<F>::MissingNode(peak_pos))?;
1018        peaks.push(digest);
1019    }
1020
1021    // Validate bitmap invariants (pending <= 1, pruned <= graftable).
1022    let grafting_height = grafting::height::<N>();
1023    let (_complete_chunks, _graftable_chunks) =
1024        graftable_chunk_window::<F, B, N>(status, ops_leaves, grafting_height)?;
1025
1026    let inactive_peaks =
1027        grafting::chunk_aligned_inactive_peaks::<F>(leaves, inactivity_floor, grafting_height)?;
1028    let hasher = qmdb::hasher::<H>();
1029
1030    // Every peak the storage layer surfaces is either a grafted-tree node (graftable chunks already
1031    // incorporate `hash(chunk || h_G_node)`), an ops node above G (hashed normally), or an ops node
1032    // below G (raw, because its chunk is pending and its digest is hashed directly into the
1033    // canonical root rather than through the tree). Bagging is a straight fold; no per-chunk
1034    // transformation is needed.
1035    Ok(hasher.root(leaves, inactive_peaks, peaks.iter())?)
1036}
1037
1038/// Resolve each bitmap chunk's covering ops-tree node, returning
1039/// `(chunk_idx, chunk_ops_digest, chunk)` triples ready for
1040/// [`grafting::graft_chunk_digests`].
1041///
1042/// Callers must pass only **graftable** chunks (those whose h=G ancestor has already been born in
1043/// the ops tree). Each graftable chunk has exactly one covering ops node at height G, looked up via
1044/// [`merkle::Graftable::subtree_root_position`].
1045pub(super) async fn read_graft_inputs<F: merkle::Graftable, D: Digest, const N: usize>(
1046    ops_tree: &impl MerkleStorage<F, Digest = D>,
1047    chunks: impl IntoIterator<Item = (usize, [u8; N])>,
1048) -> Result<Vec<(usize, D, [u8; N])>, Error<F>> {
1049    let grafting_height = grafting::height::<N>();
1050
1051    // Each graftable chunk has a single h=G ancestor at the deterministic
1052    // `subtree_root_position(chunk_idx << G, G)`. Look it up directly.
1053    try_join_all(chunks.into_iter().map(|(chunk_idx, chunk)| async move {
1054        let leaf_start = Location::<F>::new((chunk_idx as u64) << grafting_height);
1055        let pos = F::subtree_root_position(leaf_start, grafting_height);
1056        let chunk_ops_digest = ops_tree
1057            .get_node(pos)
1058            .await?
1059            .ok_or(merkle::Error::<F>::MissingGraftedLeaf(pos))?;
1060        Ok::<_, Error<F>>((chunk_idx, chunk_ops_digest, chunk))
1061    }))
1062    .await
1063}
1064
1065/// Compute grafted leaf digests for the given bitmap chunks as `(chunk_idx, digest)` pairs.
1066///
1067/// See [`read_graft_inputs`] for the chunk requirements. The grafted leaf digest is `hash(chunk ||
1068/// ops_h_G_node)`; for all-zero chunks the grafted leaf equals the ops digest directly (zero-chunk
1069/// identity).
1070///
1071/// The provided strategy determines if or how to parallelize merkleization.
1072pub(super) async fn compute_grafted_leaves<
1073    F: merkle::Graftable,
1074    H: Hasher,
1075    S: Strategy,
1076    const N: usize,
1077>(
1078    ops_tree: &impl MerkleStorage<F, Digest = H::Digest>,
1079    chunks: impl IntoIterator<Item = (usize, [u8; N])>,
1080    strategy: &S,
1081) -> Result<Vec<(usize, H::Digest)>, Error<F>> {
1082    let inputs = read_graft_inputs::<F, _, N>(ops_tree, chunks).await?;
1083    Ok(grafting::graft_chunk_digests::<H, _, N>(strategy, inputs))
1084}
1085
1086/// Build a grafted [Mem] from scratch using bitmap chunks and the ops tree.
1087///
1088/// For each non-pruned **graftable** chunk (index in `pruned_chunks..graftable_chunks`), reads the
1089/// ops tree node at the grafting height to compute the grafted leaf (see the
1090/// [grafted leaf formula](super) in the module documentation).
1091///
1092/// The most recently completed chunk may not yet be graftable (its h=G ancestor not yet born);
1093/// that chunk is **excluded** from the grafted tree and its digest is hashed directly into
1094/// the canonical root as the pending chunk. The caller must ensure that all ops tree nodes
1095/// for chunks `>= pruned_chunks` are still accessible in the ops tree (i.e., not pruned from
1096/// the journal).
1097///
1098/// `ops_leaves` must be a single consistent snapshot of `ops_tree.size()` taken in the same
1099/// instant as the bitmap state.
1100pub(super) async fn build_grafted_tree<
1101    F: merkle::Graftable,
1102    H: Hasher,
1103    S: Strategy,
1104    const N: usize,
1105>(
1106    bitmap: &impl bitmap::Readable<N>,
1107    pinned_nodes: &[H::Digest],
1108    ops_tree: &impl MerkleStorage<F, Digest = H::Digest>,
1109    ops_leaves: Location<F>,
1110    strategy: &S,
1111) -> Result<Mem<F, H::Digest>, Error<F>> {
1112    let grafting_height = grafting::height::<N>();
1113    let pruned_chunks = bitmap.pruned_chunks();
1114    let complete_chunks = bitmap.complete_chunks();
1115    let graftable_chunks = grafting::graftable_chunks::<F>(*ops_leaves, grafting_height)
1116        .min(complete_chunks as u64) as usize;
1117    assert!(
1118        pruned_chunks <= graftable_chunks && graftable_chunks <= complete_chunks,
1119        "invariant violated: pruned={pruned_chunks} graftable={graftable_chunks} complete={complete_chunks}"
1120    );
1121
1122    // Compute grafted leaves for each unpruned graftable chunk. The pending chunk (if any)
1123    // sits at index `graftable_chunks` and is excluded; its digest is hashed directly into
1124    // the canonical root.
1125    let leaves = compute_grafted_leaves::<F, H, S, N>(
1126        ops_tree,
1127        (pruned_chunks..graftable_chunks).map(|chunk_idx| (chunk_idx, bitmap.get_chunk(chunk_idx))),
1128        strategy,
1129    )
1130    .await?;
1131
1132    // Build the base grafted tree: either from pruned components or empty.
1133    let mut grafted_tree = if pruned_chunks > 0 {
1134        let grafted_pruning_boundary = Location::<F>::new(pruned_chunks as u64);
1135        Mem::from_components(Vec::new(), grafted_pruning_boundary, pinned_nodes.to_vec())
1136            .map_err(|_| Error::<F>::DataCorrupted("grafted tree rebuild failed"))?
1137    } else {
1138        Mem::new()
1139    };
1140
1141    // Add each grafted leaf digest.
1142    if !leaves.is_empty() {
1143        let batch = {
1144            let batch = grafted_tree.new_batch_with_strategy(strategy.clone());
1145            let batch = batch.add_leaf_digests(leaves.iter().map(|&(_, digest)| digest));
1146            let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
1147            batch.merkleize(&grafted_tree, &grafted_hasher)
1148        };
1149        grafted_tree.apply_batch(&batch)?;
1150    }
1151
1152    Ok(grafted_tree)
1153}
1154
1155/// Load the metadata and recover the pruning state persisted by previous runs.
1156///
1157/// The metadata store holds two kinds of entries (keyed by prefix):
1158/// - **Pruned chunks count** ([PRUNED_CHUNKS_PREFIX]): the number of bitmap chunks that have been
1159///   pruned. This tells us where the active portion of the bitmap begins.
1160/// - **Pinned node digests** ([NODE_PREFIX]): grafted tree digests at peak positions whose
1161///   underlying data has been pruned. These are needed to recompute the grafted tree root without
1162///   the pruned chunks.
1163///
1164/// Returns `(metadata_handle, pruned_chunks, pinned_node_digests)`.
1165pub(super) async fn init_metadata<F: merkle::Graftable, E: Context, D: Digest>(
1166    context: E,
1167    partition: &str,
1168) -> Result<(Metadata<E, U64, Vec<u8>>, usize, Vec<D>), Error<F>> {
1169    let metadata_cfg = MConfig {
1170        partition: partition.into(),
1171        codec_config: ((0..).into(), ()),
1172    };
1173    let metadata =
1174        Metadata::<_, U64, Vec<u8>>::init(context.child("metadata"), metadata_cfg).await?;
1175
1176    let key = U64::new(PRUNED_CHUNKS_PREFIX, 0);
1177    let pruned_chunks = match metadata.get(&key) {
1178        Some(bytes) => u64::from_be_bytes(bytes.as_slice().try_into().map_err(|_| {
1179            error!("pruned chunks value not a valid u64");
1180            Error::<F>::DataCorrupted("pruned chunks value not a valid u64")
1181        })?),
1182        None => {
1183            warn!("bitmap metadata does not contain pruned chunks, initializing as empty");
1184            0
1185        }
1186    } as usize;
1187
1188    // Load pinned nodes if database was pruned. We use nodes_to_pin on the grafted leaf count
1189    // to determine how many peaks to read. (Multiplying pruned_chunks by chunk_size is a
1190    // left-shift, preserving popcount, so the peak count is the same in grafted or ops space.)
1191    let pinned_nodes = if pruned_chunks > 0 {
1192        let pruned_loc = Location::<F>::new(pruned_chunks as u64);
1193        if !pruned_loc.is_valid() {
1194            return Err(Error::DataCorrupted("pruned chunks exceeds MAX_LEAVES"));
1195        }
1196        let mut pinned = Vec::new();
1197        for (index, _pos) in F::nodes_to_pin(pruned_loc).enumerate() {
1198            let metadata_key = U64::new(NODE_PREFIX, index as u64);
1199            let Some(bytes) = metadata.get(&metadata_key) else {
1200                return Err(Error::DataCorrupted(
1201                    "missing pinned node in grafted tree metadata",
1202                ));
1203            };
1204            let digest = D::decode(bytes.as_ref())
1205                .map_err(|_| Error::<F>::DataCorrupted("invalid pinned node digest"))?;
1206            pinned.push(digest);
1207        }
1208        pinned
1209    } else {
1210        Vec::new()
1211    };
1212
1213    Ok((metadata, pruned_chunks, pinned_nodes))
1214}
1215
1216#[cfg(test)]
1217mod tests {
1218    use super::*;
1219    use crate::{
1220        merkle::{hasher::Standard as StandardHasher, mmb, mmr, Bagging::ForwardFold},
1221        qmdb::{
1222            any::traits::{DbAny, UnmerkleizedBatch as _},
1223            current::{tests::fixed_config, unordered::fixed},
1224        },
1225        translator::OneCap,
1226    };
1227    use commonware_codec::FixedSize;
1228    use commonware_cryptography::{sha256, Sha256};
1229    use commonware_macros::test_traced;
1230    use commonware_runtime::{deterministic, Runner as _, Supervisor as _};
1231    use commonware_utils::bitmap::Prunable as PrunableBitMap;
1232
1233    const N: usize = sha256::Digest::SIZE;
1234
1235    #[test]
1236    fn partial_chunk_single_bit() {
1237        let mut bm = PrunableBitMap::<N>::new();
1238        bm.push(true);
1239        let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
1240        assert!(result.is_some());
1241        let (chunk, next_bit) = result.unwrap();
1242        assert_eq!(next_bit, 1);
1243        assert_eq!(chunk[0], 1); // bit 0 set
1244    }
1245
1246    #[test]
1247    fn partial_chunk_aligned() {
1248        let mut bm = PrunableBitMap::<N>::new();
1249        for _ in 0..PrunableBitMap::<N>::CHUNK_SIZE_BITS {
1250            bm.push(true);
1251        }
1252        let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
1253        assert!(result.is_none());
1254    }
1255
1256    #[test]
1257    fn partial_chunk_partial() {
1258        let mut bm = PrunableBitMap::<N>::new();
1259        for _ in 0..(PrunableBitMap::<N>::CHUNK_SIZE_BITS + 5) {
1260            bm.push(true);
1261        }
1262        let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
1263        assert!(result.is_some());
1264        let (_chunk, next_bit) = result.unwrap();
1265        assert_eq!(next_bit, 5);
1266    }
1267
1268    #[test]
1269    fn combine_roots_deterministic() {
1270        let ops = Sha256::hash(b"ops");
1271        let grafted = Sha256::hash(b"grafted");
1272        let r1 = combine_roots::<Sha256>(&ops, &grafted, None, None);
1273        let r2 = combine_roots::<Sha256>(&ops, &grafted, None, None);
1274        assert_eq!(r1, r2);
1275    }
1276
1277    #[test]
1278    fn combine_roots_with_partial_differs() {
1279        let ops = Sha256::hash(b"ops");
1280        let grafted = Sha256::hash(b"grafted");
1281        let partial_digest = Sha256::hash(b"partial");
1282
1283        let without = combine_roots::<Sha256>(&ops, &grafted, None, None);
1284        let with = combine_roots::<Sha256>(&ops, &grafted, None, Some((5, &partial_digest)));
1285        assert_ne!(without, with);
1286    }
1287
1288    #[test]
1289    fn combine_roots_with_pending_differs() {
1290        let ops = Sha256::hash(b"ops");
1291        let grafted = Sha256::hash(b"grafted");
1292        let pending_digest = Sha256::hash(b"pending");
1293
1294        let without = combine_roots::<Sha256>(&ops, &grafted, None, None);
1295        let with = combine_roots::<Sha256>(&ops, &grafted, Some(&pending_digest), None);
1296        assert_ne!(without, with);
1297    }
1298
1299    #[test]
1300    fn combine_roots_pending_and_partial_independent() {
1301        let ops = Sha256::hash(b"ops");
1302        let grafted = Sha256::hash(b"grafted");
1303        let pending_digest = Sha256::hash(b"pending");
1304        let partial_digest = Sha256::hash(b"partial");
1305
1306        let only_pending = combine_roots::<Sha256>(&ops, &grafted, Some(&pending_digest), None);
1307        let only_partial =
1308            combine_roots::<Sha256>(&ops, &grafted, None, Some((5, &partial_digest)));
1309        let both = combine_roots::<Sha256>(
1310            &ops,
1311            &grafted,
1312            Some(&pending_digest),
1313            Some((5, &partial_digest)),
1314        );
1315        assert_ne!(only_pending, only_partial);
1316        assert_ne!(only_pending, both);
1317        assert_ne!(only_partial, both);
1318    }
1319
1320    #[test]
1321    fn combine_roots_different_ops_root() {
1322        let ops_a = Sha256::hash(b"ops_a");
1323        let ops_b = Sha256::hash(b"ops_b");
1324        let grafted = Sha256::hash(b"grafted");
1325
1326        let r1 = combine_roots::<Sha256>(&ops_a, &grafted, None, None);
1327        let r2 = combine_roots::<Sha256>(&ops_b, &grafted, None, None);
1328        assert_ne!(r1, r2);
1329    }
1330
1331    /// Pin the canonical-root format down to the byte. A change to `combine_roots`'s hash
1332    /// pre-image (e.g., reordering, dropping a length tag, swapping pending/partial order)
1333    /// would silently break wire compatibility; this test catches that.
1334    #[test]
1335    fn combine_roots_format_golden() {
1336        let hasher = StandardHasher::<Sha256>::new(ForwardFold);
1337        let ops = Sha256::hash(b"ops");
1338        let grafted = Sha256::hash(b"grafted");
1339        let pending = Sha256::hash(b"pending");
1340        let partial = Sha256::hash(b"partial");
1341        let next_bit: u64 = 0x1122_3344_5566_7788;
1342
1343        // Neither pending nor partial.
1344        assert_eq!(
1345            combine_roots::<Sha256>(&ops, &grafted, None, None),
1346            hasher.hash([ops.as_ref(), grafted.as_ref()])
1347        );
1348
1349        // Pending only.
1350        assert_eq!(
1351            combine_roots::<Sha256>(&ops, &grafted, Some(&pending), None),
1352            hasher.hash([ops.as_ref(), grafted.as_ref(), pending.as_ref()])
1353        );
1354
1355        // Partial only.
1356        assert_eq!(
1357            combine_roots::<Sha256>(&ops, &grafted, None, Some((next_bit, &partial))),
1358            hasher.hash([
1359                ops.as_ref(),
1360                grafted.as_ref(),
1361                next_bit.to_be_bytes().as_slice(),
1362                partial.as_ref(),
1363            ])
1364        );
1365
1366        // Both: pending precedes partial.
1367        assert_eq!(
1368            combine_roots::<Sha256>(&ops, &grafted, Some(&pending), Some((next_bit, &partial))),
1369            hasher.hash([
1370                ops.as_ref(),
1371                grafted.as_ref(),
1372                pending.as_ref(),
1373                next_bit.to_be_bytes().as_slice(),
1374                partial.as_ref(),
1375            ])
1376        );
1377    }
1378
1379    type MmrDb = fixed::Db<
1380        mmr::Family,
1381        deterministic::Context,
1382        sha256::Digest,
1383        sha256::Digest,
1384        Sha256,
1385        OneCap,
1386        32,
1387        commonware_parallel::Sequential,
1388    >;
1389    type MmbDb = fixed::Db<
1390        mmb::Family,
1391        deterministic::Context,
1392        sha256::Digest,
1393        sha256::Digest,
1394        Sha256,
1395        OneCap,
1396        32,
1397        commonware_parallel::Sequential,
1398    >;
1399
1400    async fn populate_fixed_db<F, DB>(db: &mut DB, start: u64, count: u64)
1401    where
1402        F: merkle::Graftable,
1403        DB: DbAny<F, Key = sha256::Digest, Value = sha256::Digest>,
1404    {
1405        let mut batch = db.new_batch();
1406        for idx in start..start + count {
1407            let key = Sha256::hash(&idx.to_be_bytes());
1408            let value = Sha256::hash(&(idx + count).to_be_bytes());
1409            batch = batch.write(key, Some(value));
1410        }
1411        let merkleized = batch.merkleize(db, None).await.unwrap();
1412        db.apply_batch(merkleized).await.unwrap();
1413        db.commit().await.unwrap();
1414    }
1415
1416    /// A prune dropped between the pruning-metadata sync and the log prune must remain
1417    /// recoverable: the metadata durably records a bitmap boundary derived from a floor that
1418    /// may exist only in buffered operations, and reopening panics if the recovered floor
1419    /// lies below that boundary.
1420    #[test_traced]
1421    fn test_current_prune_dropped_before_log_prune() {
1422        let executor = deterministic::Runner::default();
1423        executor.start(|ctx| async move {
1424            let mut db = MmrDb::init(
1425                ctx.child("storage"),
1426                fixed_config::<OneCap>("prune-park", &ctx),
1427            )
1428            .await
1429            .unwrap();
1430
1431            // Establish a durable state, then apply (but do not commit) a batch that rewrites
1432            // every key, advancing the in-memory floor well past the durable commit's floor.
1433            populate_fixed_db::<mmr::Family, _>(&mut db, 0, 512).await;
1434            let durable_floor = db.inactivity_floor_loc();
1435            {
1436                let mut batch = db.new_batch();
1437                for idx in 0..512u64 {
1438                    let key = Sha256::hash(&idx.to_be_bytes());
1439                    let value = Sha256::hash(&(idx + 1024).to_be_bytes());
1440                    batch = batch.write(key, Some(value));
1441                }
1442                let merkleized = batch.merkleize(&db, None).await.unwrap();
1443                db.apply_batch(merkleized).await.unwrap();
1444            }
1445            assert!(db.sync_boundary() > durable_floor);
1446            let bounds = db.bounds();
1447            let floor = db.inactivity_floor_loc();
1448            let root = db.root();
1449
1450            // Drop the production prune future while it is parked after the metadata sync,
1451            // before the log prune: a genuine cancellation at that await.
1452            db.halt_before_prune_log = true;
1453            {
1454                let fut = db.prune(db.sync_boundary());
1455                futures::pin_mut!(fut);
1456                assert!(
1457                    futures::poll!(fut.as_mut()).is_pending(),
1458                    "prune must park before the log prune"
1459                );
1460            }
1461            let pruned_bits = db.any.bitmap.pruned_bits();
1462            assert!(pruned_bits > *durable_floor);
1463            drop(db);
1464
1465            // Reopening must succeed and recover the post-batch state: prune committed the
1466            // buffered operations before durably recording the pruning metadata that depends
1467            // on them. Asserting the advanced floor, root, and persisted pruned boundary
1468            // proves the drop happened after both the commit and the metadata sync.
1469            let db = MmrDb::init(
1470                ctx.child("reopen"),
1471                fixed_config::<OneCap>("prune-park", &ctx),
1472            )
1473            .await
1474            .expect("prune crash must leave the db recoverable");
1475            assert_eq!(db.bounds(), bounds);
1476            assert_eq!(db.inactivity_floor_loc(), floor);
1477            assert_eq!(db.root(), root);
1478            assert_eq!(db.any.bitmap.pruned_bits(), pruned_bits);
1479            db.destroy().await.unwrap();
1480        });
1481    }
1482
1483    #[test_traced]
1484    fn test_ops_root_witness_verifies_without_partial_chunk() {
1485        let executor = deterministic::Runner::default();
1486        executor.start(|ctx| async move {
1487            let mut db = MmrDb::init(
1488                ctx.child("storage"),
1489                fixed_config::<OneCap>("ops-root-witness-full", &ctx),
1490            )
1491            .await
1492            .unwrap();
1493            let mut next_idx = 0;
1494            populate_fixed_db::<mmr::Family, _>(&mut db, next_idx, 256).await;
1495            next_idx += 256;
1496            while partial_chunk::<_, 32>(db.any.bitmap.as_ref()).is_some() {
1497                populate_fixed_db::<mmr::Family, _>(&mut db, next_idx, 1).await;
1498                next_idx += 1;
1499            }
1500            let witness = db.ops_root_witness().await.unwrap();
1501            let ops_root = db.ops_root();
1502            let canonical_root = db.root();
1503
1504            assert!(witness.partial_chunk.is_none());
1505            assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
1506
1507            let wrong_ops_root = Sha256::hash(b"wrong ops root");
1508            assert!(!witness.verify::<Sha256>(&wrong_ops_root, &canonical_root));
1509
1510            let wrong_canonical_root = Sha256::hash(b"wrong canonical root");
1511            assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
1512
1513            let mut tampered = witness;
1514            tampered.grafted_root = Sha256::hash(b"wrong grafted root");
1515            assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
1516        });
1517    }
1518
1519    #[test_traced]
1520    fn test_ops_root_witness_verifies_with_partial_chunk() {
1521        let executor = deterministic::Runner::default();
1522        executor.start(|ctx| async move {
1523            let mut db = MmbDb::init(
1524                ctx.child("storage"),
1525                fixed_config::<OneCap>("ops-root-witness-partial", &ctx),
1526            )
1527            .await
1528            .unwrap();
1529            populate_fixed_db::<mmb::Family, _>(&mut db, 0, 260).await;
1530            let witness = db.ops_root_witness().await.unwrap();
1531            let ops_root = db.ops_root();
1532            let canonical_root = db.root();
1533
1534            assert!(witness.partial_chunk.is_some());
1535            assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
1536
1537            let wrong_ops_root = Sha256::hash(b"wrong ops root");
1538            assert!(!witness.verify::<Sha256>(&wrong_ops_root, &canonical_root));
1539
1540            let wrong_canonical_root = Sha256::hash(b"wrong canonical root");
1541            assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
1542
1543            let mut tampered = witness.clone();
1544            tampered.grafted_root = Sha256::hash(b"wrong grafted root");
1545            assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
1546
1547            let mut tampered = witness.clone();
1548            tampered.partial_chunk.as_mut().unwrap().0 += 1;
1549            assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
1550
1551            let mut tampered = witness;
1552            tampered.partial_chunk.as_mut().unwrap().1 = Sha256::hash(b"wrong partial chunk");
1553            assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
1554        });
1555    }
1556
1557    #[test_traced]
1558    fn test_ops_root_witness_verifies_with_pruned_db() {
1559        let executor = deterministic::Runner::default();
1560        executor.start(|ctx| async move {
1561            let mut db = MmrDb::init(
1562                ctx.child("storage"),
1563                fixed_config::<OneCap>("ops-root-witness-pruned", &ctx),
1564            )
1565            .await
1566            .unwrap();
1567
1568            // Churn the same keys repeatedly to drive the inactivity floor past chunk boundaries.
1569            for _ in 0..5 {
1570                populate_fixed_db::<mmr::Family, _>(&mut db, 0, 512).await;
1571            }
1572            db.prune(db.sync_boundary()).await.unwrap();
1573            assert!(
1574                db.any.bitmap.pruned_chunks() > 0,
1575                "test requires at least one pruned chunk to exercise the zero-chunk path"
1576            );
1577            let witness = db.ops_root_witness().await.unwrap();
1578            let ops_root = db.ops_root();
1579            let canonical_root = db.root();
1580
1581            assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
1582
1583            let wrong_canonical_root = Sha256::hash(b"wrong canonical root");
1584            assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
1585
1586            let mut tampered = witness;
1587            tampered.grafted_root = Sha256::hash(b"wrong grafted root");
1588            assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
1589        });
1590    }
1591
1592    #[test_traced]
1593    fn test_ops_root_witness_verifies_on_fresh_db() {
1594        let executor = deterministic::Runner::default();
1595        executor.start(|ctx| async move {
1596            let db = MmrDb::init(
1597                ctx.child("storage"),
1598                fixed_config::<OneCap>("ops-root-witness-fresh", &ctx),
1599            )
1600            .await
1601            .unwrap();
1602            let witness = db.ops_root_witness().await.unwrap();
1603            let ops_root = db.ops_root();
1604            let canonical_root = db.root();
1605
1606            assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
1607        });
1608    }
1609}