1use 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
49const NODE_PREFIX: u8 = 0;
51
52const PRUNED_CHUNKS_PREFIX: u8 = 1;
54
55pub(crate) struct Metrics<E: Context> {
57 pruned_chunks: Gauge,
59 sync_boundary: Gauge,
61 pub apply_batch_calls: Counter,
63 apply_batch_duration: Timed,
65 pub sync_calls: Counter,
67 sync_duration: Timed,
69 pub prune_calls: Counter,
71 prune_duration: Timed,
73 clock: Arc<E>,
75}
76
77impl<E: Context> Metrics<E> {
78 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 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
117pub 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 pub(super) any: any::db::Db<F, E, C, I, H, U, N, S>,
132
133 pub(super) grafted_tree: Arc<Mem<F, H::Digest>>,
144
145 pub(super) metadata: Metadata<E, U64, Vec<u8>>,
149
150 pub(super) strategy: S,
153
154 pub(super) root: DigestOf<H>,
157
158 pub(super) metrics: Metrics<E>,
160
161 #[cfg(test)]
164 pub(super) halt_before_prune_log: bool,
165}
166
167impl<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 #[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 pub const fn is_empty(&self) -> bool {
188 self.any.is_empty()
189 }
190
191 pub async fn get_metadata(&self) -> Result<Option<U::Value>, Error<F>> {
193 self.any.get_metadata().await
194 }
195
196 pub fn bounds(&self) -> std::ops::Range<Location<F>> {
199 self.any.bounds()
200 }
201
202 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
215impl<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 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 pub const fn root(&self) -> H::Digest {
241 self.root
242 }
243
244 pub const fn strategy(&self) -> &S {
246 &self.strategy
247 }
248
249 pub const fn ops_root(&self) -> H::Digest {
260 self.any.root()
261 }
262
263 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 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 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 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 #[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
372impl<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 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 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 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 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 fn delayed_merge_rewind_floor(&self) -> Option<u64> {
459 pair_absorption_threshold::<F, N>(self.any.bitmap.pruned_chunks() as u64)
460 }
461
462 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 #[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 self.any.log.commit().await?;
532
533 self.any.prune_bitmap(sync_boundary);
535 self.prune_grafted_tree_to_bitmap()?;
536
537 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 #[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 if rewind_size == current_size {
582 return Ok(());
583 }
584 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 {
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 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 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 pub(crate) async fn sync_metadata(&mut self) -> Result<(), Error<F>> {
676 self.metadata.clear();
677
678 let pruned_chunks_u64 = self.any.bitmap.pruned_chunks() as u64;
680
681 let key = U64::new(PRUNED_CHUNKS_PREFIX, 0);
683 self.metadata
684 .put(key, pruned_chunks_u64.to_be_bytes().to_vec());
685
686 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
703pub(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
732fn 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
757impl<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 #[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 #[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 self.sync_metadata().await?;
786 self.update_metrics();
787 Ok(())
788 }
789
790 #[boxed]
792 pub async fn destroy(self) -> Result<(), Error<F>> {
793 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 #[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
837pub(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
850fn 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
876pub(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
901pub(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#[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
987pub(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 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 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 Ok(hasher.root(leaves, inactive_peaks, peaks.iter())?)
1036}
1037
1038pub(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 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
1065pub(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
1086pub(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 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 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 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
1155pub(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 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); }
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 #[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 assert_eq!(
1345 combine_roots::<Sha256>(&ops, &grafted, None, None),
1346 hasher.hash([ops.as_ref(), grafted.as_ref()])
1347 );
1348
1349 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 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 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 #[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 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 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 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 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}