use crate::{
Context,
index::Unordered as UnorderedIndex,
journal::contiguous::{Contiguous, Mutable},
merkle::{
self, Graftable, Location, Position, Readable,
batch::MerkleizedBatch as GenericMerkleizedBatch, mem::Mem,
storage::Storage as MerkleStorage,
},
qmdb::{
Error,
any::{
self, ValueEncoding,
batch::{DiffCursors, DiffEntry, Staged as AnyStaged, StagedUpdates},
operation::{Operation, update},
},
batch_chain::Bounds,
bitmap::{Shared, fill_from},
current::{
db::{compute_db_root, partial_chunk, read_graft_inputs},
grafting,
},
operation::Key,
},
};
use ahash::AHashMap;
use commonware_codec::Codec;
use commonware_cryptography::{Digest, Hasher};
use commonware_parallel::Strategy;
use commonware_utils::bitmap::{self, Readable as _};
use core::ops::Range;
use std::sync::Arc;
#[derive(Clone, Debug, Default)]
pub(crate) struct ChunkOverlay<const N: usize> {
pub(crate) chunks: AHashMap<usize, [u8; N]>,
pub(crate) len: u64,
parent: Dimensions,
}
#[derive(Clone, Copy, Debug, Default)]
struct Dimensions {
len: u64,
complete_chunks: usize,
pruned_chunks: usize,
}
impl Dimensions {
fn of<B: bitmap::Readable<N>, const N: usize>(base: &B) -> Self {
Self {
len: base.len(),
complete_chunks: base.complete_chunks(),
pruned_chunks: base.pruned_chunks(),
}
}
}
impl<const N: usize> ChunkOverlay<N> {
const CHUNK_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
fn new<B: bitmap::Readable<N>>(base: &B, len: u64, capacity: usize) -> Self {
Self {
chunks: AHashMap::with_capacity(capacity),
len,
parent: Dimensions::of(base),
}
}
fn chunk_mut<B: bitmap::Readable<N>>(&mut self, base: &B, idx: usize) -> &mut [u8; N] {
let parent = self.parent;
self.chunks.entry(idx).or_insert_with(|| {
let base_has_partial = !parent.len.is_multiple_of(Self::CHUNK_BITS);
if idx < parent.complete_chunks {
base.get_chunk(idx)
} else if idx == parent.complete_chunks && base_has_partial {
base.last_chunk().0
} else {
bitmap::BitMap::<N>::EMPTY_CHUNK
}
})
}
fn set_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
let rel = (loc % Self::CHUNK_BITS) as usize;
let chunk = self.chunk_mut(base, idx);
chunk[rel / 8] |= 1 << (rel % 8);
}
fn clear_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
if idx < self.parent.pruned_chunks {
return;
}
let rel = (loc % Self::CHUNK_BITS) as usize;
let chunk = self.chunk_mut(base, idx);
chunk[rel / 8] &= !(1 << (rel % 8));
}
pub(crate) fn get(&self, idx: usize) -> Option<&[u8; N]> {
self.chunks.get(&idx)
}
pub(crate) const fn complete_chunks(&self) -> usize {
(self.len / Self::CHUNK_BITS) as usize
}
}
pub(crate) fn fill_candidates<F: Graftable, const N: usize>(
bitmap: &BitmapBatch<N>,
floor: Location<F>,
tip: u64,
limit: usize,
out: &mut Vec<Location<F>>,
) -> Location<F> {
Location::new(fill_from(bitmap, *floor, tip, limit, out))
}
struct BatchStorageAdapter<
'a,
F: Graftable,
D: Digest,
R: Readable<Family = F, Digest = D>,
S: MerkleStorage<F, Digest = D>,
> {
batch: &'a R,
base: &'a S,
_phantom: core::marker::PhantomData<(F, D)>,
}
impl<
'a,
F: Graftable,
D: Digest,
R: Readable<Family = F, Digest = D>,
S: MerkleStorage<F, Digest = D>,
> BatchStorageAdapter<'a, F, D, R, S>
{
const fn new(batch: &'a R, base: &'a S) -> Self {
Self {
batch,
base,
_phantom: core::marker::PhantomData,
}
}
}
impl<F: Graftable, D: Digest, R: Readable<Family = F, Digest = D>, S: MerkleStorage<F, Digest = D>>
MerkleStorage<F> for BatchStorageAdapter<'_, F, D, R, S>
{
type Digest = D;
fn size(&self) -> Position<F> {
self.batch.size()
}
async fn get_node(&self, pos: Position<F>) -> Result<Option<D>, merkle::Error<F>> {
if let Some(node) = self.batch.get_node(pos) {
return Ok(Some(node));
}
self.base.get_node(pos).await
}
async fn get_nodes(&self, positions: &[Position<F>]) -> Result<Vec<D>, merkle::Error<F>> {
let mut nodes = vec![None; positions.len()];
let mut base_positions = Vec::with_capacity(positions.len());
for (slot, &pos) in nodes.iter_mut().zip(positions) {
match self.batch.get_node(pos) {
Some(node) => *slot = Some(node),
None => base_positions.push(pos),
}
}
let base_nodes = if base_positions.is_empty() {
Vec::new()
} else {
self.base.get_nodes(&base_positions).await?
};
let mut base_nodes = base_nodes.into_iter();
Ok(nodes
.into_iter()
.map(|node| node.unwrap_or_else(|| base_nodes.next().expect("one node per base read")))
.collect())
}
}
struct BatchOverMem<'a, F: Graftable, D: Digest, S: Strategy> {
batch: &'a GenericMerkleizedBatch<F, D, S>,
mem: &'a Mem<F, D>,
}
impl<F: Graftable, D: Digest, S: Strategy> Readable for BatchOverMem<'_, F, D, S> {
type Family = F;
type Digest = D;
fn size(&self) -> Position<F> {
self.batch.size()
}
fn get_node(&self, pos: Position<F>) -> Option<D> {
if let Some(d) = self.batch.get_node(pos) {
return Some(d);
}
self.mem.get_node(pos)
}
}
pub struct UnmerkleizedBatch<F, H, U, const N: usize, S: Strategy>
where
F: Graftable,
U: update::Update,
H: Hasher,
Operation<F, U>: Codec,
{
inner: any::batch::UnmerkleizedBatch<F, H, U, S>,
grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
bitmap_parent: BitmapBatch<N>,
}
pub struct Staged<F, H, U, const N: usize, S: Strategy>
where
F: Graftable,
U: update::Update,
H: Hasher,
Operation<F, U>: Codec,
{
inner: AnyStaged<F, H, U, S>,
grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
bitmap_parent: BitmapBatch<N>,
}
pub struct MerkleizedBatch<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
{
pub(crate) inner: Arc<any::batch::MerkleizedBatch<F, D, U, S>>,
pub(crate) grafted: Arc<merkle::batch::MerkleizedBatch<F, D, S>>,
pub(crate) bitmap: BitmapBatch<N>,
pub(crate) canonical_root: D,
}
impl<F, H, U, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, U, N, S>
where
F: Graftable,
U: update::Update,
H: Hasher,
Operation<F, U>: Codec,
{
pub(super) const fn new(
inner: any::batch::UnmerkleizedBatch<F, H, U, S>,
grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
bitmap_parent: BitmapBatch<N>,
) -> Self {
Self {
inner,
grafted_parent,
bitmap_parent,
}
}
pub fn write(mut self, key: U::Key, value: Option<U::Value>) -> Self {
self.inner = self.inner.write(key, value);
self
}
pub async fn get<E, C, I>(
&self,
key: &U::Key,
db: &super::db::Db<F, E, C, I, H, U, N, S>,
) -> Result<Option<U::Value>, Error<F>>
where
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
{
self.inner.get(key, &db.any).await
}
pub async fn get_many<E, C, I>(
&self,
keys: &[&U::Key],
db: &super::db::Db<F, E, C, I, H, U, N, S>,
) -> Result<Vec<Option<U::Value>>, Error<F>>
where
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
{
self.inner.get_many(keys, &db.any).await
}
pub async fn stage<E, C, I>(
self,
keys: &[&U::Key],
db: &super::db::Db<F, E, C, I, H, U, N, S>,
) -> Result<(Vec<Option<U::Value>>, Staged<F, H, U, N, S>), Error<F>>
where
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
{
let Self {
inner,
grafted_parent,
bitmap_parent,
} = self;
let (values, inner) = inner.stage(keys, &db.any).await?;
Ok((
values,
Staged {
inner,
grafted_parent,
bitmap_parent,
},
))
}
}
impl<F, H, U, const N: usize, S: Strategy> Staged<F, H, U, N, S>
where
F: Graftable,
U: update::Update,
H: Hasher,
Operation<F, U>: Codec,
{
pub async fn expand<E, C, I>(
self,
keys: &[&U::Key],
db: &super::db::Db<F, E, C, I, H, U, N, S>,
) -> Result<(Range<usize>, Vec<Option<U::Value>>, Self), Error<F>>
where
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
{
let Self {
inner,
grafted_parent,
bitmap_parent,
} = self;
let (range, values, inner) = inner.expand(keys, &db.any).await?;
Ok((
range,
values,
Self {
inner,
grafted_parent,
bitmap_parent,
},
))
}
}
impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Unordered<K, V>, N, S>
where
F: Graftable,
K: Key,
V: ValueEncoding,
H: Hasher,
Operation<F, update::Unordered<K, V>>: Codec,
{
#[allow(clippy::type_complexity)]
#[tracing::instrument(
name = "qmdb.current.unordered.batch.merkleize.staged",
level = "info",
skip_all,
fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
)]
pub async fn merkleize<E, C, I>(
self,
updates: Vec<(usize, Option<V::Value>)>,
upserts: Vec<(K, Option<V::Value>)>,
metadata: Option<V::Value>,
db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
where
E: Context,
C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
{
let Self {
inner,
grafted_parent,
bitmap_parent,
} = self;
let (inner, staged_updates, prefetched) = inner
.resolve_updates_prefetched(updates, upserts, &db.any, |floor, tip, limit, out| {
fill_candidates(&bitmap_parent, floor, tip, limit, out)
})
.await?;
let inner = inner
.merkleize_with_floor_scan(
&db.any,
metadata,
staged_updates,
Some(prefetched),
|floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
)
.await?;
compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
}
}
impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Ordered<K, V>, N, S>
where
F: Graftable,
K: Key,
V: ValueEncoding,
H: Hasher,
Operation<F, update::Ordered<K, V>>: Codec,
{
#[allow(clippy::type_complexity)]
#[tracing::instrument(
name = "qmdb.current.ordered.batch.merkleize.staged",
level = "info",
skip_all,
fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
)]
pub async fn merkleize<E, C, I>(
self,
updates: Vec<(usize, Option<V::Value>)>,
upserts: Vec<(K, Option<V::Value>)>,
metadata: Option<V::Value>,
db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
where
E: Context,
C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
I: crate::index::Ordered<Value = Location<F>> + 'static,
{
let Self {
inner,
grafted_parent,
bitmap_parent,
} = self;
let (inner, staged_updates) = inner.resolve_updates(updates, upserts, db.any.strategy());
let inner = inner
.merkleize_with_floor_scan(
&db.any,
metadata,
staged_updates,
|floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
)
.await?;
compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
}
}
impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
where
F: Graftable,
K: Key,
V: ValueEncoding,
H: Hasher,
Operation<F, update::Unordered<K, V>>: Codec,
{
#[allow(clippy::type_complexity)]
#[tracing::instrument(
name = "qmdb.current.unordered.batch.merkleize",
level = "info",
skip_all
)]
pub async fn merkleize<E, C, I>(
self,
db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
metadata: Option<V::Value>,
) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
where
E: Context,
C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
{
let Self {
inner,
grafted_parent,
bitmap_parent,
} = self;
let inner = inner
.merkleize_with_floor_scan(
&db.any,
metadata,
StagedUpdates::<F, update::Unordered<K, V>>::new(),
None,
|floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
)
.await?;
compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
}
}
impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
where
F: Graftable,
K: Key,
V: ValueEncoding,
H: Hasher,
Operation<F, update::Ordered<K, V>>: Codec,
{
#[allow(clippy::type_complexity)]
#[tracing::instrument(
name = "qmdb.current.ordered.batch.merkleize",
level = "info",
skip_all
)]
pub async fn merkleize<E, C, I>(
self,
db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
metadata: Option<V::Value>,
) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
where
E: Context,
C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
I: crate::index::Ordered<Value = Location<F>> + 'static,
{
let Self {
inner,
grafted_parent,
bitmap_parent,
} = self;
let inner = inner
.merkleize_with_floor_scan(
&db.any,
metadata,
StagedUpdates::<F, update::Ordered<K, V>>::new(),
|floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
)
.await?;
compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
}
}
#[allow(clippy::type_complexity)]
fn build_chunk_overlay<F: Graftable, U, B: bitmap::Readable<N>, const N: usize>(
base: &B,
batch_len: usize,
batch_base: u64,
diff: &[(U::Key, DiffEntry<F, U::Value>)],
ancestor_diffs: &[Arc<Vec<(U::Key, DiffEntry<F, U::Value>)>>],
) -> ChunkOverlay<N>
where
U: update::Update,
{
let total_bits = base.len() + batch_len as u64;
let appended_chunks = (batch_len as u64).div_ceil(ChunkOverlay::<N>::CHUNK_BITS) as usize;
let mut overlay = ChunkOverlay::new(base, total_bits, diff.len() + appended_chunks + 1);
let commit_loc = batch_base + batch_len as u64 - 1;
overlay.set_bit(base, commit_loc);
overlay.clear_bit(base, batch_base - 1);
let mut ancestors = DiffCursors::new(ancestor_diffs.iter().map(|d| d.as_slice()));
for (key, entry) in diff {
if let Some(loc) = entry.loc()
&& *loc >= batch_base
&& *loc < batch_base + batch_len as u64
{
overlay.set_bit(base, *loc);
}
let mut prev_loc = entry.base_old_loc();
if let Some(ancestor_entry) = ancestors.resolve(key) {
prev_loc = ancestor_entry.loc();
}
if let Some(old) = prev_loc {
overlay.clear_bit(base, *old);
}
}
let parent_complete = overlay.parent.complete_chunks;
let new_complete = overlay.complete_chunks();
for idx in parent_complete..new_complete {
overlay.chunk_mut(base, idx);
}
overlay
}
async fn merkleize_grafted_batch<F, H, S, const N: usize>(
strategy: &S,
grafted_parent: Arc<GenericMerkleizedBatch<F, H::Digest, S>>,
grafted_tree: &Arc<Mem<F, H::Digest>>,
graft_inputs: Vec<(usize, H::Digest, [u8; N])>,
grafting_height: u32,
) -> Arc<GenericMerkleizedBatch<F, H::Digest, S>>
where
F: Graftable,
H: Hasher,
S: Strategy,
{
let old_grafted_leaves = *grafted_parent.leaves() as usize;
let mut grafted_batch = grafted_parent.new_batch();
let ancestors = grafted_batch.retain_ancestors();
let grafted_tree = Arc::clone(grafted_tree);
strategy
.clone()
.spawn(graft_inputs.len(), move |strategy| {
let new_leaves = grafting::graft_chunk_digests::<H, _, N>(&strategy, graft_inputs);
for (chunk_idx, digest) in new_leaves {
if chunk_idx < old_grafted_leaves {
grafted_batch = grafted_batch
.update_leaf_digest(Location::<F>::new(chunk_idx as u64), digest)
.expect("update_leaf_digest failed");
} else {
grafted_batch = grafted_batch.add_leaf_digest(digest);
}
}
let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
let merkleized = grafted_batch.merkleize(&grafted_tree, &grafted_hasher);
drop(ancestors);
merkleized
})
.await
}
async fn compute_current_layer<F, E, U, C, I, H, const N: usize, S>(
inner: Arc<any::batch::MerkleizedBatch<F, H::Digest, U, S>>,
current_db: &super::db::Db<F, E, C, I, H, U, N, S>,
grafted_parent: &Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
bitmap_parent: &BitmapBatch<N>,
) -> Result<Arc<MerkleizedBatch<F, H::Digest, U, N, S>>, Error<F>>
where
F: Graftable,
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
U: update::Update,
S: Strategy,
Operation<F, U>: Codec,
{
let batch_len = inner.journal_batch.items().len();
let batch_base = *inner.bounds.tip.size - batch_len as u64;
let overlay = build_chunk_overlay::<F, U, _, N>(
bitmap_parent,
batch_len,
batch_base,
&inner.diff,
&inner.ancestor_diffs,
);
let grafting_height = grafting::height::<N>();
let ops_tree_adapter =
BatchStorageAdapter::new(&inner.journal_batch, ¤t_db.any.log.merkle);
let overlay_ops_leaves = inner.bounds.tip.size;
let new_complete_chunks = overlay.complete_chunks();
let graftable_overlay = grafting::graftable_chunks::<F>(*overlay_ops_leaves, grafting_height)
.min(new_complete_chunks as u64) as usize;
let graftable_parent = *grafted_parent.leaves() as usize;
let pruned_chunks = bitmap_parent.pruned_chunks();
assert!(
pruned_chunks <= graftable_parent
&& graftable_parent <= graftable_overlay
&& graftable_overlay <= new_complete_chunks,
"invariant violated: pruned={pruned_chunks} graftable_parent={graftable_parent} graftable_overlay={graftable_overlay} new_complete={new_complete_chunks}"
);
let mut chunk_indices_to_update: Vec<usize> = overlay
.chunks
.iter()
.filter(|&(&idx, _)| idx < graftable_overlay && idx >= pruned_chunks)
.map(|(&idx, _)| idx)
.collect();
chunk_indices_to_update.extend(graftable_parent..graftable_overlay);
chunk_indices_to_update.sort_unstable();
chunk_indices_to_update.dedup();
let chunks_to_update = chunk_indices_to_update.into_iter().map(|idx| {
let chunk = overlay
.get(idx)
.copied()
.unwrap_or_else(|| bitmap_parent.get_chunk(idx));
(idx, chunk)
});
let graft_inputs = read_graft_inputs::<F, _, N>(&ops_tree_adapter, chunks_to_update).await?;
let grafted_batch = if graft_inputs.is_empty() {
let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
grafted_parent
.new_batch()
.merkleize(¤t_db.grafted_tree, &grafted_hasher)
} else {
merkleize_grafted_batch::<F, H, S, N>(
¤t_db.strategy,
Arc::clone(grafted_parent),
¤t_db.grafted_tree,
graft_inputs,
grafting_height,
)
.await
};
let bitmap_batch = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: bitmap_parent.clone(),
overlay: Arc::new(overlay),
shared: Arc::clone(bitmap_parent.shared()),
}));
let ops_root = inner.root();
let layered = BatchOverMem {
batch: &grafted_batch,
mem: ¤t_db.grafted_tree,
};
let grafted_storage =
grafting::Storage::<F, H, _, _>::new(&layered, grafting_height, &ops_tree_adapter);
let partial = partial_chunk::<_, N>(&bitmap_batch);
let canonical_root = compute_db_root::<F, H, _, _, N>(
&bitmap_batch,
&grafted_storage,
overlay_ops_leaves,
partial,
inner.bounds.inactivity_floor,
&ops_root,
)
.await?;
Ok(Arc::new(MerkleizedBatch {
inner,
grafted: grafted_batch,
bitmap: bitmap_batch,
canonical_root,
}))
}
#[derive(Clone, Debug)]
pub(crate) enum BitmapBatch<const N: usize> {
Base(Arc<Shared<N>>),
Layer(Arc<BitmapBatchLayer<N>>),
}
#[derive(Debug)]
pub(crate) struct BitmapBatchLayer<const N: usize> {
pub(crate) parent: BitmapBatch<N>,
pub(crate) overlay: Arc<ChunkOverlay<N>>,
pub(crate) shared: Arc<Shared<N>>,
}
impl<const N: usize> BitmapBatch<N> {
const CHUNK_SIZE_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
fn shared(&self) -> &Arc<Shared<N>> {
match self {
Self::Base(s) => s,
Self::Layer(layer) => &layer.shared,
}
}
fn trim_committed(&self) -> Self {
let shared = self.shared();
let committed = bitmap::Readable::<N>::len(shared.as_ref());
let mut kept = Vec::new();
let mut current = self;
while let Self::Layer(layer) = current {
if layer.overlay.len <= committed {
break;
}
kept.push(Arc::clone(&layer.overlay));
current = &layer.parent;
}
let mut result = Self::Base(Arc::clone(shared));
for overlay in kept.into_iter().rev() {
result = Self::Layer(Arc::new(BitmapBatchLayer {
parent: result,
overlay,
shared: Arc::clone(shared),
}));
}
result
}
}
impl<const N: usize> bitmap::Readable<N> for BitmapBatch<N> {
fn complete_chunks(&self) -> usize {
(self.len() / Self::CHUNK_SIZE_BITS) as usize
}
fn get_chunk(&self, idx: usize) -> [u8; N] {
let mut current = self;
loop {
match current {
Self::Base(shared) => return shared.get_chunk(idx),
Self::Layer(layer) => {
if let Some(&chunk) = layer.overlay.get(idx) {
return chunk;
}
current = &layer.parent;
}
}
}
}
fn last_chunk(&self) -> ([u8; N], u64) {
let total = self.len();
if total == 0 {
return (bitmap::BitMap::<N>::EMPTY_CHUNK, 0);
}
let rem = total % Self::CHUNK_SIZE_BITS;
let bits_in_last = if rem == 0 { Self::CHUNK_SIZE_BITS } else { rem };
let idx = if rem == 0 {
self.complete_chunks().saturating_sub(1)
} else {
self.complete_chunks()
};
(self.get_chunk(idx), bits_in_last)
}
fn pruned_chunks(&self) -> usize {
self.shared().pruned_chunks()
}
fn len(&self) -> u64 {
match self {
Self::Base(shared) => bitmap::Readable::<N>::len(shared.as_ref()),
Self::Layer(layer) => layer.overlay.len,
}
}
}
impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
MerkleizedBatch<F, D, U, N, S>
{
pub const fn root(&self) -> D {
self.canonical_root
}
pub fn ops_root(&self) -> D {
self.inner.root()
}
pub fn bounds(&self) -> &Bounds<F, D> {
self.inner.bounds()
}
pub fn operations(&self) -> (Location<F>, Arc<Vec<Operation<F, U>>>) {
self.inner.operations()
}
pub fn sync_boundary(&self) -> Location<F> {
super::db::sync_boundary::<F, N>(
*self.inner.bounds().inactivity_floor / bitmap::Prunable::<N>::CHUNK_SIZE_BITS,
*self.inner.bounds().tip.size,
)
}
}
impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
MerkleizedBatch<F, D, U, N, S>
where
Operation<F, U>: Codec,
{
pub fn new_batch<H>(self: &Arc<Self>) -> UnmerkleizedBatch<F, H, U, N, S>
where
H: Hasher<Digest = D>,
{
UnmerkleizedBatch::new(
self.inner.new_batch::<H>(),
Arc::clone(&self.grafted),
self.bitmap.trim_committed(),
)
}
pub async fn get<E, C, I, H>(
&self,
key: &U::Key,
db: &super::db::Db<F, E, C, I, H, U, N, S>,
) -> Result<Option<U::Value>, Error<F>>
where
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
H: Hasher<Digest = D>,
{
self.inner.get(key, &db.any).await
}
pub async fn get_many<E, C, I, H>(
&self,
keys: &[&U::Key],
db: &super::db::Db<F, E, C, I, H, U, N, S>,
) -> Result<Vec<Option<U::Value>>, Error<F>>
where
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
H: Hasher<Digest = D>,
{
self.inner.get_many(keys, &db.any).await
}
}
impl<F, E, C, I, H, U, const N: usize, S> super::db::Db<F, E, C, I, H, U, N, S>
where
F: Graftable,
E: Context,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
U: update::Update,
S: Strategy,
Operation<F, U>: Codec,
{
pub fn to_batch(&self) -> Arc<MerkleizedBatch<F, H::Digest, U, N, S>> {
let grafted = self.grafted_snapshot();
Arc::new(MerkleizedBatch {
inner: self.any.to_batch(),
grafted,
bitmap: BitmapBatch::Base(Arc::clone(&self.any.bitmap)),
canonical_root: self.root,
})
}
}
#[cfg(any(test, feature = "test-traits"))]
mod trait_impls {
use super::*;
use crate::{
journal::contiguous::Mutable,
qmdb::any::traits::{
ApplyBatchResult, BatchableDb, MerkleizedBatch as MerkleizedBatchTrait,
UnmerkleizedBatch as UnmerkleizedBatchTrait,
},
};
use std::future::Future;
type CurrentDb<F, E, C, I, H, U, const N: usize, S> =
crate::qmdb::current::db::Db<F, E, C, I, H, U, N, S>;
impl<F, K, V, H, E, C, I, const N: usize, S>
UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>>
for UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
where
F: Graftable,
K: Key,
V: ValueEncoding + 'static,
H: Hasher,
E: Context,
C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
S: Strategy,
Operation<F, update::Unordered<K, V>>: Codec,
{
type Family = F;
type K = K;
type V = V::Value;
type Metadata = V::Value;
type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;
fn write(self, key: K, value: Option<V::Value>) -> Self {
Self::write(self, key, value)
}
async fn merkleize(
self,
db: &CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>,
metadata: Option<V::Value>,
) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
self.merkleize(db, metadata).await
}
}
impl<F, K, V, H, E, C, I, const N: usize, S>
UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>>
for UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
where
F: Graftable,
K: Key,
V: ValueEncoding + 'static,
H: Hasher,
E: Context,
C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
I: crate::index::Ordered<Value = Location<F>> + 'static,
S: Strategy,
Operation<F, update::Ordered<K, V>>: Codec,
{
type Family = F;
type K = K;
type V = V::Value;
type Metadata = V::Value;
type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;
fn write(self, key: K, value: Option<V::Value>) -> Self {
Self::write(self, key, value)
}
async fn merkleize(
self,
db: &CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>,
metadata: Option<V::Value>,
) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
self.merkleize(db, metadata).await
}
}
impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
MerkleizedBatchTrait for Arc<MerkleizedBatch<F, D, U, N, S>>
where
Operation<F, U>: Codec,
{
type Digest = D;
fn root(&self) -> D {
MerkleizedBatch::root(self)
}
}
impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
for CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>
where
F: Graftable,
E: Context,
K: Key,
V: ValueEncoding + 'static,
C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
I: UnorderedIndex<Value = Location<F>> + 'static,
H: Hasher,
S: Strategy,
Operation<F, update::Unordered<K, V>>: Codec,
{
type Family = F;
type K = K;
type V = V::Value;
type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;
type Batch = UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>;
fn new_batch(&self) -> Self::Batch {
self.new_batch()
}
fn apply_batch(
self,
batch: Self::Merkleized,
) -> impl Future<Output = ApplyBatchResult<Self>> {
self.apply_batch(batch)
}
}
impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
for CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>
where
F: Graftable,
E: Context,
K: Key,
V: ValueEncoding + 'static,
C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
I: crate::index::Ordered<Value = Location<F>> + 'static,
H: Hasher,
S: Strategy,
Operation<F, update::Ordered<K, V>>: Codec,
{
type Family = F;
type K = K;
type V = V::Value;
type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;
type Batch = UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>;
fn new_batch(&self) -> Self::Batch {
self.new_batch()
}
fn apply_batch(
self,
batch: Self::Merkleized,
) -> impl Future<Output = ApplyBatchResult<Self>> {
self.apply_batch(batch)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{mmb, mmr, utils::detached::block_strategy};
use commonware_cryptography::Sha256;
use commonware_macros::test_traced;
use commonware_parallel::{Manual, Rayon};
use commonware_utils::{NZUsize, bitmap::Prunable as BitMap};
use std::{
future::Future as _,
task::Context as TaskContext,
time::{Duration, Instant},
};
const N: usize = 4;
type Bm = BitMap<N>;
type GraftedBatch =
Arc<GenericMerkleizedBatch<mmb::Family, <Sha256 as Hasher>::Digest, Manual<Rayon>>>;
type Location = mmr::Location;
fn make_bitmap(bits: &[bool]) -> Bm {
let mut bm = Bm::new();
for &b in bits {
bm.push(b);
}
bm
}
fn grafted_chain(
strategy: &Manual<Rayon>,
mem: &Arc<Mem<mmb::Family, <Sha256 as Hasher>::Digest>>,
) -> (GraftedBatch, GraftedBatch) {
let hasher = grafting::hasher::<mmb::Family, Sha256>(grafting::height::<1>());
let a = mem
.new_batch_with_strategy(strategy.clone())
.add_leaf_digest(Sha256::hash(&[b"a-0"]))
.add_leaf_digest(Sha256::hash(&[b"a-1"]))
.merkleize(mem, &hasher);
let b = a
.new_batch()
.add_leaf_digest(Sha256::hash(&[b"b-0"]))
.merkleize(mem, &hasher);
(a, b)
}
#[test_traced]
fn test_grafted_merkleize_retains_ancestors_after_cancellation() {
let strategy = Rayon::new(NZUsize!(2)).unwrap();
let manual = strategy.manual();
let mem = Arc::new(Mem::<mmb::Family, <Sha256 as Hasher>::Digest>::new());
let grafting_height = grafting::height::<1>();
let graft_inputs = || vec![(0, Sha256::hash(&[b"replacement"]), [1u8; 1])];
let waker = futures::task::noop_waker();
let mut context = TaskContext::from_waker(&waker);
let (a, b) = grafted_chain(&manual, &mem);
let ancestor = Arc::downgrade(&a);
let release = block_strategy(&strategy, 2);
let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
&manual,
Arc::clone(&b),
&mem,
graft_inputs(),
grafting_height,
));
assert!(merkleize.as_mut().poll(&mut context).is_pending());
drop(b);
drop(a);
drop(release);
let _ = futures::executor::block_on(merkleize);
assert!(ancestor.upgrade().is_none());
let (a, b) = grafted_chain(&manual, &mem);
let ancestor = Arc::downgrade(&a);
let release = block_strategy(&strategy, 2);
let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
&manual,
Arc::clone(&b),
&mem,
graft_inputs(),
grafting_height,
));
assert!(merkleize.as_mut().poll(&mut context).is_pending());
drop(merkleize);
drop(b);
drop(a);
assert!(ancestor.upgrade().is_some());
drop(release);
let deadline = Instant::now() + Duration::from_secs(10);
while ancestor.upgrade().is_some() {
assert!(
Instant::now() < deadline,
"detached grafted merkleization did not release its ancestors"
);
std::thread::yield_now();
}
}
#[test]
fn chunk_overlay_pushes() {
use crate::qmdb::any::value::FixedEncoding;
use commonware_utils::sequence::FixedBytes;
type K = FixedBytes<4>;
type V = FixedEncoding<u64>;
type U = crate::qmdb::any::operation::update::Unordered<K, V>;
let key1 = FixedBytes::from([1, 0, 0, 0]);
let key2 = FixedBytes::from([2, 0, 0, 0]);
let base = make_bitmap(&[true; 4]);
let mut diff = vec![
(
key1,
DiffEntry::Active {
value: 100u64,
loc: Location::new(4), base_old_loc: None,
},
),
(
key2,
DiffEntry::Active {
value: 200u64,
loc: Location::new(99), base_old_loc: None,
},
),
];
diff.sort_by(|a, b| a.0.cmp(&b.0));
let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 4, 4, &diff, &[]);
let c0 = overlay.get(0).expect("chunk 0 should be dirty");
assert_ne!(c0[0] & (1 << 4), 0); assert_eq!(c0[0] & (1 << 5), 0); assert_eq!(c0[0] & (1 << 6), 0); assert_ne!(c0[0] & (1 << 7), 0); assert_eq!(c0[0] & (1 << 3), 0); }
#[test]
fn chunk_overlay_clears() {
use crate::qmdb::any::value::FixedEncoding;
use commonware_utils::sequence::FixedBytes;
type K = FixedBytes<4>;
type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;
let key1 = FixedBytes::from([1, 0, 0, 0]);
let key2 = FixedBytes::from([2, 0, 0, 0]);
let key3 = FixedBytes::from([3, 0, 0, 0]);
let base = make_bitmap(&[true; 64]);
let mut diff: Vec<(K, DiffEntry<mmr::Family, u64>)> = vec![
(
key1,
DiffEntry::Active {
value: 100,
loc: Location::new(70),
base_old_loc: Some(Location::new(5)),
},
),
(
key2,
DiffEntry::Deleted {
base_old_loc: Some(Location::new(10)),
},
),
(
key3,
DiffEntry::Active {
value: 300,
loc: Location::new(71),
base_old_loc: None,
},
),
];
diff.sort_by(|a, b| a.0.cmp(&b.0));
let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 8, 64, &diff, &[]);
let c0 = overlay.get(0).expect("chunk 0 should be dirty");
assert_eq!(c0[0] & (1 << 5), 0); assert_eq!(c0[1] & (1 << 2), 0);
assert_eq!(c0[0] & (1 << 4), 1 << 4); assert_eq!(c0[1] & (1 << 3), 1 << 3); }
#[test]
fn chunk_overlay_preserves_partial_parent_chunk() {
use crate::qmdb::any::value::FixedEncoding;
use commonware_utils::sequence::FixedBytes;
type K = FixedBytes<4>;
type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;
let base = make_bitmap(&[true; 20]);
assert_eq!(base.complete_chunks(), 0);
let key1 = FixedBytes::from([1, 0, 0, 0]);
let mut diff = vec![(
key1,
DiffEntry::Active {
value: 42u64,
loc: Location::new(35),
base_old_loc: None,
},
)];
diff.sort_by(|a, b| a.0.cmp(&b.0));
let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 20, 20, &diff, &[]);
let c0 = overlay.get(0).expect("chunk 0 should be in overlay");
assert_eq!(c0[0], 0xFF);
assert_eq!(c0[1], 0xFF);
assert_eq!(c0[2], 0x07);
}
fn next_candidate<B: bitmap::Readable<N2>, const N2: usize>(
bitmap: &B,
floor: Location,
tip: u64,
) -> Option<Location> {
let floor = *floor;
let bitmap_len = bitmap.len();
let committed_end = bitmap_len.min(tip);
if floor < committed_end
&& let Some(idx) = bitmap.ones_iter_from(floor).next()
&& idx < committed_end
{
return Some(Location::new(idx));
}
let candidate = floor.max(bitmap_len);
(candidate < tip).then(|| Location::new(candidate))
}
#[test]
fn bitmap_scan_all_active() {
let bm = make_bitmap(&[true; 8]);
for i in 0..8 {
assert_eq!(
next_candidate(&bm, Location::new(i), 8),
Some(Location::new(i))
);
}
assert_eq!(next_candidate(&bm, Location::new(8), 8), None);
}
#[test]
fn bitmap_scan_all_inactive() {
let bm = make_bitmap(&[false; 8]);
assert_eq!(next_candidate(&bm, Location::new(0), 8), None);
}
#[test]
fn bitmap_scan_skips_inactive() {
let bm = make_bitmap(&[false, false, true, false, true]);
assert_eq!(
next_candidate(&bm, Location::new(0), 5),
Some(Location::new(2))
);
assert_eq!(
next_candidate(&bm, Location::new(3), 5),
Some(Location::new(4))
);
assert_eq!(next_candidate(&bm, Location::new(5), 5), None);
}
#[test]
fn bitmap_scan_beyond_bitmap_len_returns_candidate() {
let bm = make_bitmap(&[false; 4]);
assert_eq!(
next_candidate(&bm, Location::new(0), 8),
Some(Location::new(4))
);
assert_eq!(
next_candidate(&bm, Location::new(6), 8),
Some(Location::new(6))
);
}
#[test]
fn bitmap_scan_respects_tip() {
let bm = make_bitmap(&[false, false, false, true]);
assert_eq!(next_candidate(&bm, Location::new(0), 3), None);
assert_eq!(
next_candidate(&bm, Location::new(0), 4),
Some(Location::new(3))
);
}
#[test]
fn bitmap_scan_floor_at_tip() {
let bm = make_bitmap(&[true; 4]);
assert_eq!(next_candidate(&bm, Location::new(4), 4), None);
}
#[test]
fn bitmap_scan_empty_bitmap() {
let bm = Bm::new();
assert_eq!(
next_candidate(&bm, Location::new(0), 5),
Some(Location::new(0))
);
assert_eq!(next_candidate(&bm, Location::new(0), 0), None);
}
#[test]
fn fill_candidates_matches_oracle() {
fn assert_matches(name: &str, chain: &BitmapBatch<N>, tip: u64) {
for floor in 0..=tip {
let mut want = Vec::new();
let mut scan = Location::new(floor);
while let Some(c) = next_candidate(chain, scan, tip) {
want.push(c);
scan = c + 1;
}
for split in 0..=want.len() {
let mut got = Vec::new();
let next = fill_candidates(chain, Location::new(floor), tip, split, &mut got);
fill_candidates(chain, next, tip, want.len() + 1, &mut got);
assert_eq!(got, want, "{name} floor={floor} split={split}");
}
}
}
let bits = [true, false, true, true, false, false, true, false];
let base = make_bitmap(&bits);
let flat = BitmapBatch::Base(Arc::new(Shared::new(make_bitmap(&bits))));
let shared = Arc::new(Shared::new(make_bitmap(&bits)));
let mut overlay = ChunkOverlay::new(&base, 12, 1);
overlay.clear_bit(&base, 3);
overlay.clear_bit(&base, 6);
overlay.set_bit(&base, 9);
let one_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: BitmapBatch::Base(Arc::clone(&shared)),
overlay: Arc::new(overlay),
shared,
}));
let shared = Arc::new(Shared::new(make_bitmap(&bits)));
let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
overlay1.clear_bit(&base, 3);
overlay1.set_bit(&base, 9);
let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: BitmapBatch::Base(Arc::clone(&shared)),
overlay: Arc::new(overlay1),
shared: Arc::clone(&shared),
}));
let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
overlay2.clear_bit(&chain1, 6);
overlay2.clear_bit(&chain1, 9);
overlay2.set_bit(&chain1, 13);
let two_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: chain1,
overlay: Arc::new(overlay2),
shared,
}));
let make_pruned = || {
let mut bits = [false; 40];
bits[33] = true;
bits[38] = true;
let mut bm = make_bitmap(&bits);
bm.prune_to_bit(32);
bm
};
let pruned_base = make_pruned();
let shared = Arc::new(Shared::new(make_pruned()));
let mut overlay = ChunkOverlay::new(&pruned_base, 46, 1);
overlay.clear_bit(&pruned_base, 38);
overlay.set_bit(&pruned_base, 41);
overlay.set_bit(&pruned_base, 44);
let pruned = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: BitmapBatch::Base(Arc::clone(&shared)),
overlay: Arc::new(overlay),
shared,
}));
for (name, chain, committed) in [
("flat", flat, 8),
("one-layer", one_layer, 8),
("two-layer", two_layer, 8),
("pruned-base", pruned, 40),
] {
let len = bitmap::Readable::<N>::len(&chain);
for tip in [committed, len, len + 3] {
assert_matches(name, &chain, tip);
}
let tip = len + 3;
let cap = tip as usize;
let pruned_bits = bitmap::Readable::<N>::pruned_bits(&chain);
for floor in pruned_bits..=committed {
let mut got = Vec::new();
let next = fill_candidates(&chain, Location::new(floor), committed, cap, &mut got);
fill_candidates(&chain, next, tip, cap, &mut got);
assert!(got.is_sorted_by(|a, b| a < b), "{name} floor={floor}");
for loc in floor..tip {
let must_emit = loc >= len || bitmap::Readable::<N>::get_bit(&chain, loc);
assert!(
!must_emit || got.contains(&Location::new(loc)),
"{name} floor={floor} lost {loc}"
);
}
}
}
}
#[test]
fn fill_candidates_filters_ancestor_clears() {
let bits = [true, false, true, true, false, false, true, false];
let base = make_bitmap(&bits);
let shared = Arc::new(Shared::new(make_bitmap(&bits)));
let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
overlay1.clear_bit(&base, 3);
overlay1.set_bit(&base, 9);
let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: BitmapBatch::Base(Arc::clone(&shared)),
overlay: Arc::new(overlay1),
shared: Arc::clone(&shared),
}));
let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
overlay2.clear_bit(&chain1, 6);
overlay2.clear_bit(&chain1, 9);
overlay2.set_bit(&chain1, 13);
let chain2 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: chain1.clone(),
overlay: Arc::new(overlay2),
shared,
}));
let scan = |chain: &BitmapBatch<N>, tip: u64| {
let mut got = Vec::new();
fill_candidates(chain, Location::new(0), tip, 16, &mut got);
got
};
let want = |locs: &[u64]| locs.iter().copied().map(Location::new).collect::<Vec<_>>();
assert_eq!(scan(&chain1, 12), want(&[0, 2, 6, 9]));
assert_eq!(scan(&chain2, 14), want(&[0, 2, 13]));
assert_eq!(scan(&chain2, 16), want(&[0, 2, 13, 14, 15]));
}
#[test]
fn fill_candidates_mixes_overlay_and_base_chunks() {
let mut bits = [false; 40];
for i in [1, 30, 33, 35, 38] {
bits[i] = true;
}
let base = make_bitmap(&bits);
let shared = Arc::new(Shared::new(make_bitmap(&bits)));
let mut overlay = ChunkOverlay::new(&base, 44, 1);
overlay.clear_bit(&base, 35);
overlay.set_bit(&base, 41);
let chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: BitmapBatch::Base(Arc::clone(&shared)),
overlay: Arc::new(overlay),
shared,
}));
let mut got = Vec::new();
fill_candidates(&chain, Location::new(0), 44, 16, &mut got);
let want: Vec<Location> = [1, 30, 33, 38, 41].into_iter().map(Location::new).collect();
assert_eq!(got, want);
}
fn make_chain(shared: &Arc<Shared<N>>, overlay_lens: &[u64]) -> BitmapBatch<N> {
let mut chain = BitmapBatch::Base(Arc::clone(shared));
for &len in overlay_lens {
let overlay = Arc::new(ChunkOverlay::new(&chain, len, 0));
chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
parent: chain,
overlay,
shared: Arc::clone(shared),
}));
}
chain
}
fn chain_overlays(batch: &BitmapBatch<N>) -> Vec<u64> {
let mut lens = Vec::new();
let mut current = batch;
while let BitmapBatch::Layer(layer) = current {
lens.push(layer.overlay.len);
current = &layer.parent;
}
assert!(matches!(current, BitmapBatch::Base(_)));
lens.reverse();
lens
}
#[test]
fn trim_committed_already_base() {
let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
let base = BitmapBatch::Base(Arc::clone(&shared));
let result = base.trim_committed();
match result {
BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
BitmapBatch::Layer(_) => panic!("expected Base"),
}
}
#[test]
fn trim_committed_all_committed() {
let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
let chain = make_chain(&shared, &[32]);
let result = chain.trim_committed();
match result {
BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
BitmapBatch::Layer(_) => panic!("expected Base after full trim"),
}
}
#[test]
fn trim_committed_none_committed() {
let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 32])));
let chain = make_chain(&shared, &[64, 96]);
let result = chain.trim_committed();
assert_eq!(chain_overlays(&result), vec![64, 96]);
}
#[test]
fn trim_committed_exactly_one_uncommitted() {
let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
let chain = make_chain(&shared, &[64, 96]);
let result = chain.trim_committed();
assert_eq!(chain_overlays(&result), vec![96]);
assert!(Arc::ptr_eq(result.shared(), &shared));
}
#[test]
fn trim_committed_multiple_uncommitted() {
let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
let chain = make_chain(&shared, &[64, 96, 128]);
let result = chain.trim_committed();
assert_eq!(chain_overlays(&result), vec![96, 128]);
assert!(Arc::ptr_eq(result.shared(), &shared));
}
}