1use crate::{
6 Context,
7 index::Unordered as UnorderedIndex,
8 journal::contiguous::{Contiguous, Mutable},
9 merkle::{
10 self, Graftable, Location, Position, Readable,
11 batch::MerkleizedBatch as GenericMerkleizedBatch, mem::Mem,
12 storage::Storage as MerkleStorage,
13 },
14 qmdb::{
15 Error,
16 any::{
17 self, ValueEncoding,
18 batch::{DiffCursors, DiffEntry, Staged as AnyStaged, StagedUpdates},
19 operation::{Operation, update},
20 },
21 batch_chain::Bounds,
22 bitmap::{Shared, fill_from},
23 current::{
24 db::{compute_db_root, partial_chunk, read_graft_inputs},
25 grafting,
26 },
27 operation::Key,
28 },
29};
30use ahash::AHashMap;
31use commonware_codec::Codec;
32use commonware_cryptography::{Digest, Hasher};
33use commonware_parallel::Strategy;
34use commonware_utils::bitmap::{self, Readable as _};
35use core::ops::Range;
36use std::sync::Arc;
37
38#[derive(Clone, Debug, Default)]
44pub(crate) struct ChunkOverlay<const N: usize> {
45 pub(crate) chunks: AHashMap<usize, [u8; N]>,
49 pub(crate) len: u64,
51 parent: Dimensions,
53}
54
55#[derive(Clone, Copy, Debug, Default)]
59struct Dimensions {
60 len: u64,
61 complete_chunks: usize,
62 pruned_chunks: usize,
63}
64
65impl Dimensions {
66 fn of<B: bitmap::Readable<N>, const N: usize>(base: &B) -> Self {
67 Self {
68 len: base.len(),
69 complete_chunks: base.complete_chunks(),
70 pruned_chunks: base.pruned_chunks(),
71 }
72 }
73}
74
75impl<const N: usize> ChunkOverlay<N> {
76 const CHUNK_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
77
78 fn new<B: bitmap::Readable<N>>(base: &B, len: u64, capacity: usize) -> Self {
81 Self {
82 chunks: AHashMap::with_capacity(capacity),
83 len,
84 parent: Dimensions::of(base),
85 }
86 }
87
88 fn chunk_mut<B: bitmap::Readable<N>>(&mut self, base: &B, idx: usize) -> &mut [u8; N] {
91 let parent = self.parent;
92 self.chunks.entry(idx).or_insert_with(|| {
93 let base_has_partial = !parent.len.is_multiple_of(Self::CHUNK_BITS);
94 if idx < parent.complete_chunks {
95 base.get_chunk(idx)
96 } else if idx == parent.complete_chunks && base_has_partial {
97 base.last_chunk().0
98 } else {
99 bitmap::BitMap::<N>::EMPTY_CHUNK
100 }
101 })
102 }
103
104 fn set_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
106 let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
107 let rel = (loc % Self::CHUNK_BITS) as usize;
108 let chunk = self.chunk_mut(base, idx);
109 chunk[rel / 8] |= 1 << (rel % 8);
110 }
111
112 fn clear_bit<B: bitmap::Readable<N>>(&mut self, base: &B, loc: u64) {
115 let idx = bitmap::Prunable::<N>::to_chunk_index(loc);
116 if idx < self.parent.pruned_chunks {
117 return;
118 }
119 let rel = (loc % Self::CHUNK_BITS) as usize;
120 let chunk = self.chunk_mut(base, idx);
121 chunk[rel / 8] &= !(1 << (rel % 8));
122 }
123
124 pub(crate) fn get(&self, idx: usize) -> Option<&[u8; N]> {
126 self.chunks.get(&idx)
127 }
128
129 pub(crate) const fn complete_chunks(&self) -> usize {
131 (self.len / Self::CHUNK_BITS) as usize
132 }
133}
134
135pub(crate) fn fill_candidates<F: Graftable, const N: usize>(
146 bitmap: &BitmapBatch<N>,
147 floor: Location<F>,
148 tip: u64,
149 limit: usize,
150 out: &mut Vec<Location<F>>,
151) -> Location<F> {
152 Location::new(fill_from(bitmap, *floor, tip, limit, out))
153}
154
155struct BatchStorageAdapter<
161 'a,
162 F: Graftable,
163 D: Digest,
164 R: Readable<Family = F, Digest = D>,
165 S: MerkleStorage<F, Digest = D>,
166> {
167 batch: &'a R,
168 base: &'a S,
169 _phantom: core::marker::PhantomData<(F, D)>,
170}
171
172impl<
173 'a,
174 F: Graftable,
175 D: Digest,
176 R: Readable<Family = F, Digest = D>,
177 S: MerkleStorage<F, Digest = D>,
178> BatchStorageAdapter<'a, F, D, R, S>
179{
180 const fn new(batch: &'a R, base: &'a S) -> Self {
181 Self {
182 batch,
183 base,
184 _phantom: core::marker::PhantomData,
185 }
186 }
187}
188
189impl<F: Graftable, D: Digest, R: Readable<Family = F, Digest = D>, S: MerkleStorage<F, Digest = D>>
190 MerkleStorage<F> for BatchStorageAdapter<'_, F, D, R, S>
191{
192 type Digest = D;
193
194 fn size(&self) -> Position<F> {
195 self.batch.size()
196 }
197 async fn get_node(&self, pos: Position<F>) -> Result<Option<D>, merkle::Error<F>> {
198 if let Some(node) = self.batch.get_node(pos) {
199 return Ok(Some(node));
200 }
201 self.base.get_node(pos).await
202 }
203
204 async fn get_nodes(&self, positions: &[Position<F>]) -> Result<Vec<D>, merkle::Error<F>> {
205 let mut nodes = vec![None; positions.len()];
206 let mut base_positions = Vec::with_capacity(positions.len());
207
208 for (slot, &pos) in nodes.iter_mut().zip(positions) {
210 match self.batch.get_node(pos) {
211 Some(node) => *slot = Some(node),
212 None => base_positions.push(pos),
213 }
214 }
215
216 let base_nodes = if base_positions.is_empty() {
218 Vec::new()
219 } else {
220 self.base.get_nodes(&base_positions).await?
221 };
222 let mut base_nodes = base_nodes.into_iter();
223 Ok(nodes
224 .into_iter()
225 .map(|node| node.unwrap_or_else(|| base_nodes.next().expect("one node per base read")))
226 .collect())
227 }
228}
229
230struct BatchOverMem<'a, F: Graftable, D: Digest, S: Strategy> {
235 batch: &'a GenericMerkleizedBatch<F, D, S>,
236 mem: &'a Mem<F, D>,
237}
238
239impl<F: Graftable, D: Digest, S: Strategy> Readable for BatchOverMem<'_, F, D, S> {
240 type Family = F;
241 type Digest = D;
242
243 fn size(&self) -> Position<F> {
244 self.batch.size()
245 }
246
247 fn get_node(&self, pos: Position<F>) -> Option<D> {
248 if let Some(d) = self.batch.get_node(pos) {
249 return Some(d);
250 }
251 self.mem.get_node(pos)
252 }
253}
254
255pub struct UnmerkleizedBatch<F, H, U, const N: usize, S: Strategy>
261where
262 F: Graftable,
263 U: update::Update,
264 H: Hasher,
265 Operation<F, U>: Codec,
266{
267 inner: any::batch::UnmerkleizedBatch<F, H, U, S>,
269
270 grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
272
273 bitmap_parent: BitmapBatch<N>,
275}
276
277pub struct Staged<F, H, U, const N: usize, S: Strategy>
279where
280 F: Graftable,
281 U: update::Update,
282 H: Hasher,
283 Operation<F, U>: Codec,
284{
285 inner: AnyStaged<F, H, U, S>,
286 grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
287 bitmap_parent: BitmapBatch<N>,
288}
289
290pub struct MerkleizedBatch<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
325{
326 pub(crate) inner: Arc<any::batch::MerkleizedBatch<F, D, U, S>>,
328
329 pub(crate) grafted: Arc<merkle::batch::MerkleizedBatch<F, D, S>>,
331
332 pub(crate) bitmap: BitmapBatch<N>,
334
335 pub(crate) canonical_root: D,
337}
338
339impl<F, H, U, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, U, N, S>
340where
341 F: Graftable,
342 U: update::Update,
343 H: Hasher,
344 Operation<F, U>: Codec,
345{
346 pub(super) const fn new(
347 inner: any::batch::UnmerkleizedBatch<F, H, U, S>,
348 grafted_parent: Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
349 bitmap_parent: BitmapBatch<N>,
350 ) -> Self {
351 Self {
352 inner,
353 grafted_parent,
354 bitmap_parent,
355 }
356 }
357
358 pub fn write(mut self, key: U::Key, value: Option<U::Value>) -> Self {
363 self.inner = self.inner.write(key, value);
364 self
365 }
366
367 pub async fn get<E, C, I>(
369 &self,
370 key: &U::Key,
371 db: &super::db::Db<F, E, C, I, H, U, N, S>,
372 ) -> Result<Option<U::Value>, Error<F>>
373 where
374 E: Context,
375 C: Contiguous<Item = Operation<F, U>>,
376 I: UnorderedIndex<Value = Location<F>> + 'static,
377 {
378 self.inner.get(key, &db.any).await
379 }
380
381 pub async fn get_many<E, C, I>(
389 &self,
390 keys: &[&U::Key],
391 db: &super::db::Db<F, E, C, I, H, U, N, S>,
392 ) -> Result<Vec<Option<U::Value>>, Error<F>>
393 where
394 E: Context,
395 C: Contiguous<Item = Operation<F, U>>,
396 I: UnorderedIndex<Value = Location<F>> + 'static,
397 {
398 self.inner.get_many(keys, &db.any).await
399 }
400
401 pub async fn stage<E, C, I>(
407 self,
408 keys: &[&U::Key],
409 db: &super::db::Db<F, E, C, I, H, U, N, S>,
410 ) -> Result<(Vec<Option<U::Value>>, Staged<F, H, U, N, S>), Error<F>>
411 where
412 E: Context,
413 C: Contiguous<Item = Operation<F, U>>,
414 I: UnorderedIndex<Value = Location<F>> + 'static,
415 {
416 let Self {
417 inner,
418 grafted_parent,
419 bitmap_parent,
420 } = self;
421 let (values, inner) = inner.stage(keys, &db.any).await?;
422 Ok((
423 values,
424 Staged {
425 inner,
426 grafted_parent,
427 bitmap_parent,
428 },
429 ))
430 }
431}
432
433impl<F, H, U, const N: usize, S: Strategy> Staged<F, H, U, N, S>
434where
435 F: Graftable,
436 U: update::Update,
437 H: Hasher,
438 Operation<F, U>: Codec,
439{
440 pub async fn expand<E, C, I>(
449 self,
450 keys: &[&U::Key],
451 db: &super::db::Db<F, E, C, I, H, U, N, S>,
452 ) -> Result<(Range<usize>, Vec<Option<U::Value>>, Self), Error<F>>
453 where
454 E: Context,
455 C: Contiguous<Item = Operation<F, U>>,
456 I: UnorderedIndex<Value = Location<F>> + 'static,
457 {
458 let Self {
459 inner,
460 grafted_parent,
461 bitmap_parent,
462 } = self;
463 let (range, values, inner) = inner.expand(keys, &db.any).await?;
464 Ok((
465 range,
466 values,
467 Self {
468 inner,
469 grafted_parent,
470 bitmap_parent,
471 },
472 ))
473 }
474}
475
476impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Unordered<K, V>, N, S>
477where
478 F: Graftable,
479 K: Key,
480 V: ValueEncoding,
481 H: Hasher,
482 Operation<F, update::Unordered<K, V>>: Codec,
483{
484 #[allow(clippy::type_complexity)]
497 #[tracing::instrument(
498 name = "qmdb.current.unordered.batch.merkleize.staged",
499 level = "info",
500 skip_all,
501 fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
502 )]
503 pub async fn merkleize<E, C, I>(
504 self,
505 updates: Vec<(usize, Option<V::Value>)>,
506 upserts: Vec<(K, Option<V::Value>)>,
507 metadata: Option<V::Value>,
508 db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
509 ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
510 where
511 E: Context,
512 C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
513 I: UnorderedIndex<Value = Location<F>> + 'static,
514 {
515 let Self {
516 inner,
517 grafted_parent,
518 bitmap_parent,
519 } = self;
520
521 let (inner, staged_updates, prefetched) = inner
525 .resolve_updates_prefetched(updates, upserts, &db.any, |floor, tip, limit, out| {
526 fill_candidates(&bitmap_parent, floor, tip, limit, out)
527 })
528 .await?;
529 let inner = inner
530 .merkleize_with_floor_scan(
531 &db.any,
532 metadata,
533 staged_updates,
534 Some(prefetched),
535 |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
536 )
537 .await?;
538 compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
539 }
540}
541
542impl<F, K, V, H, const N: usize, S: Strategy> Staged<F, H, update::Ordered<K, V>, N, S>
543where
544 F: Graftable,
545 K: Key,
546 V: ValueEncoding,
547 H: Hasher,
548 Operation<F, update::Ordered<K, V>>: Codec,
549{
550 #[allow(clippy::type_complexity)]
563 #[tracing::instrument(
564 name = "qmdb.current.ordered.batch.merkleize.staged",
565 level = "info",
566 skip_all,
567 fields(updates = updates.len() as u64, upserts = upserts.len() as u64),
568 )]
569 pub async fn merkleize<E, C, I>(
570 self,
571 updates: Vec<(usize, Option<V::Value>)>,
572 upserts: Vec<(K, Option<V::Value>)>,
573 metadata: Option<V::Value>,
574 db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
575 ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
576 where
577 E: Context,
578 C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
579 I: crate::index::Ordered<Value = Location<F>> + 'static,
580 {
581 let Self {
582 inner,
583 grafted_parent,
584 bitmap_parent,
585 } = self;
586 let (inner, staged_updates) = inner.resolve_updates(updates, upserts, db.any.strategy());
587 let inner = inner
588 .merkleize_with_floor_scan(
589 &db.any,
590 metadata,
591 staged_updates,
592 |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
593 )
594 .await?;
595 compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
596 }
597}
598
599impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
601where
602 F: Graftable,
603 K: Key,
604 V: ValueEncoding,
605 H: Hasher,
606 Operation<F, update::Unordered<K, V>>: Codec,
607{
608 #[allow(clippy::type_complexity)]
610 #[tracing::instrument(
611 name = "qmdb.current.unordered.batch.merkleize",
612 level = "info",
613 skip_all
614 )]
615 pub async fn merkleize<E, C, I>(
616 self,
617 db: &super::db::Db<F, E, C, I, H, update::Unordered<K, V>, N, S>,
618 metadata: Option<V::Value>,
619 ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>, Error<F>>
620 where
621 E: Context,
622 C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
623 I: UnorderedIndex<Value = Location<F>> + 'static,
624 {
625 let Self {
626 inner,
627 grafted_parent,
628 bitmap_parent,
629 } = self;
630 let inner = inner
632 .merkleize_with_floor_scan(
633 &db.any,
634 metadata,
635 StagedUpdates::<F, update::Unordered<K, V>>::new(),
636 None,
637 |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
638 )
639 .await?;
640 compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
641 }
642}
643
644impl<F, K, V, H, const N: usize, S: Strategy> UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
646where
647 F: Graftable,
648 K: Key,
649 V: ValueEncoding,
650 H: Hasher,
651 Operation<F, update::Ordered<K, V>>: Codec,
652{
653 #[allow(clippy::type_complexity)]
655 #[tracing::instrument(
656 name = "qmdb.current.ordered.batch.merkleize",
657 level = "info",
658 skip_all
659 )]
660 pub async fn merkleize<E, C, I>(
661 self,
662 db: &super::db::Db<F, E, C, I, H, update::Ordered<K, V>, N, S>,
663 metadata: Option<V::Value>,
664 ) -> Result<Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>, Error<F>>
665 where
666 E: Context,
667 C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
668 I: crate::index::Ordered<Value = Location<F>> + 'static,
669 {
670 let Self {
671 inner,
672 grafted_parent,
673 bitmap_parent,
674 } = self;
675 let inner = inner
677 .merkleize_with_floor_scan(
678 &db.any,
679 metadata,
680 StagedUpdates::<F, update::Ordered<K, V>>::new(),
681 |floor, tip, limit, out| fill_candidates(&bitmap_parent, floor, tip, limit, out),
682 )
683 .await?;
684 compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
685 }
686}
687
688#[allow(clippy::type_complexity)]
698fn build_chunk_overlay<F: Graftable, U, B: bitmap::Readable<N>, const N: usize>(
699 base: &B,
700 batch_len: usize,
701 batch_base: u64,
702 diff: &[(U::Key, DiffEntry<F, U::Value>)],
703 ancestor_diffs: &[Arc<Vec<(U::Key, DiffEntry<F, U::Value>)>>],
704) -> ChunkOverlay<N>
705where
706 U: update::Update,
707{
708 let total_bits = base.len() + batch_len as u64;
709 let appended_chunks = (batch_len as u64).div_ceil(ChunkOverlay::<N>::CHUNK_BITS) as usize;
710 let mut overlay = ChunkOverlay::new(base, total_bits, diff.len() + appended_chunks + 1);
711
712 let commit_loc = batch_base + batch_len as u64 - 1;
714 overlay.set_bit(base, commit_loc);
715
716 overlay.clear_bit(base, batch_base - 1);
718
719 let mut ancestors = DiffCursors::new(ancestor_diffs.iter().map(|d| d.as_slice()));
722 for (key, entry) in diff {
723 if let Some(loc) = entry.loc()
725 && *loc >= batch_base
726 && *loc < batch_base + batch_len as u64
727 {
728 overlay.set_bit(base, *loc);
729 }
730
731 let mut prev_loc = entry.base_old_loc();
734 if let Some(ancestor_entry) = ancestors.resolve(key) {
735 prev_loc = ancestor_entry.loc();
736 }
737 if let Some(old) = prev_loc {
738 overlay.clear_bit(base, *old);
739 }
740 }
741
742 let parent_complete = overlay.parent.complete_chunks;
746 let new_complete = overlay.complete_chunks();
747 for idx in parent_complete..new_complete {
748 overlay.chunk_mut(base, idx);
749 }
750
751 overlay
752}
753
754async fn merkleize_grafted_batch<F, H, S, const N: usize>(
756 strategy: &S,
757 grafted_parent: Arc<GenericMerkleizedBatch<F, H::Digest, S>>,
758 grafted_tree: &Arc<Mem<F, H::Digest>>,
759 graft_inputs: Vec<(usize, H::Digest, [u8; N])>,
760 grafting_height: u32,
761) -> Arc<GenericMerkleizedBatch<F, H::Digest, S>>
762where
763 F: Graftable,
764 H: Hasher,
765 S: Strategy,
766{
767 let old_grafted_leaves = *grafted_parent.leaves() as usize;
768 let mut grafted_batch = grafted_parent.new_batch();
769 let ancestors = grafted_batch.retain_ancestors();
770 let grafted_tree = Arc::clone(grafted_tree);
771 strategy
772 .clone()
773 .spawn(graft_inputs.len(), move |strategy| {
774 let new_leaves = grafting::graft_chunk_digests::<H, _, N>(&strategy, graft_inputs);
775 for (chunk_idx, digest) in new_leaves {
776 if chunk_idx < old_grafted_leaves {
777 grafted_batch = grafted_batch
778 .update_leaf_digest(Location::<F>::new(chunk_idx as u64), digest)
779 .expect("update_leaf_digest failed");
780 } else {
781 grafted_batch = grafted_batch.add_leaf_digest(digest);
782 }
783 }
784 let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
785 let merkleized = grafted_batch.merkleize(&grafted_tree, &grafted_hasher);
786 drop(ancestors);
787 merkleized
788 })
789 .await
790}
791
792async fn compute_current_layer<F, E, U, C, I, H, const N: usize, S>(
798 inner: Arc<any::batch::MerkleizedBatch<F, H::Digest, U, S>>,
799 current_db: &super::db::Db<F, E, C, I, H, U, N, S>,
800 grafted_parent: &Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>>,
801 bitmap_parent: &BitmapBatch<N>,
802) -> Result<Arc<MerkleizedBatch<F, H::Digest, U, N, S>>, Error<F>>
803where
804 F: Graftable,
805 E: Context,
806 C: Contiguous<Item = Operation<F, U>>,
807 I: UnorderedIndex<Value = Location<F>>,
808 H: Hasher,
809 U: update::Update,
810 S: Strategy,
811 Operation<F, U>: Codec,
812{
813 let batch_len = inner.journal_batch.items().len();
814 let batch_base = *inner.bounds.tip.size - batch_len as u64;
815
816 let overlay = build_chunk_overlay::<F, U, _, N>(
818 bitmap_parent,
819 batch_len,
820 batch_base,
821 &inner.diff,
822 &inner.ancestor_diffs,
823 );
824
825 let grafting_height = grafting::height::<N>();
826 let ops_tree_adapter =
827 BatchStorageAdapter::new(&inner.journal_batch, ¤t_db.any.log.merkle);
828
829 let overlay_ops_leaves = inner.bounds.tip.size;
832
833 let new_complete_chunks = overlay.complete_chunks();
841 let graftable_overlay = grafting::graftable_chunks::<F>(*overlay_ops_leaves, grafting_height)
842 .min(new_complete_chunks as u64) as usize;
843 let graftable_parent = *grafted_parent.leaves() as usize;
844 let pruned_chunks = bitmap_parent.pruned_chunks();
845 assert!(
846 pruned_chunks <= graftable_parent
847 && graftable_parent <= graftable_overlay
848 && graftable_overlay <= new_complete_chunks,
849 "invariant violated: pruned={pruned_chunks} graftable_parent={graftable_parent} graftable_overlay={graftable_overlay} new_complete={new_complete_chunks}"
850 );
851
852 let mut chunk_indices_to_update: Vec<usize> = overlay
858 .chunks
859 .iter()
860 .filter(|&(&idx, _)| idx < graftable_overlay && idx >= pruned_chunks)
861 .map(|(&idx, _)| idx)
862 .collect();
863 chunk_indices_to_update.extend(graftable_parent..graftable_overlay);
864 chunk_indices_to_update.sort_unstable();
865 chunk_indices_to_update.dedup();
866 let chunks_to_update = chunk_indices_to_update.into_iter().map(|idx| {
867 let chunk = overlay
868 .get(idx)
869 .copied()
870 .unwrap_or_else(|| bitmap_parent.get_chunk(idx));
871 (idx, chunk)
872 });
873
874 let graft_inputs = read_graft_inputs::<F, _, N>(&ops_tree_adapter, chunks_to_update).await?;
879 let grafted_batch = if graft_inputs.is_empty() {
880 let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
881 grafted_parent
882 .new_batch()
883 .merkleize(¤t_db.grafted_tree, &grafted_hasher)
884 } else {
885 merkleize_grafted_batch::<F, H, S, N>(
886 ¤t_db.strategy,
887 Arc::clone(grafted_parent),
888 ¤t_db.grafted_tree,
889 graft_inputs,
890 grafting_height,
891 )
892 .await
893 };
894
895 let bitmap_batch = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
899 parent: bitmap_parent.clone(),
900 overlay: Arc::new(overlay),
901 shared: Arc::clone(bitmap_parent.shared()),
902 }));
903
904 let ops_root = inner.root();
907 let layered = BatchOverMem {
908 batch: &grafted_batch,
909 mem: ¤t_db.grafted_tree,
910 };
911 let grafted_storage =
912 grafting::Storage::<F, H, _, _>::new(&layered, grafting_height, &ops_tree_adapter);
913 let partial = partial_chunk::<_, N>(&bitmap_batch);
919 let canonical_root = compute_db_root::<F, H, _, _, N>(
920 &bitmap_batch,
921 &grafted_storage,
922 overlay_ops_leaves,
923 partial,
924 inner.bounds.inactivity_floor,
925 &ops_root,
926 )
927 .await?;
928
929 Ok(Arc::new(MerkleizedBatch {
930 inner,
931 grafted: grafted_batch,
932 bitmap: bitmap_batch,
933 canonical_root,
934 }))
935}
936
937#[derive(Clone, Debug)]
943pub(crate) enum BitmapBatch<const N: usize> {
944 Base(Arc<Shared<N>>),
946 Layer(Arc<BitmapBatchLayer<N>>),
948}
949
950#[derive(Debug)]
952pub(crate) struct BitmapBatchLayer<const N: usize> {
953 pub(crate) parent: BitmapBatch<N>,
954 pub(crate) overlay: Arc<ChunkOverlay<N>>,
956 pub(crate) shared: Arc<Shared<N>>,
959}
960
961impl<const N: usize> BitmapBatch<N> {
962 const CHUNK_SIZE_BITS: u64 = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
963
964 fn shared(&self) -> &Arc<Shared<N>> {
966 match self {
967 Self::Base(s) => s,
968 Self::Layer(layer) => &layer.shared,
969 }
970 }
971
972 fn trim_committed(&self) -> Self {
976 let shared = self.shared();
977 let committed = bitmap::Readable::<N>::len(shared.as_ref());
978 let mut kept = Vec::new();
979 let mut current = self;
980 while let Self::Layer(layer) = current {
981 if layer.overlay.len <= committed {
982 break;
983 }
984 kept.push(Arc::clone(&layer.overlay));
985 current = &layer.parent;
986 }
987 let mut result = Self::Base(Arc::clone(shared));
988 for overlay in kept.into_iter().rev() {
989 result = Self::Layer(Arc::new(BitmapBatchLayer {
990 parent: result,
991 overlay,
992 shared: Arc::clone(shared),
993 }));
994 }
995 result
996 }
997}
998
999impl<const N: usize> bitmap::Readable<N> for BitmapBatch<N> {
1000 fn complete_chunks(&self) -> usize {
1001 (self.len() / Self::CHUNK_SIZE_BITS) as usize
1002 }
1003
1004 fn get_chunk(&self, idx: usize) -> [u8; N] {
1005 let mut current = self;
1008 loop {
1009 match current {
1010 Self::Base(shared) => return shared.get_chunk(idx),
1011 Self::Layer(layer) => {
1012 if let Some(&chunk) = layer.overlay.get(idx) {
1013 return chunk;
1014 }
1015 current = &layer.parent;
1016 }
1017 }
1018 }
1019 }
1020
1021 fn last_chunk(&self) -> ([u8; N], u64) {
1022 let total = self.len();
1023 if total == 0 {
1024 return (bitmap::BitMap::<N>::EMPTY_CHUNK, 0);
1025 }
1026 let rem = total % Self::CHUNK_SIZE_BITS;
1027 let bits_in_last = if rem == 0 { Self::CHUNK_SIZE_BITS } else { rem };
1028 let idx = if rem == 0 {
1029 self.complete_chunks().saturating_sub(1)
1030 } else {
1031 self.complete_chunks()
1032 };
1033 (self.get_chunk(idx), bits_in_last)
1034 }
1035
1036 fn pruned_chunks(&self) -> usize {
1037 self.shared().pruned_chunks()
1038 }
1039
1040 fn len(&self) -> u64 {
1041 match self {
1042 Self::Base(shared) => bitmap::Readable::<N>::len(shared.as_ref()),
1043 Self::Layer(layer) => layer.overlay.len,
1044 }
1045 }
1046}
1047
1048impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
1049 MerkleizedBatch<F, D, U, N, S>
1050{
1051 pub const fn root(&self) -> D {
1053 self.canonical_root
1054 }
1055
1056 pub fn ops_root(&self) -> D {
1058 self.inner.root()
1059 }
1060
1061 pub fn bounds(&self) -> &Bounds<F, D> {
1063 self.inner.bounds()
1064 }
1065
1066 pub fn operations(&self) -> (Location<F>, Arc<Vec<Operation<F, U>>>) {
1073 self.inner.operations()
1074 }
1075
1076 pub fn sync_boundary(&self) -> Location<F> {
1080 super::db::sync_boundary::<F, N>(
1084 *self.inner.bounds().inactivity_floor / bitmap::Prunable::<N>::CHUNK_SIZE_BITS,
1085 *self.inner.bounds().tip.size,
1086 )
1087 }
1088}
1089
1090impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
1091 MerkleizedBatch<F, D, U, N, S>
1092where
1093 Operation<F, U>: Codec,
1094{
1095 pub fn new_batch<H>(self: &Arc<Self>) -> UnmerkleizedBatch<F, H, U, N, S>
1105 where
1106 H: Hasher<Digest = D>,
1107 {
1108 UnmerkleizedBatch::new(
1109 self.inner.new_batch::<H>(),
1110 Arc::clone(&self.grafted),
1111 self.bitmap.trim_committed(),
1112 )
1113 }
1114
1115 pub async fn get<E, C, I, H>(
1120 &self,
1121 key: &U::Key,
1122 db: &super::db::Db<F, E, C, I, H, U, N, S>,
1123 ) -> Result<Option<U::Value>, Error<F>>
1124 where
1125 E: Context,
1126 C: Contiguous<Item = Operation<F, U>>,
1127 I: UnorderedIndex<Value = Location<F>> + 'static,
1128 H: Hasher<Digest = D>,
1129 {
1130 self.inner.get(key, &db.any).await
1131 }
1132
1133 pub async fn get_many<E, C, I, H>(
1137 &self,
1138 keys: &[&U::Key],
1139 db: &super::db::Db<F, E, C, I, H, U, N, S>,
1140 ) -> Result<Vec<Option<U::Value>>, Error<F>>
1141 where
1142 E: Context,
1143 C: Contiguous<Item = Operation<F, U>>,
1144 I: UnorderedIndex<Value = Location<F>> + 'static,
1145 H: Hasher<Digest = D>,
1146 {
1147 self.inner.get_many(keys, &db.any).await
1148 }
1149}
1150
1151impl<F, E, C, I, H, U, const N: usize, S> super::db::Db<F, E, C, I, H, U, N, S>
1152where
1153 F: Graftable,
1154 E: Context,
1155 C: Contiguous<Item = Operation<F, U>>,
1156 I: UnorderedIndex<Value = Location<F>>,
1157 H: Hasher,
1158 U: update::Update,
1159 S: Strategy,
1160 Operation<F, U>: Codec,
1161{
1162 pub fn to_batch(&self) -> Arc<MerkleizedBatch<F, H::Digest, U, N, S>> {
1168 let grafted = self.grafted_snapshot();
1169 Arc::new(MerkleizedBatch {
1170 inner: self.any.to_batch(),
1171 grafted,
1172 bitmap: BitmapBatch::Base(Arc::clone(&self.any.bitmap)),
1173 canonical_root: self.root,
1174 })
1175 }
1176}
1177
1178#[cfg(any(test, feature = "test-traits"))]
1179mod trait_impls {
1180 use super::*;
1181 use crate::{
1182 journal::contiguous::Mutable,
1183 qmdb::any::traits::{
1184 ApplyBatchResult, BatchableDb, MerkleizedBatch as MerkleizedBatchTrait,
1185 UnmerkleizedBatch as UnmerkleizedBatchTrait,
1186 },
1187 };
1188 use std::future::Future;
1189
1190 type CurrentDb<F, E, C, I, H, U, const N: usize, S> =
1191 crate::qmdb::current::db::Db<F, E, C, I, H, U, N, S>;
1192
1193 impl<F, K, V, H, E, C, I, const N: usize, S>
1194 UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>>
1195 for UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>
1196 where
1197 F: Graftable,
1198 K: Key,
1199 V: ValueEncoding + 'static,
1200 H: Hasher,
1201 E: Context,
1202 C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
1203 I: UnorderedIndex<Value = Location<F>> + 'static,
1204 S: Strategy,
1205 Operation<F, update::Unordered<K, V>>: Codec,
1206 {
1207 type Family = F;
1208 type K = K;
1209 type V = V::Value;
1210 type Metadata = V::Value;
1211 type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;
1212
1213 fn write(self, key: K, value: Option<V::Value>) -> Self {
1214 Self::write(self, key, value)
1215 }
1216
1217 async fn merkleize(
1218 self,
1219 db: &CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>,
1220 metadata: Option<V::Value>,
1221 ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
1222 self.merkleize(db, metadata).await
1223 }
1224 }
1225
1226 impl<F, K, V, H, E, C, I, const N: usize, S>
1227 UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>>
1228 for UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>
1229 where
1230 F: Graftable,
1231 K: Key,
1232 V: ValueEncoding + 'static,
1233 H: Hasher,
1234 E: Context,
1235 C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
1236 I: crate::index::Ordered<Value = Location<F>> + 'static,
1237 S: Strategy,
1238 Operation<F, update::Ordered<K, V>>: Codec,
1239 {
1240 type Family = F;
1241 type K = K;
1242 type V = V::Value;
1243 type Metadata = V::Value;
1244 type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;
1245
1246 fn write(self, key: K, value: Option<V::Value>) -> Self {
1247 Self::write(self, key, value)
1248 }
1249
1250 async fn merkleize(
1251 self,
1252 db: &CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>,
1253 metadata: Option<V::Value>,
1254 ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
1255 self.merkleize(db, metadata).await
1256 }
1257 }
1258
1259 impl<F: Graftable, D: Digest, U: update::Update, const N: usize, S: Strategy>
1260 MerkleizedBatchTrait for Arc<MerkleizedBatch<F, D, U, N, S>>
1261 where
1262 Operation<F, U>: Codec,
1263 {
1264 type Digest = D;
1265
1266 fn root(&self) -> D {
1267 MerkleizedBatch::root(self)
1268 }
1269 }
1270
1271 impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
1272 for CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N, S>
1273 where
1274 F: Graftable,
1275 E: Context,
1276 K: Key,
1277 V: ValueEncoding + 'static,
1278 C: Mutable<Item = Operation<F, update::Unordered<K, V>>>,
1279 I: UnorderedIndex<Value = Location<F>> + 'static,
1280 H: Hasher,
1281 S: Strategy,
1282 Operation<F, update::Unordered<K, V>>: Codec,
1283 {
1284 type Family = F;
1285 type K = K;
1286 type V = V::Value;
1287 type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Unordered<K, V>, N, S>>;
1288 type Batch = UnmerkleizedBatch<F, H, update::Unordered<K, V>, N, S>;
1289
1290 fn new_batch(&self) -> Self::Batch {
1291 self.new_batch()
1292 }
1293
1294 fn apply_batch(
1295 self,
1296 batch: Self::Merkleized,
1297 ) -> impl Future<Output = ApplyBatchResult<Self>> {
1298 self.apply_batch(batch)
1299 }
1300 }
1301
1302 impl<F, E, K, V, C, I, H, const N: usize, S> BatchableDb
1303 for CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N, S>
1304 where
1305 F: Graftable,
1306 E: Context,
1307 K: Key,
1308 V: ValueEncoding + 'static,
1309 C: Mutable<Item = Operation<F, update::Ordered<K, V>>>,
1310 I: crate::index::Ordered<Value = Location<F>> + 'static,
1311 H: Hasher,
1312 S: Strategy,
1313 Operation<F, update::Ordered<K, V>>: Codec,
1314 {
1315 type Family = F;
1316 type K = K;
1317 type V = V::Value;
1318 type Merkleized = Arc<MerkleizedBatch<F, H::Digest, update::Ordered<K, V>, N, S>>;
1319 type Batch = UnmerkleizedBatch<F, H, update::Ordered<K, V>, N, S>;
1320
1321 fn new_batch(&self) -> Self::Batch {
1322 self.new_batch()
1323 }
1324
1325 fn apply_batch(
1326 self,
1327 batch: Self::Merkleized,
1328 ) -> impl Future<Output = ApplyBatchResult<Self>> {
1329 self.apply_batch(batch)
1330 }
1331 }
1332}
1333
1334#[cfg(test)]
1335mod tests {
1336 use super::*;
1337 use crate::{mmb, mmr, utils::detached::block_strategy};
1338 use commonware_cryptography::Sha256;
1339 use commonware_macros::test_traced;
1340 use commonware_parallel::{Manual, Rayon};
1341 use commonware_utils::{NZUsize, bitmap::Prunable as BitMap};
1342 use std::{
1343 future::Future as _,
1344 task::Context as TaskContext,
1345 time::{Duration, Instant},
1346 };
1347
1348 const N: usize = 4;
1350 type Bm = BitMap<N>;
1351 type GraftedBatch =
1352 Arc<GenericMerkleizedBatch<mmb::Family, <Sha256 as Hasher>::Digest, Manual<Rayon>>>;
1353 type Location = mmr::Location;
1354
1355 fn make_bitmap(bits: &[bool]) -> Bm {
1356 let mut bm = Bm::new();
1357 for &b in bits {
1358 bm.push(b);
1359 }
1360 bm
1361 }
1362
1363 fn grafted_chain(
1364 strategy: &Manual<Rayon>,
1365 mem: &Arc<Mem<mmb::Family, <Sha256 as Hasher>::Digest>>,
1366 ) -> (GraftedBatch, GraftedBatch) {
1367 let hasher = grafting::hasher::<mmb::Family, Sha256>(grafting::height::<1>());
1368 let a = mem
1369 .new_batch_with_strategy(strategy.clone())
1370 .add_leaf_digest(Sha256::hash(&[b"a-0"]))
1371 .add_leaf_digest(Sha256::hash(&[b"a-1"]))
1372 .merkleize(mem, &hasher);
1373 let b = a
1374 .new_batch()
1375 .add_leaf_digest(Sha256::hash(&[b"b-0"]))
1376 .merkleize(mem, &hasher);
1377 (a, b)
1378 }
1379
1380 #[test_traced]
1382 fn test_grafted_merkleize_retains_ancestors_after_cancellation() {
1383 let strategy = Rayon::new(NZUsize!(2)).unwrap();
1384 let manual = strategy.manual();
1385 let mem = Arc::new(Mem::<mmb::Family, <Sha256 as Hasher>::Digest>::new());
1386 let grafting_height = grafting::height::<1>();
1387 let graft_inputs = || vec![(0, Sha256::hash(&[b"replacement"]), [1u8; 1])];
1388 let waker = futures::task::noop_waker();
1389 let mut context = TaskContext::from_waker(&waker);
1390
1391 let (a, b) = grafted_chain(&manual, &mem);
1393 let ancestor = Arc::downgrade(&a);
1394 let release = block_strategy(&strategy, 2);
1395 let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
1396 &manual,
1397 Arc::clone(&b),
1398 &mem,
1399 graft_inputs(),
1400 grafting_height,
1401 ));
1402 assert!(merkleize.as_mut().poll(&mut context).is_pending());
1403 drop(b);
1404 drop(a);
1405 drop(release);
1406 let _ = futures::executor::block_on(merkleize);
1407 assert!(ancestor.upgrade().is_none());
1408
1409 let (a, b) = grafted_chain(&manual, &mem);
1411 let ancestor = Arc::downgrade(&a);
1412 let release = block_strategy(&strategy, 2);
1413 let mut merkleize = Box::pin(merkleize_grafted_batch::<mmb::Family, Sha256, _, 1>(
1414 &manual,
1415 Arc::clone(&b),
1416 &mem,
1417 graft_inputs(),
1418 grafting_height,
1419 ));
1420 assert!(merkleize.as_mut().poll(&mut context).is_pending());
1421 drop(merkleize);
1422 drop(b);
1423 drop(a);
1424 assert!(ancestor.upgrade().is_some());
1425
1426 drop(release);
1427 let deadline = Instant::now() + Duration::from_secs(10);
1428 while ancestor.upgrade().is_some() {
1429 assert!(
1430 Instant::now() < deadline,
1431 "detached grafted merkleization did not release its ancestors"
1432 );
1433 std::thread::yield_now();
1434 }
1435 }
1436
1437 #[test]
1440 fn chunk_overlay_pushes() {
1441 use crate::qmdb::any::value::FixedEncoding;
1442 use commonware_utils::sequence::FixedBytes;
1443
1444 type K = FixedBytes<4>;
1445 type V = FixedEncoding<u64>;
1446 type U = crate::qmdb::any::operation::update::Unordered<K, V>;
1447
1448 let key1 = FixedBytes::from([1, 0, 0, 0]);
1449 let key2 = FixedBytes::from([2, 0, 0, 0]);
1450
1451 let base = make_bitmap(&[true; 4]);
1456 let mut diff = vec![
1457 (
1458 key1,
1459 DiffEntry::Active {
1460 value: 100u64,
1461 loc: Location::new(4), base_old_loc: None,
1463 },
1464 ),
1465 (
1466 key2,
1467 DiffEntry::Active {
1468 value: 200u64,
1469 loc: Location::new(99), base_old_loc: None,
1471 },
1472 ),
1473 ];
1474 diff.sort_by(|a, b| a.0.cmp(&b.0));
1475
1476 let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 4, 4, &diff, &[]);
1477
1478 let c0 = overlay.get(0).expect("chunk 0 should be dirty");
1481 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); }
1487
1488 #[test]
1489 fn chunk_overlay_clears() {
1490 use crate::qmdb::any::value::FixedEncoding;
1491 use commonware_utils::sequence::FixedBytes;
1492
1493 type K = FixedBytes<4>;
1494 type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;
1495
1496 let key1 = FixedBytes::from([1, 0, 0, 0]);
1497 let key2 = FixedBytes::from([2, 0, 0, 0]);
1498 let key3 = FixedBytes::from([3, 0, 0, 0]);
1499
1500 let base = make_bitmap(&[true; 64]);
1502
1503 let mut diff: Vec<(K, DiffEntry<mmr::Family, u64>)> = vec![
1504 (
1505 key1,
1506 DiffEntry::Active {
1507 value: 100,
1508 loc: Location::new(70),
1509 base_old_loc: Some(Location::new(5)),
1510 },
1511 ),
1512 (
1513 key2,
1514 DiffEntry::Deleted {
1515 base_old_loc: Some(Location::new(10)),
1516 },
1517 ),
1518 (
1519 key3,
1520 DiffEntry::Active {
1521 value: 300,
1522 loc: Location::new(71),
1523 base_old_loc: None,
1524 },
1525 ),
1526 ];
1527 diff.sort_by(|a, b| a.0.cmp(&b.0));
1528
1529 let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 8, 64, &diff, &[]);
1531
1532 let c0 = overlay.get(0).expect("chunk 0 should be dirty");
1534 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); }
1541
1542 #[test]
1546 fn chunk_overlay_preserves_partial_parent_chunk() {
1547 use crate::qmdb::any::value::FixedEncoding;
1548 use commonware_utils::sequence::FixedBytes;
1549
1550 type K = FixedBytes<4>;
1551 type U = crate::qmdb::any::operation::update::Unordered<K, FixedEncoding<u64>>;
1552
1553 let base = make_bitmap(&[true; 20]);
1555 assert_eq!(base.complete_chunks(), 0); let key1 = FixedBytes::from([1, 0, 0, 0]);
1561 let mut diff = vec![(
1562 key1,
1563 DiffEntry::Active {
1564 value: 42u64,
1565 loc: Location::new(35),
1566 base_old_loc: None,
1567 },
1568 )];
1569 diff.sort_by(|a, b| a.0.cmp(&b.0));
1570
1571 let overlay = build_chunk_overlay::<mmr::Family, U, _, N>(&base, 20, 20, &diff, &[]);
1572
1573 let c0 = overlay.get(0).expect("chunk 0 should be in overlay");
1575 assert_eq!(c0[0], 0xFF);
1577 assert_eq!(c0[1], 0xFF);
1579 assert_eq!(c0[2], 0x07);
1581 }
1582
1583 fn next_candidate<B: bitmap::Readable<N2>, const N2: usize>(
1589 bitmap: &B,
1590 floor: Location,
1591 tip: u64,
1592 ) -> Option<Location> {
1593 let floor = *floor;
1594 let bitmap_len = bitmap.len();
1595 let committed_end = bitmap_len.min(tip);
1596 if floor < committed_end
1597 && let Some(idx) = bitmap.ones_iter_from(floor).next()
1598 && idx < committed_end
1599 {
1600 return Some(Location::new(idx));
1601 }
1602 let candidate = floor.max(bitmap_len);
1603 (candidate < tip).then(|| Location::new(candidate))
1604 }
1605
1606 #[test]
1607 fn bitmap_scan_all_active() {
1608 let bm = make_bitmap(&[true; 8]);
1609 for i in 0..8 {
1610 assert_eq!(
1611 next_candidate(&bm, Location::new(i), 8),
1612 Some(Location::new(i))
1613 );
1614 }
1615 assert_eq!(next_candidate(&bm, Location::new(8), 8), None);
1616 }
1617
1618 #[test]
1619 fn bitmap_scan_all_inactive() {
1620 let bm = make_bitmap(&[false; 8]);
1621 assert_eq!(next_candidate(&bm, Location::new(0), 8), None);
1622 }
1623
1624 #[test]
1625 fn bitmap_scan_skips_inactive() {
1626 let bm = make_bitmap(&[false, false, true, false, true]);
1628 assert_eq!(
1629 next_candidate(&bm, Location::new(0), 5),
1630 Some(Location::new(2))
1631 );
1632 assert_eq!(
1633 next_candidate(&bm, Location::new(3), 5),
1634 Some(Location::new(4))
1635 );
1636 assert_eq!(next_candidate(&bm, Location::new(5), 5), None);
1637 }
1638
1639 #[test]
1640 fn bitmap_scan_beyond_bitmap_len_returns_candidate() {
1641 let bm = make_bitmap(&[false; 4]);
1644 assert_eq!(
1646 next_candidate(&bm, Location::new(0), 8),
1647 Some(Location::new(4))
1648 );
1649 assert_eq!(
1650 next_candidate(&bm, Location::new(6), 8),
1651 Some(Location::new(6))
1652 );
1653 }
1654
1655 #[test]
1656 fn bitmap_scan_respects_tip() {
1657 let bm = make_bitmap(&[false, false, false, true]);
1658 assert_eq!(next_candidate(&bm, Location::new(0), 3), None);
1660 assert_eq!(
1662 next_candidate(&bm, Location::new(0), 4),
1663 Some(Location::new(3))
1664 );
1665 }
1666
1667 #[test]
1668 fn bitmap_scan_floor_at_tip() {
1669 let bm = make_bitmap(&[true; 4]);
1670 assert_eq!(next_candidate(&bm, Location::new(4), 4), None);
1671 }
1672
1673 #[test]
1674 fn bitmap_scan_empty_bitmap() {
1675 let bm = Bm::new();
1676 assert_eq!(
1678 next_candidate(&bm, Location::new(0), 5),
1679 Some(Location::new(0))
1680 );
1681 assert_eq!(next_candidate(&bm, Location::new(0), 0), None);
1683 }
1684
1685 #[test]
1686 fn fill_candidates_matches_oracle() {
1687 fn assert_matches(name: &str, chain: &BitmapBatch<N>, tip: u64) {
1691 for floor in 0..=tip {
1692 let mut want = Vec::new();
1693 let mut scan = Location::new(floor);
1694 while let Some(c) = next_candidate(chain, scan, tip) {
1695 want.push(c);
1696 scan = c + 1;
1697 }
1698 for split in 0..=want.len() {
1699 let mut got = Vec::new();
1700 let next = fill_candidates(chain, Location::new(floor), tip, split, &mut got);
1701 fill_candidates(chain, next, tip, want.len() + 1, &mut got);
1702 assert_eq!(got, want, "{name} floor={floor} split={split}");
1703 }
1704 }
1705 }
1706
1707 let bits = [true, false, true, true, false, false, true, false];
1708 let base = make_bitmap(&bits);
1709
1710 let flat = BitmapBatch::Base(Arc::new(Shared::new(make_bitmap(&bits))));
1712
1713 let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1715 let mut overlay = ChunkOverlay::new(&base, 12, 1);
1716 overlay.clear_bit(&base, 3);
1717 overlay.clear_bit(&base, 6);
1718 overlay.set_bit(&base, 9);
1719 let one_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1720 parent: BitmapBatch::Base(Arc::clone(&shared)),
1721 overlay: Arc::new(overlay),
1722 shared,
1723 }));
1724
1725 let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1727 let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
1728 overlay1.clear_bit(&base, 3);
1729 overlay1.set_bit(&base, 9);
1730 let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1731 parent: BitmapBatch::Base(Arc::clone(&shared)),
1732 overlay: Arc::new(overlay1),
1733 shared: Arc::clone(&shared),
1734 }));
1735 let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
1736 overlay2.clear_bit(&chain1, 6);
1737 overlay2.clear_bit(&chain1, 9);
1738 overlay2.set_bit(&chain1, 13);
1739 let two_layer = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1740 parent: chain1,
1741 overlay: Arc::new(overlay2),
1742 shared,
1743 }));
1744
1745 let make_pruned = || {
1748 let mut bits = [false; 40];
1749 bits[33] = true;
1750 bits[38] = true;
1751 let mut bm = make_bitmap(&bits);
1752 bm.prune_to_bit(32);
1753 bm
1754 };
1755 let pruned_base = make_pruned();
1756 let shared = Arc::new(Shared::new(make_pruned()));
1757 let mut overlay = ChunkOverlay::new(&pruned_base, 46, 1);
1758 overlay.clear_bit(&pruned_base, 38);
1759 overlay.set_bit(&pruned_base, 41);
1760 overlay.set_bit(&pruned_base, 44);
1761 let pruned = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1762 parent: BitmapBatch::Base(Arc::clone(&shared)),
1763 overlay: Arc::new(overlay),
1764 shared,
1765 }));
1766
1767 for (name, chain, committed) in [
1768 ("flat", flat, 8),
1769 ("one-layer", one_layer, 8),
1770 ("two-layer", two_layer, 8),
1771 ("pruned-base", pruned, 40),
1772 ] {
1773 let len = bitmap::Readable::<N>::len(&chain);
1774 for tip in [committed, len, len + 3] {
1775 assert_matches(name, &chain, tip);
1776 }
1777
1778 let tip = len + 3;
1784 let cap = tip as usize;
1785 let pruned_bits = bitmap::Readable::<N>::pruned_bits(&chain);
1786 for floor in pruned_bits..=committed {
1787 let mut got = Vec::new();
1788 let next = fill_candidates(&chain, Location::new(floor), committed, cap, &mut got);
1789 fill_candidates(&chain, next, tip, cap, &mut got);
1790 assert!(got.is_sorted_by(|a, b| a < b), "{name} floor={floor}");
1791 for loc in floor..tip {
1792 let must_emit = loc >= len || bitmap::Readable::<N>::get_bit(&chain, loc);
1793 assert!(
1794 !must_emit || got.contains(&Location::new(loc)),
1795 "{name} floor={floor} lost {loc}"
1796 );
1797 }
1798 }
1799 }
1800 }
1801
1802 #[test]
1803 fn fill_candidates_filters_ancestor_clears() {
1804 let bits = [true, false, true, true, false, false, true, false];
1805 let base = make_bitmap(&bits);
1806 let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1807
1808 let mut overlay1 = ChunkOverlay::new(&base, 12, 2);
1810 overlay1.clear_bit(&base, 3);
1811 overlay1.set_bit(&base, 9);
1812 let chain1 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1813 parent: BitmapBatch::Base(Arc::clone(&shared)),
1814 overlay: Arc::new(overlay1),
1815 shared: Arc::clone(&shared),
1816 }));
1817
1818 let mut overlay2 = ChunkOverlay::new(&chain1, 14, 2);
1821 overlay2.clear_bit(&chain1, 6);
1822 overlay2.clear_bit(&chain1, 9);
1823 overlay2.set_bit(&chain1, 13);
1824 let chain2 = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1825 parent: chain1.clone(),
1826 overlay: Arc::new(overlay2),
1827 shared,
1828 }));
1829
1830 let scan = |chain: &BitmapBatch<N>, tip: u64| {
1835 let mut got = Vec::new();
1836 fill_candidates(chain, Location::new(0), tip, 16, &mut got);
1837 got
1838 };
1839 let want = |locs: &[u64]| locs.iter().copied().map(Location::new).collect::<Vec<_>>();
1840 assert_eq!(scan(&chain1, 12), want(&[0, 2, 6, 9]));
1841 assert_eq!(scan(&chain2, 14), want(&[0, 2, 13]));
1842 assert_eq!(scan(&chain2, 16), want(&[0, 2, 13, 14, 15]));
1843 }
1844
1845 #[test]
1846 fn fill_candidates_mixes_overlay_and_base_chunks() {
1847 let mut bits = [false; 40];
1849 for i in [1, 30, 33, 35, 38] {
1850 bits[i] = true;
1851 }
1852 let base = make_bitmap(&bits);
1853 let shared = Arc::new(Shared::new(make_bitmap(&bits)));
1854
1855 let mut overlay = ChunkOverlay::new(&base, 44, 1);
1859 overlay.clear_bit(&base, 35);
1860 overlay.set_bit(&base, 41);
1861 let chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1862 parent: BitmapBatch::Base(Arc::clone(&shared)),
1863 overlay: Arc::new(overlay),
1864 shared,
1865 }));
1866
1867 let mut got = Vec::new();
1870 fill_candidates(&chain, Location::new(0), 44, 16, &mut got);
1871 let want: Vec<Location> = [1, 30, 33, 38, 41].into_iter().map(Location::new).collect();
1872 assert_eq!(got, want);
1873 }
1874
1875 fn make_chain(shared: &Arc<Shared<N>>, overlay_lens: &[u64]) -> BitmapBatch<N> {
1888 let mut chain = BitmapBatch::Base(Arc::clone(shared));
1889 for &len in overlay_lens {
1890 let overlay = Arc::new(ChunkOverlay::new(&chain, len, 0));
1891 chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
1892 parent: chain,
1893 overlay,
1894 shared: Arc::clone(shared),
1895 }));
1896 }
1897 chain
1898 }
1899
1900 fn chain_overlays(batch: &BitmapBatch<N>) -> Vec<u64> {
1904 let mut lens = Vec::new();
1905 let mut current = batch;
1906 while let BitmapBatch::Layer(layer) = current {
1907 lens.push(layer.overlay.len);
1908 current = &layer.parent;
1909 }
1910 assert!(matches!(current, BitmapBatch::Base(_)));
1911 lens.reverse();
1912 lens
1913 }
1914
1915 #[test]
1920 fn trim_committed_already_base() {
1921 let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1922 let base = BitmapBatch::Base(Arc::clone(&shared));
1923 let result = base.trim_committed();
1924 match result {
1926 BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
1927 BitmapBatch::Layer(_) => panic!("expected Base"),
1928 }
1929 }
1930
1931 #[test]
1937 fn trim_committed_all_committed() {
1938 let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1940 let chain = make_chain(&shared, &[32]);
1941 let result = chain.trim_committed();
1942 match result {
1944 BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
1945 BitmapBatch::Layer(_) => panic!("expected Base after full trim"),
1946 }
1947 }
1948
1949 #[test]
1954 fn trim_committed_none_committed() {
1955 let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 32])));
1957 let chain = make_chain(&shared, &[64, 96]);
1958 let result = chain.trim_committed();
1959 assert_eq!(chain_overlays(&result), vec![64, 96]);
1961 }
1962
1963 #[test]
1969 fn trim_committed_exactly_one_uncommitted() {
1970 let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1972 let chain = make_chain(&shared, &[64, 96]);
1973 let result = chain.trim_committed();
1974 assert_eq!(chain_overlays(&result), vec![96]);
1976 assert!(Arc::ptr_eq(result.shared(), &shared));
1978 }
1979
1980 #[test]
1985 fn trim_committed_multiple_uncommitted() {
1986 let shared = Arc::new(Shared::<N>::new(make_bitmap(&[true; 64])));
1988 let chain = make_chain(&shared, &[64, 96, 128]);
1989 let result = chain.trim_committed();
1990 assert_eq!(chain_overlays(&result), vec![96, 128]);
1992 assert!(Arc::ptr_eq(result.shared(), &shared));
1994 }
1995}