PR 3627 nits

Author: danlaine

storage/src/qmdb/benches/chained_growth.rs

@@ -3,8 +3,8 @@
 //! Setup (untimed): seed `NUM_KEYS` keys, then grow a chain of `PREBUILT_CHAIN` batches applying
 //! each parent while the child is still alive.
 //!
-//! Timed: do `batches` more merkleize + apply iterations on top of the pre-built chain, with a
-//! single random update per batch so each overlay covers a tiny fraction of chunks.
+//! Timed: do `batches` more merkleize + apply iterations on top of the pre-built chain, with
+//! `updates` random updates per batch (varied across `UPDATES_PER_BATCH_SET`).
 
 use crate::common::{make_fixed_value, Digest, WRITE_BUFFER_SIZE};
 use commonware_cryptography::{Hasher, Sha256};
@@ -87,17 +87,17 @@ fn cur_fix_cfg(
 /// Number of pre-populated keys in the seeded database.
 const NUM_KEYS: u64 = 1_000_000;
 
-/// Random updates per batch. One update means each batch's chunk overlay covers ~1 / num_chunks
-/// of the bitmap, forcing chain reads to walk deep before finding a matching layer.
-const UPDATES_PER_BATCH: u64 = 1;
-
 /// Number of batches grown during the untimed seed phase, producing a Db::status chain of this
 /// depth that subsequent reads must walk through.
-const PREBUILT_CHAIN: u64 = 10_000;
+const PREBUILT_CHAIN: u64 = 1_000;
 
 /// Number of additional batches to grow during the timed region.
 const GROW_COUNTS: [u64; 1] = [100];
 
+/// Random updates per batch. The sparse value exercises the chain-walk read path; the dense value
+/// exercises the merkleize path under many per-batch chunk writes.
+const UPDATES_PER_BATCH_SET: [u64; 2] = [1, 10_000];
+
 #[derive(Debug, Clone, Copy)]
 enum CurrentVariant {
     UnorderedFixed32,
@@ -191,18 +191,19 @@ async fn run_chained_growth<
     Fork: Fn(&C::Merkleized) -> C::Batch,
 >(
     mut db: C,
+    updates_per_batch: u64,
     grow: u64,
     fork_child: Fork,
 ) -> Duration {
     seed_db(&mut db, NUM_KEYS).await;
     let mut rng = StdRng::seed_from_u64(99);
 
     // Pre-build a deep chain (untimed).
-    let initial = write_random_updates(db.new_batch(), UPDATES_PER_BATCH, NUM_KEYS, &mut rng);
+    let initial = write_random_updates(db.new_batch(), updates_per_batch, NUM_KEYS, &mut rng);
     let mut parent = initial.merkleize(&db, None).await.unwrap();
     for _ in 0..PREBUILT_CHAIN {
         let child_batch =
-            write_random_updates(fork_child(&parent), UPDATES_PER_BATCH, NUM_KEYS, &mut rng);
+            write_random_updates(fork_child(&parent), updates_per_batch, NUM_KEYS, &mut rng);
         let child = child_batch.merkleize(&db, None).await.unwrap();
         db.apply_batch(parent).await.unwrap();
         parent = child;
@@ -216,7 +217,7 @@ async fn run_chained_growth<
     let start = Instant::now();
     for _ in 0..grow {
         let child_batch =
-            write_random_updates(fork_child(&parent), UPDATES_PER_BATCH, NUM_KEYS, &mut rng);
+            write_random_updates(fork_child(&parent), updates_per_batch, NUM_KEYS, &mut rng);
         let child = child_batch.merkleize(&db, None).await.unwrap();
         black_box(child.root());
         db.apply_batch(parent).await.unwrap();
@@ -232,26 +233,30 @@ async fn run_chained_growth<
 fn bench_chained_growth(c: &mut Criterion) {
     let runner = tokio::Runner::new(Config::default());
     for batches in GROW_COUNTS {
-        for &variant in &CURRENT_VARIANTS {
-            c.bench_function(
-                &format!(
-                    "{}/variant={} batches={batches}",
-                    module_path!(),
-                    variant.name()
-                ),
-                |b| {
-                    b.to_async(&runner).iter_custom(|iters| async move {
-                        let ctx = context::get::<Context>();
-                        let mut total = Duration::ZERO;
-                        for _ in 0..iters {
-                            with_current_db!(ctx.clone(), variant, |mut db| {
-                                total += run_chained_growth(db, batches, |p| p.new_batch()).await;
-                            });
-                        }
-                        total
-                    });
-                },
-            );
+        for updates in UPDATES_PER_BATCH_SET {
+            for &variant in &CURRENT_VARIANTS {
+                c.bench_function(
+                    &format!(
+                        "{}/variant={} updates={updates} batches={batches}",
+                        module_path!(),
+                        variant.name()
+                    ),
+                    |b| {
+                        b.to_async(&runner).iter_custom(|iters| async move {
+                            let ctx = context::get::<Context>();
+                            let mut total = Duration::ZERO;
+                            for _ in 0..iters {
+                                with_current_db!(ctx.clone(), variant, |mut db| {
+                                    total +=
+                                        run_chained_growth(db, updates, batches, |p| p.new_batch())
+                                            .await;
+                                });
+                            }
+                            total
+                        });
+                    },
+                );
+            }
         }
     }
 }

storage/src/qmdb/current/batch.rs

@@ -634,6 +634,7 @@ where
     let bitmap_batch = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
         parent: bitmap_parent.clone(),
         overlay: Arc::new(overlay),
+        shared: Arc::clone(bitmap_parent.shared()),
     }));
 
     // Compute canonical root. The grafted batch alone cannot resolve committed nodes,
@@ -763,19 +764,19 @@ pub(crate) struct BitmapBatchLayer<const N: usize> {
     pub(crate) parent: BitmapBatch<N>,
     /// Chunk-level overlay: materialized bytes for every chunk that differs from parent.
     pub(crate) overlay: Arc<ChunkOverlay<N>>,
+    /// Cached terminal [`SharedBitmap`] so [`BitmapBatch::shared`] and
+    /// [`BitmapBatch::pruned_chunks`] answer in O(1) instead of walking the chain.
+    pub(crate) shared: Arc<SharedBitmap<N>>,
 }
 
 impl<const N: usize> BitmapBatch<N> {
     const CHUNK_SIZE_BITS: u64 = BitMap::<N>::CHUNK_SIZE_BITS;
 
-    /// Walk to the terminal [`SharedBitmap`] at the bottom of the chain.
+    /// Return the terminal [`SharedBitmap`] at the bottom of the chain.
     fn shared(&self) -> &Arc<SharedBitmap<N>> {
-        let mut current = self;
-        loop {
-            match current {
-                Self::Base(s) => return s,
-                Self::Layer(layer) => current = &layer.parent,
-            }
+        match self {
+            Self::Base(s) => s,
+            Self::Layer(layer) => &layer.shared,
         }
     }
 
@@ -817,6 +818,7 @@ impl<const N: usize> BitmapBatch<N> {
                 return Self::Layer(Arc::new(BitmapBatchLayer {
                     parent: Self::Base(Arc::clone(shared)),
                     overlay: Arc::clone(&layer.overlay),
+                    shared: Arc::clone(shared),
                 }));
             }
         }
@@ -835,6 +837,7 @@ impl<const N: usize> BitmapBatch<N> {
             result = Self::Layer(Arc::new(BitmapBatchLayer {
                 parent: result,
                 overlay,
+                shared: Arc::clone(shared),
             }));
         }
         result
@@ -879,13 +882,7 @@ impl<const N: usize> BitmapReadable<N> for BitmapBatch<N> {
     }
 
     fn pruned_chunks(&self) -> usize {
-        let mut current = self;
-        loop {
-            match current {
-                Self::Base(shared) => return shared.pruned_chunks(),
-                Self::Layer(layer) => current = &layer.parent,
-            }
-        }
+        self.shared().pruned_chunks()
     }
 
     fn len(&self) -> u64 {
@@ -1416,4 +1413,127 @@ mod tests {
         // Empty bitmap, tip = 0: no candidates.
         assert_eq!(scan.next_candidate(Location::new(0), 0), None);
     }
+
+    // ---- trim_committed tests ----
+    //
+    // `trim_committed` is called from `MerkleizedBatch::new_batch` to strip any `Layer`s whose
+    // overlays have already been absorbed into the shared committed bitmap by a prior apply.
+    // It has five branches; each test below exercises one of them directly, without going
+    // through the full Db/batch machinery, so the function's structural output can be asserted.
+
+    /// Build a chain `Base(shared) -> Layer(len=L1) -> Layer(len=L2) -> ...` from a list of
+    /// overlay lengths (bottom to top). Each constructed `Layer` caches `shared` per the
+    /// struct's invariant.
+    fn make_chain(shared: &Arc<SharedBitmap<N>>, overlay_lens: &[u64]) -> BitmapBatch<N> {
+        let mut chain = BitmapBatch::Base(Arc::clone(shared));
+        for &len in overlay_lens {
+            chain = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
+                parent: chain,
+                overlay: Arc::new(ChunkOverlay::new(len)),
+                shared: Arc::clone(shared),
+            }));
+        }
+        chain
+    }
+
+    /// Walk a chain and return its overlay lengths in bottom-to-top order. Used to assert the
+    /// structural output of `trim_committed` without touching private fields. Panics if the
+    /// chain isn't terminated by a single `Base` at the bottom.
+    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
+    }
+
+    /// Branch 1: input is already a bare `Base` with no speculative layers on top. This is the
+    /// input `MerkleizedBatch::new_batch` sees when a caller builds a child off a batch whose
+    /// chain was previously trimmed flat (e.g., immediately after apply collapsed everything).
+    /// `trim_committed` must short-circuit and return an equivalent `Base` pointing at the
+    /// same `SharedBitmap`; no new allocation should be needed for the terminal.
+    #[test]
+    fn trim_committed_already_base() {
+        let shared = Arc::new(SharedBitmap::<N>::new(make_bitmap(&[true; 64])));
+        let base = BitmapBatch::Base(Arc::clone(&shared));
+        let result = base.trim_committed();
+        // Must still be `Base`, and pointer-equal to the input's shared terminal.
+        match result {
+            BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
+            BitmapBatch::Layer(_) => panic!("expected Base"),
+        }
+    }
+
+    /// Branch 2: the shared committed bitmap has caught up to (or past) the top layer's
+    /// `overlay.len`, meaning every layer in the chain has been absorbed by prior applies.
+    /// This is the steady-state "extend a just-applied batch" flow: after `apply_batch(A)`,
+    /// `A`'s own layer has `overlay.len == committed`. `trim_committed` must collapse the
+    /// entire chain down to a bare `Base` so the new child starts from a clean terminal.
+    #[test]
+    fn trim_committed_all_committed() {
+        // `shared.len() == 64`; the single layer's `overlay.len == 32 (<= 64)`, so it's committed.
+        let shared = Arc::new(SharedBitmap::<N>::new(make_bitmap(&[true; 64])));
+        let chain = make_chain(&shared, &[32]);
+        let result = chain.trim_committed();
+        // Collapsed to a bare Base, pointing at the original shared.
+        match result {
+            BitmapBatch::Base(s) => assert!(Arc::ptr_eq(&s, &shared)),
+            BitmapBatch::Layer(_) => panic!("expected Base after full trim"),
+        }
+    }
+
+    /// Branch 3: every layer in the chain is still speculative (none have been applied), so
+    /// the committed bitmap hasn't caught up to any layer's `overlay.len`. `trim_committed`
+    /// must return the chain unchanged. This is the "speculate ahead without applying" case:
+    /// a caller building multiple batches deep (A, then B off A, then C off B) without
+    /// `apply_batch` in between.
+    #[test]
+    fn trim_committed_none_committed() {
+        // `shared.len() == 32`; both overlays have `len > 32`, so neither is committed.
+        let shared = Arc::new(SharedBitmap::<N>::new(make_bitmap(&[true; 32])));
+        let chain = make_chain(&shared, &[64, 96]);
+        let result = chain.trim_committed();
+        // Structure must be preserved in bottom-to-top order.
+        assert_eq!(chain_overlays(&result), vec![64, 96]);
+    }
+
+    /// Branch 4: exactly one layer is uncommitted (the newest), and everything below it has
+    /// been absorbed. This is the dominant pattern in chained growth: apply parent A, then
+    /// B held alive off A, then call `B.new_batch()` to build C. Roberto's fast path avoids
+    /// the general-case allocation dance and constructs the result directly as
+    /// `Layer(Base, overlay_B)` with a single `Arc::new`. The test also verifies the rebuilt
+    /// layer carries the cached `shared` reference correctly.
+    #[test]
+    fn trim_committed_exactly_one_uncommitted() {
+        // `shared.len() == 64`; committed layer (`overlay.len == 64`) + uncommitted (`96`).
+        let shared = Arc::new(SharedBitmap::<N>::new(make_bitmap(&[true; 64])));
+        let chain = make_chain(&shared, &[64, 96]);
+        let result = chain.trim_committed();
+        // The committed layer is gone; only the uncommitted overlay remains.
+        assert_eq!(chain_overlays(&result), vec![96]);
+        // And the rebuilt layer's `shared` field still points at the original terminal.
+        assert!(Arc::ptr_eq(result.shared(), &shared));
+    }
+
+    /// Branch 5: two or more uncommitted layers on top of a committed prefix. This is the
+    /// general-case path: the function has to walk the uncommitted tail, collect the overlays,
+    /// and rebuild a fresh chain rooted at `Base`. Exercises the loop that wires the cached
+    /// `shared` field into each reconstructed layer. Real-world trigger: build A, then B off A,
+    /// then C off B; apply only A; then call `C.new_batch()` — the trim sees B and C as
+    /// uncommitted, with A already absorbed.
+    #[test]
+    fn trim_committed_multiple_uncommitted() {
+        // `shared.len() == 64`; committed layer (64), then two uncommitted (96, 128).
+        let shared = Arc::new(SharedBitmap::<N>::new(make_bitmap(&[true; 64])));
+        let chain = make_chain(&shared, &[64, 96, 128]);
+        let result = chain.trim_committed();
+        // Committed layer dropped; uncommitted pair preserved in order.
+        assert_eq!(chain_overlays(&result), vec![96, 128]);
+        // Every reconstructed layer must still cache the original shared terminal.
+        assert!(Arc::ptr_eq(result.shared(), &shared));
+    }
 }

storage/src/qmdb/current/mod.rs

@@ -2969,6 +2969,65 @@ pub mod tests {
         });
     }
 
+    /// Build a child batch from a still-live parent whose apply was followed by a prune, then
+    /// merkleize and apply the child. The parent's `BitmapBatch` chain terminates in the shared
+    /// committed bitmap, and `prune` mutates that bitmap's pruning boundary in place. When the
+    /// child is constructed via `parent.new_batch()`, the internal `trim_committed` call must
+    /// observe the advanced boundary and produce a correct child chain; merkleize and apply must
+    /// then produce correct state for keys at and beyond the advanced floor.
+    #[test_traced("INFO")]
+    fn test_current_live_batch_child_after_prune() {
+        let executor = deterministic::Runner::default();
+        executor.start(|context| async move {
+            let ctx = context.with_label("db");
+            let mut db: UnorderedVariableDb = UnorderedVariableDb::init(
+                ctx.clone(),
+                variable_config::<OneCap>("child-after-prune", &ctx),
+            )
+            .await
+            .unwrap();
+
+            // Seed enough ops to span multiple bitmap chunks.
+            let mut seed = db.new_batch();
+            for i in 0u64..300 {
+                seed = seed.write(key(i), Some(val(i)));
+            }
+            let seed_m = seed.merkleize(&db, None).await.unwrap();
+            db.apply_batch(seed_m).await.unwrap();
+            db.commit().await.unwrap();
+
+            // Overwrite keys 0..250 so the inactivity floor advances past chunk 0.
+            let mut a_batch = db.new_batch();
+            for i in 0u64..250 {
+                a_batch = a_batch.write(key(i), Some(val(i + 10_000)));
+            }
+            let a = a_batch.merkleize(&db, None).await.unwrap();
+            db.apply_batch(Arc::clone(&a)).await.unwrap();
+            db.commit().await.unwrap();
+
+            // Prune while `a` is still live. Mutates the shared bitmap's pruning boundary in place.
+            let floor = db.inactivity_floor_loc();
+            db.prune(floor).await.unwrap();
+
+            // Extend `a` into `b` AFTER the prune. Building `b` off `a` triggers
+            // `trim_committed` on `a`'s chain, which must correctly see the advanced pruning
+            // boundary on the shared bitmap.
+            let b = a
+                .new_batch::<Sha256>()
+                .write(key(300), Some(val(300)))
+                .merkleize(&db, None)
+                .await
+                .unwrap();
+
+            db.apply_batch(b).await.unwrap();
+            assert_eq!(db.get(&key(0)).await.unwrap(), Some(val(10_000)));
+            assert_eq!(db.get(&key(249)).await.unwrap(), Some(val(10_249)));
+            assert_eq!(db.get(&key(300)).await.unwrap(), Some(val(300)));
+
+            db.destroy().await.unwrap();
+        });
+    }
+
     /// Regression: applying a batch after its ancestor Arc is dropped (without
     /// committing) must still apply the ancestor's bitmap pushes/clears and
     /// snapshot diffs.