use crate::{
Context,
index::Unordered as UnorderedIndex,
journal::{
Error as JournalError,
contiguous::{Contiguous, Mutable},
},
merkle::{
self, Graftable, Location, Position, hasher::Hasher as _, mem::Mem,
storage::Storage as MerkleStorage,
},
metadata::{Config as MConfig, Metadata},
qmdb::{
self, Error,
any::{
self,
operation::{Operation, update::Update},
},
current::{
batch::BitmapBatch,
grafting,
proof::{OperationProof, OpsRootWitness, RangeProof, RangeProofSpec},
},
operation::Floored as _,
},
};
use commonware_codec::{Codec, CodecShared, DecodeExt};
use commonware_cryptography::{Digest, DigestOf, Hasher};
use commonware_macros::boxed;
use commonware_parallel::Strategy;
use commonware_runtime::{
Handle,
telemetry::metrics::{
Counter, Gauge, GaugeExt as _, MetricsExt as _,
histogram::{ScopedTimer, Timed},
},
};
use commonware_utils::{
bitmap::{self, Readable as _},
sequence::prefixed_u64::U64,
};
use core::{num::NonZeroU64, ops::Range};
use std::{collections::BTreeMap, sync::Arc};
use tracing::{error, warn};
const NODE_PREFIX: u8 = 0;
const PRUNED_CHUNKS_PREFIX: u8 = 1;
type GraftedPinnedNodes<F, D> = Vec<(Position<F>, D)>;
pub(crate) struct Metrics<E: Context> {
pruned_chunks: Gauge,
sync_boundary: Gauge,
pub apply_batch_calls: Counter,
apply_batch_duration: Timed,
pub sync_calls: Counter,
sync_duration: Timed,
pub prune_calls: Counter,
prune_duration: Timed,
clock: Arc<E>,
}
impl<E: Context> Metrics<E> {
pub fn new(context: E) -> Self {
Self {
pruned_chunks: context.gauge("pruned_chunks", "Number of pruned bitmap chunks"),
sync_boundary: context
.gauge("sync_boundary", "Most recent safe sync boundary location"),
apply_batch_calls: context.counter("apply_batch_calls", "Number of apply-batch calls"),
apply_batch_duration: Timed::register(
&context,
"apply_batch_duration",
"Duration of apply-batch calls",
),
sync_calls: context.counter("sync_calls", "Number of sync calls"),
sync_duration: Timed::register(&context, "sync_duration", "Duration of sync calls"),
prune_calls: context.counter("prune_calls", "Number of prune calls"),
prune_duration: Timed::register(&context, "prune_duration", "Duration of prune calls"),
clock: Arc::new(context),
}
}
pub fn apply_batch_timer(&self) -> ScopedTimer<E> {
self.apply_batch_duration.scoped(&self.clock)
}
pub fn sync_timer(&self) -> ScopedTimer<E> {
self.sync_duration.scoped(&self.clock)
}
pub fn prune_timer(&self) -> ScopedTimer<E> {
self.prune_duration.scoped(&self.clock)
}
pub fn update(&self, pruned_chunks: u64, sync_boundary: u64) {
let _ = self.pruned_chunks.try_set(pruned_chunks);
let _ = self.sync_boundary.try_set(sync_boundary);
}
}
pub struct Db<
F: merkle::Graftable,
E: Context,
C: Contiguous<Item: CodecShared>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
U: Send + Sync,
const N: usize,
S: Strategy,
> {
pub(super) any: any::db::Db<F, E, C, I, H, U, N, S>,
pub(super) grafted_tree: Arc<Mem<F, H::Digest>>,
pub(super) metadata: Metadata<E, U64, Vec<u8>>,
pub(super) strategy: S,
pub(super) root: DigestOf<H>,
pub(super) metrics: Metrics<E>,
#[cfg(test)]
pub(super) halt_before_prune_log: bool,
}
impl<F, E, C, I, H, U, const N: usize, S> std::fmt::Debug for Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
U: Update,
S: Strategy,
Operation<F, U>: Codec,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Db")
.field("bounds", &self.bounds())
.field("inactivity_floor_loc", &self.any.inactivity_floor_loc)
.finish_non_exhaustive()
}
}
impl<F, E, C, I, H, U, const N: usize, S> Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
U: Update,
S: Strategy,
Operation<F, U>: Codec,
{
#[cfg(any(test, feature = "test-traits"))]
pub(crate) const fn inactivity_floor_loc(&self) -> Location<F> {
self.any.inactivity_floor_loc()
}
pub const fn is_empty(&self) -> bool {
self.any.is_empty()
}
pub async fn get_metadata(&self) -> Result<Option<U::Value>, Error<F>> {
self.any.get_metadata().await
}
pub async fn get_many(&self, keys: &[&U::Key]) -> Result<Vec<Option<U::Value>>, Error<F>> {
self.any.get_many(keys).await
}
pub fn bounds(&self) -> std::ops::Range<Location<F>> {
self.any.bounds()
}
pub fn bitmap(&self) -> &impl bitmap::Readable<N> {
self.any.bitmap.as_ref()
}
pub fn verify_range_proof(
proof: &RangeProof<F, H::Digest>,
start_loc: Location<F>,
ops: &[Operation<F, U>],
chunks: &[[u8; N]],
root: &H::Digest,
) -> bool {
proof.verify::<H, _, N>(start_loc, ops, chunks, root)
}
pub fn grafted_storage(&self) -> impl MerkleStorage<F, Digest = H::Digest> + '_ {
grafting::Storage::<F, H, _, _>::new(
&self.grafted_tree,
grafting::height::<N>(),
&self.any.log.merkle,
)
}
pub const fn root(&self) -> H::Digest {
self.root
}
pub const fn strategy(&self) -> &S {
&self.strategy
}
pub const fn ops_root(&self) -> H::Digest {
self.any.root()
}
pub async fn ops_root_witness(&self) -> Result<OpsRootWitness<F, H::Digest>, Error<F>> {
let storage = self.grafted_storage();
let ops_size = storage.size();
let ops_leaves = Location::<F>::try_from(ops_size)?;
let grafted_root = compute_grafted_root::<F, H, _, _, N>(
self.any.bitmap.as_ref(),
&storage,
ops_leaves,
self.any.inactivity_floor_loc,
)
.await?;
let hasher = qmdb::hasher::<H>();
let partial_chunk = partial_chunk::<_, N>(self.any.bitmap.as_ref())
.map(|(chunk, next_bit)| (next_bit, hasher.digest(chunk.as_slice())));
let pending_chunk_digest: F::PendingChunk<H::Digest> = pending_chunk::<F, _, N>(
self.any.bitmap.as_ref(),
ops_leaves,
grafting::height::<N>(),
)?
.map(|chunk| hasher.digest(chunk.as_slice()))
.try_into()
.expect("pending_chunk must be consistent with family");
Ok(OpsRootWitness {
grafted_root,
pending_chunk_digest,
partial_chunk,
})
}
pub(super) fn grafted_snapshot(&self) -> Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>> {
merkle::batch::MerkleizedBatch::from_mem_with_strategy(
&self.grafted_tree,
self.strategy.clone(),
)
}
pub fn new_batch(&self) -> super::batch::UnmerkleizedBatch<F, H, U, N, S> {
super::batch::UnmerkleizedBatch::new(
self.any.new_batch(),
self.grafted_snapshot(),
BitmapBatch::Base(Arc::clone(&self.any.bitmap)),
)
}
pub(super) async fn operation_proof(
&self,
loc: Location<F>,
) -> Result<OperationProof<F, H::Digest, N>, Error<F>> {
let storage = self.grafted_storage();
let ops_root = self.any.root();
OperationProof::new::<H, _>(
self.any.bitmap.as_ref(),
&storage,
self.any.inactivity_floor_loc,
loc,
ops_root,
)
.await
}
#[allow(clippy::type_complexity)]
#[tracing::instrument(
name = "qmdb.current.db.range_proof",
level = "info",
skip_all,
fields(
start_loc = *start_loc,
max_ops = max_ops.get(),
),
)]
pub async fn range_proof(
&self,
start_loc: Location<F>,
max_ops: NonZeroU64,
) -> Result<(RangeProof<F, H::Digest>, Vec<Operation<F, U>>, Vec<[u8; N]>), Error<F>> {
let storage = self.grafted_storage();
let ops_root = self.any.root();
RangeProof::new_with_ops::<H, _, _, N>(
self.any.bitmap.as_ref(),
&storage,
&self.any.log,
RangeProofSpec {
start_loc,
max_ops,
inactivity_floor: self.any.inactivity_floor_loc,
ops_root,
},
)
.await
}
}
impl<F, E, C, I, H, U, const N: usize, S> Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
C: Mutable<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
U: Update,
S: Strategy,
Operation<F, U>: Codec,
{
pub async fn ops_historical_proof(
&self,
historical_size: Location<F>,
start_loc: Location<F>,
max_ops: NonZeroU64,
) -> Result<(merkle::Proof<F, H::Digest>, Vec<Operation<F, U>>), Error<F>> {
self.any
.historical_proof(historical_size, start_loc, max_ops)
.await
}
pub async fn pinned_nodes_at(&self, loc: Location<F>) -> Result<Vec<H::Digest>, Error<F>> {
self.any.pinned_nodes_at(loc).await
}
pub fn sync_boundary(&self) -> Location<F> {
sync_boundary::<F, N>(
*self.any.inactivity_floor_loc / bitmap::Prunable::<N>::CHUNK_SIZE_BITS,
*self.any.last_commit_loc + 1,
)
}
pub(super) fn update_metrics(&self) {
self.metrics.update(
self.any.bitmap.pruned_chunks() as u64,
*self.sync_boundary(),
);
}
fn delayed_merge_rewind_floor(&self) -> Option<u64> {
pair_absorption_threshold::<F, N>(self.any.bitmap.pruned_chunks() as u64)
}
fn grafted_pinned_nodes(
&self,
loc: Location<F>,
) -> Result<GraftedPinnedNodes<F, H::Digest>, Error<F>> {
F::nodes_to_pin(loc)
.map(|pos| {
let digest = self
.grafted_tree
.get_node(pos)
.ok_or(Error::<F>::DataCorrupted("missing grafted pinned node"))?;
Ok((pos, digest))
})
.collect()
}
fn prune_grafted_tree_to_bitmap(&mut self) -> Result<(), Error<F>> {
let pruned_chunks = self.any.bitmap.pruned_chunks() as u64;
if pruned_chunks == 0 {
return Ok(());
}
let prune_loc = Location::<F>::new(pruned_chunks);
if prune_loc <= self.grafted_tree.bounds().start {
return Ok(());
}
let prune_pos = Position::try_from(prune_loc)
.map_err(|_| Error::<F>::DataCorrupted("prune location overflow"))?;
let size = self.grafted_tree.size();
let pinned: BTreeMap<_, _> = self.grafted_pinned_nodes(prune_loc)?.into_iter().collect();
let mut retained = Vec::with_capacity((*size - *prune_pos) as usize);
for p in *prune_pos..*size {
let digest = self
.grafted_tree
.get_node(Position::new(p))
.ok_or(Error::<F>::DataCorrupted("missing retained grafted node"))?;
retained.push(digest);
}
self.grafted_tree = Arc::new(Mem::from_pruned_with_retained(prune_pos, pinned, retained));
Ok(())
}
#[tracing::instrument(name = "qmdb.current.db.prune", level = "info", skip_all)]
#[boxed]
pub async fn prune(mut self, prune_loc: Location<F>) -> Result<Self, Error<F>> {
let _timer = self.metrics.prune_timer();
self.metrics.prune_calls.inc();
let sync_boundary = self.sync_boundary();
if prune_loc > sync_boundary {
return Err(Error::PruneBeyondMinRequired(prune_loc, sync_boundary));
}
self.any.log = self.any.log.commit().await?;
self.any.prune_bitmap(sync_boundary);
self.prune_grafted_tree_to_bitmap()?;
self = self.sync_metadata().await?;
#[cfg(test)]
if self.halt_before_prune_log {
std::future::pending::<()>().await;
}
(self.any, _) = self.any.prune_log(prune_loc).await?;
self.any.update_metrics();
self.update_metrics();
Ok(self)
}
#[tracing::instrument(name = "qmdb.current.db.rewind", level = "info", skip_all)]
#[boxed]
pub async fn rewind(mut self, size: Location<F>) -> Result<Self, Error<F>> {
let rewind_size = *size;
let current_size = *self.any.last_commit_loc + 1;
if rewind_size == current_size {
return Ok(self);
}
if rewind_size == 0 || rewind_size > current_size {
return Err(Error::Journal(JournalError::InvalidRewind(rewind_size)));
}
let pruned_chunks = self.any.bitmap.pruned_chunks();
let pruned_bits = (pruned_chunks as u64)
.checked_mul(bitmap::Prunable::<N>::CHUNK_SIZE_BITS)
.ok_or_else(|| Error::DataCorrupted("pruned ops leaves overflow"))?;
if rewind_size < pruned_bits {
return Err(Error::Journal(JournalError::ItemPruned(rewind_size - 1)));
}
if let Some(rewind_floor) = self.delayed_merge_rewind_floor()
&& rewind_size < rewind_floor
{
return Err(Error::Journal(JournalError::ItemPruned(rewind_size - 1)));
}
{
let rewind_last_loc = Location::<F>::new(rewind_size - 1);
let rewind_last_op = self.any.log.read(*rewind_last_loc).await?;
let Some(rewind_floor) = rewind_last_op.has_floor() else {
return Err(Error::<F>::UnexpectedData(rewind_last_loc));
};
if *rewind_floor < pruned_bits {
return Err(Error::<F>::Journal(JournalError::ItemPruned(*rewind_floor)));
}
}
let pinned_nodes: Vec<H::Digest> = if pruned_chunks > 0 {
let grafted_leaves = Location::<F>::new(pruned_chunks as u64);
self.grafted_pinned_nodes(grafted_leaves)?
.into_iter()
.map(|(_, digest)| digest)
.collect()
} else {
Vec::new()
};
self.any = self.any.rewind(size).await?;
let (grafted_tree, root) = rebuild_grafted_tree::<F, H, S, N>(
self.any.bitmap.as_ref(),
&pinned_nodes,
&self.any.log.merkle,
self.any.inactivity_floor_loc,
self.any.root(),
&self.strategy,
)
.await?;
self.grafted_tree = Arc::new(grafted_tree);
self.root = root;
self.update_metrics();
Ok(self)
}
pub(crate) async fn sync_metadata(mut self) -> Result<Self, Error<F>> {
self.metadata.clear();
let pruned_chunks_u64 = self.any.bitmap.pruned_chunks() as u64;
let key = U64::new(PRUNED_CHUNKS_PREFIX, 0);
self.metadata
.put(key, pruned_chunks_u64.to_be_bytes().to_vec());
let pruned_chunks = Location::<F>::new(pruned_chunks_u64);
for (i, (_, digest)) in self
.grafted_pinned_nodes(pruned_chunks)?
.into_iter()
.enumerate()
{
let key = U64::new(NODE_PREFIX, i as u64);
self.metadata.put(key, digest.to_vec());
}
self.metadata = self.metadata.sync().await?;
Ok(self)
}
#[tracing::instrument(name = "qmdb.current.db.start_sync", level = "info", skip_all)]
#[boxed]
pub async fn start_sync(mut self) -> Result<(Self, Handle<()>), Error<F>> {
let (any, handle) = self.any.start_sync().await?;
self.any = any;
Ok((self, handle))
}
#[tracing::instrument(name = "qmdb.current.db.commit", level = "info", skip_all)]
#[boxed]
pub async fn commit(mut self) -> Result<Self, Error<F>> {
self.any = self.any.commit().await?;
Ok(self)
}
#[tracing::instrument(name = "qmdb.current.db.sync", level = "info", skip_all)]
#[boxed]
pub async fn sync(mut self) -> Result<Self, Error<F>> {
let _timer = self.metrics.sync_timer();
self.metrics.sync_calls.inc();
self.any = self.any.sync().await?;
self = self.sync_metadata().await?;
self.update_metrics();
Ok(self)
}
#[boxed]
pub async fn destroy(self) -> Result<(), Error<F>> {
let Self { any, metadata, .. } = self;
metadata.destroy().await?;
any.destroy().await
}
pub fn validate_batch(
&self,
batch: &super::batch::MerkleizedBatch<F, H::Digest, U, N, S>,
) -> Result<(), Error<F>> {
self.any.validate_batch(&batch.inner)
}
#[tracing::instrument(name = "qmdb.current.db.apply_batch", level = "info", skip_all)]
#[boxed]
pub async fn apply_batch(
mut self,
batch: Arc<super::batch::MerkleizedBatch<F, H::Digest, U, N, S>>,
) -> Result<(Self, Range<Location<F>>), Error<F>> {
let _timer = self.metrics.apply_batch_timer();
self.metrics.apply_batch_calls.inc();
let range;
(self.any, range) = self.any.apply_batch(Arc::clone(&batch.inner)).await?;
Arc::make_mut(&mut self.grafted_tree).apply_batch(&batch.grafted)?;
self.root = batch.canonical_root;
self.update_metrics();
Ok((self, range))
}
}
pub(crate) fn sync_boundary<F: Graftable, const N: usize>(
mut floor_chunks: u64,
ops_leaves: u64,
) -> Location<F> {
let chunk_bits = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
let grafting_height = grafting::height::<N>();
while floor_chunks > 0 {
let required_ops = pair_absorption_threshold::<F, N>(floor_chunks).unwrap_or_else(|| {
let youngest_start = (floor_chunks - 1) * chunk_bits;
let pos = F::subtree_root_position(Location::<F>::new(youngest_start), grafting_height);
F::peak_birth_size(pos, grafting_height)
});
if ops_leaves >= required_ops {
break;
}
floor_chunks -= 1;
}
Location::new(floor_chunks * chunk_bits)
}
fn pair_absorption_threshold<F: Graftable, const N: usize>(chunk_count: u64) -> Option<u64> {
if chunk_count == 0 {
return None;
}
let grafting_height = grafting::height::<N>();
let youngest = chunk_count - 1;
let youngest_start = youngest << grafting_height;
let youngest_end = (youngest + 1) << grafting_height;
let youngest_pos =
F::subtree_root_position(Location::<F>::new(youngest_start), grafting_height);
if F::peak_birth_size(youngest_pos, grafting_height) <= youngest_end {
return None;
}
let pair_chunk = youngest & !1;
let pair_start = pair_chunk << grafting_height;
let pair_pos = F::subtree_root_position(Location::<F>::new(pair_start), grafting_height + 1);
Some(F::peak_birth_size(pair_pos, grafting_height + 1))
}
pub(super) fn partial_chunk<B: bitmap::Readable<N>, const N: usize>(
bitmap: &B,
) -> Option<([u8; N], u64)> {
let next_bit = bitmap.len() % bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
if next_bit == 0 {
return None;
}
let (last_chunk, _) = bitmap.last_chunk();
Some((last_chunk, next_bit))
}
fn graftable_chunk_window<F: merkle::Graftable, B: bitmap::Readable<N>, const N: usize>(
bitmap: &B,
ops_leaves: Location<F>,
grafting_height: u32,
) -> Result<(u64, u64), Error<F>> {
let complete = bitmap.complete_chunks() as u64;
let graftable = grafting::graftable_chunks::<F>(*ops_leaves, grafting_height).min(complete);
let pending = complete - graftable;
if pending > 1 {
return Err(Error::DataCorrupted("multiple pending bitmap chunks"));
}
let pruned = bitmap.pruned_chunks() as u64;
if pruned > graftable {
return Err(Error::DataCorrupted(
"pruned chunks exceed graftable chunks",
));
}
Ok((complete, graftable))
}
pub(super) fn pending_chunk<F: merkle::Graftable, B: bitmap::Readable<N>, const N: usize>(
bitmap: &B,
ops_leaves: Location<F>,
grafting_height: u32,
) -> Result<Option<[u8; N]>, Error<F>> {
let (complete, graftable) =
graftable_chunk_window::<F, B, N>(bitmap, ops_leaves, grafting_height)?;
if complete - graftable != 1 {
return Ok(None);
}
Ok(Some(bitmap.get_chunk(graftable as usize)))
}
pub(super) fn combine_roots<H: Hasher>(
ops_root: &H::Digest,
grafted_root: &H::Digest,
pending: Option<&H::Digest>,
partial: Option<(u64, &H::Digest)>,
) -> H::Digest {
let hasher = qmdb::hasher::<H>();
match (pending, partial) {
(None, None) => hasher.hash(&[ops_root.as_ref(), grafted_root.as_ref()]),
(Some(pe), None) => hasher.hash(&[ops_root.as_ref(), grafted_root.as_ref(), pe.as_ref()]),
(None, Some((nb, p))) => {
let nb_bytes = nb.to_be_bytes();
hasher.hash(&[
ops_root.as_ref(),
grafted_root.as_ref(),
nb_bytes.as_slice(),
p.as_ref(),
])
}
(Some(pe), Some((nb, p))) => {
let nb_bytes = nb.to_be_bytes();
hasher.hash(&[
ops_root.as_ref(),
grafted_root.as_ref(),
pe.as_ref(),
nb_bytes.as_slice(),
p.as_ref(),
])
}
}
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn compute_db_root<
F: merkle::Graftable,
H: Hasher,
B: bitmap::Readable<N>,
S: MerkleStorage<F, Digest = H::Digest>,
const N: usize,
>(
status: &B,
storage: &S,
ops_leaves: Location<F>,
partial_chunk: Option<([u8; N], u64)>,
inactivity_floor: Location<F>,
ops_root: &H::Digest,
) -> Result<H::Digest, Error<F>> {
let grafted_root =
compute_grafted_root::<F, H, B, S, N>(status, storage, ops_leaves, inactivity_floor)
.await?;
let hasher = qmdb::hasher::<H>();
let pending = pending_chunk::<F, B, N>(status, ops_leaves, grafting::height::<N>())?
.map(|chunk| hasher.digest(&chunk));
let partial = partial_chunk.map(|(chunk, next_bit)| {
let digest = hasher.digest(&chunk);
(next_bit, digest)
});
Ok(combine_roots::<H>(
ops_root,
&grafted_root,
pending.as_ref(),
partial.as_ref().map(|(nb, d)| (*nb, d)),
))
}
pub(super) async fn rebuild_grafted_tree<F, H, S, const N: usize>(
bitmap: &impl bitmap::Readable<N>,
pinned_nodes: &[H::Digest],
ops_tree: &impl MerkleStorage<F, Digest = H::Digest>,
inactivity_floor: Location<F>,
ops_root: H::Digest,
strategy: &S,
) -> Result<(Mem<F, H::Digest>, H::Digest), Error<F>>
where
F: merkle::Graftable,
H: Hasher,
S: Strategy,
{
let ops_leaves = Location::<F>::try_from(ops_tree.size())?;
let grafted_tree =
build_grafted_tree::<F, H, S, N>(bitmap, pinned_nodes, ops_tree, ops_leaves, strategy)
.await?;
let storage =
grafting::Storage::<F, H, _, _>::new(&grafted_tree, grafting::height::<N>(), ops_tree);
let partial_chunk = partial_chunk(bitmap);
let root = compute_db_root::<F, H, _, _, N>(
bitmap,
&storage,
ops_leaves,
partial_chunk,
inactivity_floor,
&ops_root,
)
.await?;
Ok((grafted_tree, root))
}
pub(super) async fn compute_grafted_root<
F: merkle::Graftable,
H: Hasher,
B: bitmap::Readable<N>,
S: MerkleStorage<F, Digest = H::Digest>,
const N: usize,
>(
status: &B,
storage: &S,
ops_leaves: Location<F>,
inactivity_floor: Location<F>,
) -> Result<H::Digest, Error<F>> {
let size = storage.size();
let leaves = Location::try_from(size)?;
let mut peaks: Vec<H::Digest> = Vec::new();
for (peak_pos, _) in F::peaks(size) {
let digest = storage
.get_node(peak_pos)
.await?
.ok_or_else(|| merkle::Error::<F>::MissingNode(peak_pos))?;
peaks.push(digest);
}
let grafting_height = grafting::height::<N>();
let (_complete_chunks, _graftable_chunks) =
graftable_chunk_window::<F, B, N>(status, ops_leaves, grafting_height)?;
let inactive_peaks =
grafting::chunk_aligned_inactive_peaks::<F>(leaves, inactivity_floor, grafting_height)?;
let hasher = qmdb::hasher::<H>();
Ok(hasher.root(leaves, inactive_peaks, peaks.iter())?)
}
pub(super) async fn read_graft_inputs<F: merkle::Graftable, D: Digest, const N: usize>(
ops_tree: &impl MerkleStorage<F, Digest = D>,
chunks: impl IntoIterator<Item = (usize, [u8; N])>,
) -> Result<Vec<(usize, D, [u8; N])>, Error<F>> {
let grafting_height = grafting::height::<N>();
let chunks: Vec<(usize, [u8; N])> = chunks.into_iter().collect();
let positions: Vec<Position<F>> = chunks
.iter()
.map(|&(chunk_idx, _)| {
let leaf_start = Location::<F>::new((chunk_idx as u64) << grafting_height);
F::subtree_root_position(leaf_start, grafting_height)
})
.collect();
let nodes = ops_tree.get_nodes(&positions).await?;
Ok(chunks
.into_iter()
.zip(nodes)
.map(|((chunk_idx, chunk), chunk_ops_digest)| (chunk_idx, chunk_ops_digest, chunk))
.collect())
}
pub(super) async fn compute_grafted_leaves<
F: merkle::Graftable,
H: Hasher,
S: Strategy,
const N: usize,
>(
ops_tree: &impl MerkleStorage<F, Digest = H::Digest>,
chunks: impl IntoIterator<Item = (usize, [u8; N])>,
strategy: &S,
) -> Result<Vec<(usize, H::Digest)>, Error<F>> {
let inputs = read_graft_inputs::<F, _, N>(ops_tree, chunks).await?;
Ok(grafting::graft_chunk_digests::<H, _, N>(strategy, inputs))
}
pub(super) async fn build_grafted_tree<
F: merkle::Graftable,
H: Hasher,
S: Strategy,
const N: usize,
>(
bitmap: &impl bitmap::Readable<N>,
pinned_nodes: &[H::Digest],
ops_tree: &impl MerkleStorage<F, Digest = H::Digest>,
ops_leaves: Location<F>,
strategy: &S,
) -> Result<Mem<F, H::Digest>, Error<F>> {
let grafting_height = grafting::height::<N>();
let pruned_chunks = bitmap.pruned_chunks();
let complete_chunks = bitmap.complete_chunks();
let graftable_chunks = grafting::graftable_chunks::<F>(*ops_leaves, grafting_height)
.min(complete_chunks as u64) as usize;
assert!(
pruned_chunks <= graftable_chunks && graftable_chunks <= complete_chunks,
"invariant violated: pruned={pruned_chunks} graftable={graftable_chunks} complete={complete_chunks}"
);
let leaves = compute_grafted_leaves::<F, H, S, N>(
ops_tree,
(pruned_chunks..graftable_chunks).map(|chunk_idx| (chunk_idx, bitmap.get_chunk(chunk_idx))),
strategy,
)
.await?;
let mut grafted_tree = if pruned_chunks > 0 {
let grafted_pruning_boundary = Location::<F>::new(pruned_chunks as u64);
Mem::from_components(Vec::new(), grafted_pruning_boundary, pinned_nodes.to_vec())
.map_err(|_| Error::<F>::DataCorrupted("grafted tree rebuild failed"))?
} else {
Mem::new()
};
if !leaves.is_empty() {
let batch = {
let batch = grafted_tree.new_batch_with_strategy(strategy.clone());
let batch = batch.add_leaf_digests(leaves.iter().map(|&(_, digest)| digest));
let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
batch.merkleize(&grafted_tree, &grafted_hasher)
};
grafted_tree.apply_batch(&batch)?;
}
Ok(grafted_tree)
}
pub(super) async fn init_metadata<F: merkle::Graftable, E: Context, D: Digest>(
context: E,
partition: &str,
) -> Result<(Metadata<E, U64, Vec<u8>>, usize, Vec<D>), Error<F>> {
let metadata_cfg = MConfig {
partition: partition.into(),
codec_config: ((0..).into(), ()),
};
let metadata =
Metadata::<_, U64, Vec<u8>>::init(context.child("metadata"), metadata_cfg).await?;
let key = U64::new(PRUNED_CHUNKS_PREFIX, 0);
let pruned_chunks = match metadata.get(&key) {
Some(bytes) => u64::from_be_bytes(bytes.as_slice().try_into().map_err(|_| {
error!("pruned chunks value not a valid u64");
Error::<F>::DataCorrupted("pruned chunks value not a valid u64")
})?),
None => {
warn!("bitmap metadata does not contain pruned chunks, initializing as empty");
0
}
} as usize;
let pinned_nodes = if pruned_chunks > 0 {
let pruned_loc = Location::<F>::new(pruned_chunks as u64);
if !pruned_loc.is_valid() {
return Err(Error::DataCorrupted("pruned chunks exceeds MAX_LEAVES"));
}
let mut pinned = Vec::new();
for (index, _pos) in F::nodes_to_pin(pruned_loc).enumerate() {
let metadata_key = U64::new(NODE_PREFIX, index as u64);
let Some(bytes) = metadata.get(&metadata_key) else {
return Err(Error::DataCorrupted(
"missing pinned node in grafted tree metadata",
));
};
let digest = D::decode(bytes.as_ref())
.map_err(|_| Error::<F>::DataCorrupted("invalid pinned node digest"))?;
pinned.push(digest);
}
pinned
} else {
Vec::new()
};
Ok((metadata, pruned_chunks, pinned_nodes))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
merkle::{Bagging::ForwardFold, hasher::Standard as StandardHasher, mmb, mmr},
qmdb::{
any::traits::{DbAny, UnmerkleizedBatch as _},
current::{tests::fixed_config, unordered::fixed},
},
translator::OneCap,
};
use commonware_codec::FixedSize;
use commonware_cryptography::{Sha256, sha256};
use commonware_macros::test_traced;
use commonware_runtime::{Runner as _, Supervisor as _, deterministic};
use commonware_utils::bitmap::Prunable as PrunableBitMap;
const N: usize = sha256::Digest::SIZE;
#[test]
fn partial_chunk_single_bit() {
let mut bm = PrunableBitMap::<N>::new();
bm.push(true);
let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
assert!(result.is_some());
let (chunk, next_bit) = result.unwrap();
assert_eq!(next_bit, 1);
assert_eq!(chunk[0], 1); }
#[test]
fn partial_chunk_aligned() {
let mut bm = PrunableBitMap::<N>::new();
for _ in 0..PrunableBitMap::<N>::CHUNK_SIZE_BITS {
bm.push(true);
}
let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
assert!(result.is_none());
}
#[test]
fn partial_chunk_partial() {
let mut bm = PrunableBitMap::<N>::new();
for _ in 0..(PrunableBitMap::<N>::CHUNK_SIZE_BITS + 5) {
bm.push(true);
}
let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
assert!(result.is_some());
let (_chunk, next_bit) = result.unwrap();
assert_eq!(next_bit, 5);
}
#[test]
fn partial_chunk_empty() {
let bm = PrunableBitMap::<N>::new();
assert!(partial_chunk::<PrunableBitMap<N>, N>(&bm).is_none());
}
#[test]
fn partial_chunk_fully_pruned() {
let bm = PrunableBitMap::<N>::new_with_pruned_chunks(1).unwrap();
assert!(partial_chunk::<PrunableBitMap<N>, N>(&bm).is_none());
}
#[test]
fn combine_roots_deterministic() {
let ops = Sha256::hash(&[b"ops"]);
let grafted = Sha256::hash(&[b"grafted"]);
let r1 = combine_roots::<Sha256>(&ops, &grafted, None, None);
let r2 = combine_roots::<Sha256>(&ops, &grafted, None, None);
assert_eq!(r1, r2);
}
#[test]
fn combine_roots_with_partial_differs() {
let ops = Sha256::hash(&[b"ops"]);
let grafted = Sha256::hash(&[b"grafted"]);
let partial_digest = Sha256::hash(&[b"partial"]);
let without = combine_roots::<Sha256>(&ops, &grafted, None, None);
let with = combine_roots::<Sha256>(&ops, &grafted, None, Some((5, &partial_digest)));
assert_ne!(without, with);
}
#[test]
fn combine_roots_with_pending_differs() {
let ops = Sha256::hash(&[b"ops"]);
let grafted = Sha256::hash(&[b"grafted"]);
let pending_digest = Sha256::hash(&[b"pending"]);
let without = combine_roots::<Sha256>(&ops, &grafted, None, None);
let with = combine_roots::<Sha256>(&ops, &grafted, Some(&pending_digest), None);
assert_ne!(without, with);
}
#[test]
fn combine_roots_pending_and_partial_independent() {
let ops = Sha256::hash(&[b"ops"]);
let grafted = Sha256::hash(&[b"grafted"]);
let pending_digest = Sha256::hash(&[b"pending"]);
let partial_digest = Sha256::hash(&[b"partial"]);
let only_pending = combine_roots::<Sha256>(&ops, &grafted, Some(&pending_digest), None);
let only_partial =
combine_roots::<Sha256>(&ops, &grafted, None, Some((5, &partial_digest)));
let both = combine_roots::<Sha256>(
&ops,
&grafted,
Some(&pending_digest),
Some((5, &partial_digest)),
);
assert_ne!(only_pending, only_partial);
assert_ne!(only_pending, both);
assert_ne!(only_partial, both);
}
#[test]
fn combine_roots_different_ops_root() {
let ops_a = Sha256::hash(&[b"ops_a"]);
let ops_b = Sha256::hash(&[b"ops_b"]);
let grafted = Sha256::hash(&[b"grafted"]);
let r1 = combine_roots::<Sha256>(&ops_a, &grafted, None, None);
let r2 = combine_roots::<Sha256>(&ops_b, &grafted, None, None);
assert_ne!(r1, r2);
}
#[test]
fn combine_roots_format_golden() {
let hasher = StandardHasher::<Sha256>::new(ForwardFold);
let ops = Sha256::hash(&[b"ops"]);
let grafted = Sha256::hash(&[b"grafted"]);
let pending = Sha256::hash(&[b"pending"]);
let partial = Sha256::hash(&[b"partial"]);
let next_bit: u64 = 0x1122_3344_5566_7788;
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, None, None),
hasher.hash(&[ops.as_ref(), grafted.as_ref()])
);
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, Some(&pending), None),
hasher.hash(&[ops.as_ref(), grafted.as_ref(), pending.as_ref()])
);
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, None, Some((next_bit, &partial))),
hasher.hash(&[
ops.as_ref(),
grafted.as_ref(),
next_bit.to_be_bytes().as_slice(),
partial.as_ref(),
])
);
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, Some(&pending), Some((next_bit, &partial))),
hasher.hash(&[
ops.as_ref(),
grafted.as_ref(),
pending.as_ref(),
next_bit.to_be_bytes().as_slice(),
partial.as_ref(),
])
);
}
type MmrDb = fixed::Db<
mmr::Family,
deterministic::Context,
sha256::Digest,
sha256::Digest,
Sha256,
OneCap,
32,
commonware_parallel::Sequential,
>;
type MmbDb = fixed::Db<
mmb::Family,
deterministic::Context,
sha256::Digest,
sha256::Digest,
Sha256,
OneCap,
32,
commonware_parallel::Sequential,
>;
#[boxed]
async fn populate_fixed_db<F, DB>(db: DB, start: u64, count: u64) -> DB
where
F: merkle::Graftable,
DB: DbAny<F, Key = sha256::Digest, Value = sha256::Digest>,
{
let mut batch = db.new_batch();
for idx in start..start + count {
let key = Sha256::hash(&[&idx.to_be_bytes()]);
let value = Sha256::hash(&[&(idx + count).to_be_bytes()]);
batch = batch.write(key, Some(value));
}
let merkleized = batch.merkleize(&db, None).await.unwrap();
let (db, _) = db.apply_batch(merkleized).await.unwrap();
db.commit().await.unwrap()
}
#[test_traced]
fn test_operations_match_applied_range() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let db = MmrDb::init(
ctx.child("db"),
fixed_config::<OneCap>("operations-match-applied-range", &ctx),
)
.await
.unwrap();
let db = populate_fixed_db::<mmr::Family, _>(db, 0, 8).await;
let mut batch = db.new_batch();
for idx in 0..4u64 {
let key = Sha256::hash(&[&idx.to_be_bytes()]);
let value = Sha256::hash(&[&(idx + 100).to_be_bytes()]);
batch = batch.write(key, Some(value));
}
let merkleized = batch.merkleize(&db, None).await.unwrap();
let (start, ops) = merkleized.operations();
let (db, range) = db.apply_batch(merkleized).await.unwrap();
assert_eq!(start, range.start);
assert_eq!(*start + ops.len() as u64, *range.end);
db.destroy().await.unwrap();
});
}
#[test_traced]
fn test_start_sync_recovery() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let db = MmrDb::init(
ctx.child("first"),
fixed_config::<OneCap>("start-sync-recovery", &ctx),
)
.await
.unwrap();
let key = Sha256::hash(&[&0u64.to_be_bytes()]);
let value = Sha256::hash(&[&1u64.to_be_bytes()]);
let merkleized = db
.new_batch()
.write(key, Some(value))
.merkleize(&db, None)
.await
.unwrap();
let (db, _) = db.apply_batch(merkleized).await.unwrap();
let (db, handle) = db.start_sync().await.unwrap();
handle.await.unwrap();
let root = db.root();
drop(db);
let db = MmrDb::init(
ctx.child("second"),
fixed_config::<OneCap>("start-sync-recovery", &ctx),
)
.await
.unwrap();
assert_eq!(db.root(), root);
assert_eq!(db.get(&key).await.unwrap(), Some(value));
db.destroy().await.unwrap();
});
}
#[test_traced]
fn test_current_prune_dropped_before_log_prune() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("prune-park", &ctx),
)
.await
.unwrap();
let db = populate_fixed_db::<mmr::Family, _>(db, 0, 512).await;
let durable_floor = db.inactivity_floor_loc();
let mut batch = db.new_batch();
for idx in 0..512u64 {
let key = Sha256::hash(&[&idx.to_be_bytes()]);
let value = Sha256::hash(&[&(idx + 1024).to_be_bytes()]);
batch = batch.write(key, Some(value));
}
let merkleized = batch.merkleize(&db, None).await.unwrap();
let (mut db, _) = db.apply_batch(merkleized).await.unwrap();
assert!(db.sync_boundary() > durable_floor);
let bounds = db.bounds();
let floor = db.inactivity_floor_loc();
let root = db.root();
db.halt_before_prune_log = true;
let boundary = db.sync_boundary();
{
let fut = db.prune(boundary);
futures::pin_mut!(fut);
assert!(
futures::poll!(fut.as_mut()).is_pending(),
"prune must park before the log prune"
);
}
let db = MmrDb::init(
ctx.child("reopen"),
fixed_config::<OneCap>("prune-park", &ctx),
)
.await
.expect("prune crash must leave the db recoverable");
assert_eq!(db.bounds(), bounds);
assert_eq!(db.inactivity_floor_loc(), floor);
assert_eq!(db.root(), root);
assert!(db.any.bitmap.pruned_bits() > *durable_floor);
db.destroy().await.unwrap();
});
}
#[test_traced]
fn test_ops_root_witness_verifies_without_partial_chunk() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let mut db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-full", &ctx),
)
.await
.unwrap();
let mut next_idx = 0;
db = populate_fixed_db::<mmr::Family, _>(db, next_idx, 256).await;
next_idx += 256;
while partial_chunk::<_, 32>(db.any.bitmap.as_ref()).is_some() {
db = populate_fixed_db::<mmr::Family, _>(db, next_idx, 1).await;
next_idx += 1;
}
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.partial_chunk.is_none());
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
let wrong_ops_root = Sha256::hash(&[b"wrong ops root"]);
assert!(!witness.verify::<Sha256>(&wrong_ops_root, &canonical_root));
let wrong_canonical_root = Sha256::hash(&[b"wrong canonical root"]);
assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
let mut tampered = witness;
tampered.grafted_root = Sha256::hash(&[b"wrong grafted root"]);
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
});
}
#[test_traced]
fn test_ops_root_witness_verifies_with_partial_chunk() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let db = MmbDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-partial", &ctx),
)
.await
.unwrap();
let db = populate_fixed_db::<mmb::Family, _>(db, 0, 260).await;
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.partial_chunk.is_some());
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
let wrong_ops_root = Sha256::hash(&[b"wrong ops root"]);
assert!(!witness.verify::<Sha256>(&wrong_ops_root, &canonical_root));
let wrong_canonical_root = Sha256::hash(&[b"wrong canonical root"]);
assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
let mut tampered = witness.clone();
tampered.grafted_root = Sha256::hash(&[b"wrong grafted root"]);
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
let mut tampered = witness.clone();
tampered.partial_chunk.as_mut().unwrap().0 += 1;
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
let mut tampered = witness;
tampered.partial_chunk.as_mut().unwrap().1 = Sha256::hash(&[b"wrong partial chunk"]);
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
});
}
#[test_traced]
fn test_ops_root_witness_verifies_with_pruned_db() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let mut db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-pruned", &ctx),
)
.await
.unwrap();
for _ in 0..5 {
db = populate_fixed_db::<mmr::Family, _>(db, 0, 512).await;
}
let boundary = db.sync_boundary();
let db = db.prune(boundary).await.unwrap();
assert!(
db.any.bitmap.pruned_chunks() > 0,
"test requires at least one pruned chunk to exercise the zero-chunk path"
);
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
let wrong_canonical_root = Sha256::hash(&[b"wrong canonical root"]);
assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
let mut tampered = witness;
tampered.grafted_root = Sha256::hash(&[b"wrong grafted root"]);
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
});
}
#[test_traced]
fn test_ops_root_witness_verifies_on_fresh_db() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-fresh", &ctx),
)
.await
.unwrap();
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
});
}
}