support split bagging in qmdb-current

Author: roberto-bayardo

storage/src/merkle/mod.rs

@@ -33,6 +33,8 @@ pub use position::Position;
 #[cfg(test)]
 pub(crate) use proof::build_range_proof;
 pub use proof::Proof;
+#[cfg(feature = "std")]
+pub(crate) use proof::{build_range_collection_proof, range_collection_nodes};
 pub use read::Readable;
 use thiserror::Error;
 

storage/src/merkle/proof.rs

@@ -738,6 +738,70 @@ impl<F: Family, D: Digest> Proof<F, D> {
             )
             .ok_or(ReconstructionError::InvalidProof)
     }
+
+    /// Authenticate the proven range without reconstructing the full generic Merkle root.
+    ///
+    /// This consumes only the sibling digests needed to rebuild the proven range peaks, then returns
+    /// those peak digests in `collected`. It deliberately does not consume peak-bagging witnesses
+    /// such as prefix peaks, after peaks, or backward-fold suffix accumulators.
+    ///
+    /// Current QMDB grafted proofs use this path because their final root is rebuilt by the wrapper
+    /// from the collected range peaks plus grafted prefix/suffix witnesses. Including generic
+    /// peak-bagging witnesses here would create proof bytes that the wrapper root ignores, making
+    /// those bytes malleable.
+    pub(crate) fn reconstruct_range_collecting<H, E>(
+        &self,
+        hasher: &H,
+        elements: &[E],
+        start_loc: Location<F>,
+        collected: &mut Vec<(Position<F>, D)>,
+    ) -> Result<(), ReconstructionError>
+    where
+        H: Hasher<F, Digest = D>,
+        E: AsRef<[u8]>,
+    {
+        if elements.is_empty() {
+            return Err(ReconstructionError::MissingElements);
+        }
+        if !start_loc.is_valid_index() {
+            return Err(ReconstructionError::InvalidStartLoc);
+        }
+        let end_loc = start_loc
+            .checked_add(elements.len() as u64)
+            .ok_or(ReconstructionError::InvalidEndLoc)?;
+        if end_loc > self.leaves {
+            return Err(ReconstructionError::InvalidEndLoc);
+        }
+
+        let range = start_loc..end_loc;
+        let bp = Blueprint::new_using_policy(
+            self.leaves,
+            self.inactive_peaks,
+            Bagging::ForwardFold,
+            range,
+        )
+        .map_err(|_| ReconstructionError::InvalidSize)?;
+
+        let mut sibling_cursor = 0usize;
+        let mut elements_iter = elements.iter();
+        for peak in &bp.range_peaks {
+            let peak_digest = peak.reconstruct_digest(
+                hasher,
+                &bp.range,
+                &mut elements_iter,
+                &self.digests,
+                &mut sibling_cursor,
+                Some(collected),
+            )?;
+            collected.push((peak.pos, peak_digest));
+        }
+
+        if elements_iter.next().is_some() || sibling_cursor != self.digests.len() {
+            return Err(ReconstructionError::ExtraDigests);
+        }
+
+        Ok(())
+    }
 }
 
 /// A perfect binary subtree within a peak, identified by its root position, height,
@@ -1222,6 +1286,52 @@ where
     )
 }
 
+/// Return the node positions needed by [`build_range_collection_proof`].
+#[cfg(feature = "std")]
+pub(crate) fn range_collection_nodes<F: Family>(
+    leaves: Location<F>,
+    inactive_peaks: usize,
+    range: Range<Location<F>>,
+) -> Result<Vec<Position<F>>, super::Error<F>> {
+    let bp = Blueprint::new_using_policy(leaves, inactive_peaks, Bagging::ForwardFold, range)?;
+    let mut fetch_nodes = Vec::new();
+    for peak in &bp.range_peaks {
+        peak.collect_siblings(&bp.range, &mut fetch_nodes);
+    }
+    Ok(fetch_nodes)
+}
+
+/// Build a proof containing only the sibling digests needed to authenticate the requested range.
+///
+/// The resulting proof cannot reconstruct a complete generic Merkle root on its own. It is intended
+/// for wrappers that rebuild a custom root from separately supplied peak witnesses, such as current
+/// QMDB grafted proofs.
+#[cfg(feature = "std")]
+pub(crate) fn build_range_collection_proof<F, D, E>(
+    leaves: Location<F>,
+    inactive_peaks: usize,
+    range: Range<Location<F>>,
+    get_node: impl Fn(Position<F>) -> Option<D>,
+    element_pruned: impl Fn(Position<F>) -> E,
+) -> Result<Proof<F, D>, E>
+where
+    F: Family,
+    D: Digest,
+    E: From<super::Error<F>>,
+{
+    let fetch_nodes = range_collection_nodes(leaves, inactive_peaks, range)?;
+    let mut digests = Vec::with_capacity(fetch_nodes.len());
+    for pos in fetch_nodes {
+        digests.push(get_node(pos).ok_or_else(|| element_pruned(pos))?);
+    }
+
+    Ok(Proof {
+        leaves,
+        inactive_peaks,
+        digests,
+    })
+}
+
 /// Returns the positions of the minimal set of nodes whose digests are required to prove the
 /// inclusion of the elements at the specified `locations`, using the provided root bagging.
 #[cfg(any(feature = "std", test))]

storage/src/qmdb/current/batch.rs

@@ -10,6 +10,7 @@ use crate::{
         mem::Mem, storage::Storage as MerkleStorage, Graftable, Location, Position, Readable,
     },
     qmdb::{
+        self,
         any::{
             self,
             batch::{lookup_sorted, DiffEntry, FloorScan},
@@ -176,7 +177,7 @@ struct BatchStorageAdapter<
 
 impl<
         'a,
-        F: Graftable,
+        F: Graftable + qmdb::RootSpec,
         D: Digest,
         R: Readable<Family = F, Digest = D, Error = merkle::Error<F>>,
         S: MerkleStorage<F, Digest = D>,
@@ -192,7 +193,7 @@ impl<
 }
 
 impl<
-        F: Graftable,
+        F: Graftable + qmdb::RootSpec,
         D: Digest,
         R: Readable<Family = F, Digest = D, Error = merkle::Error<F>>,
         S: MerkleStorage<F, Digest = D>,
@@ -349,7 +350,7 @@ where
 // Unordered get + merkleize.
 impl<F, K, V, H, const N: usize> UnmerkleizedBatch<F, H, update::Unordered<K, V>, N>
 where
-    F: Graftable,
+    F: Graftable + qmdb::RootSpec,
     K: Key,
     V: ValueEncoding,
     H: Hasher,
@@ -412,7 +413,7 @@ where
 // Ordered get + merkleize.
 impl<F, K, V, H, const N: usize> UnmerkleizedBatch<F, H, update::Ordered<K, V>, N>
 where
-    F: Graftable,
+    F: Graftable + qmdb::RootSpec,
     K: Key,
     V: ValueEncoding,
     H: Hasher,
@@ -550,7 +551,7 @@ async fn compute_current_layer<F, E, U, C, I, H, const N: usize>(
     bitmap_parent: &BitmapBatch<N>,
 ) -> Result<Arc<MerkleizedBatch<F, H::Digest, U, N>>, Error<F>>
 where
-    F: Graftable,
+    F: Graftable + qmdb::RootSpec,
     E: Context,
     U: update::Update + Send + Sync,
     C: Contiguous<Item = Operation<F, U>>,
@@ -674,6 +675,7 @@ where
         &bitmap_batch,
         &grafted_storage,
         partial,
+        inner.new_inactivity_floor_loc,
         &ops_root,
     )
     .await?;
@@ -956,7 +958,7 @@ where
 
 impl<F, E, C, I, H, U, const N: usize> super::db::Db<F, E, C, I, H, U, N>
 where
-    F: Graftable,
+    F: Graftable + qmdb::RootSpec,
     E: Context,
     U: update::Update + Send + Sync,
     C: Contiguous<Item = Operation<F, U>>,
@@ -1000,7 +1002,7 @@ mod trait_impls {
         UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N>>
         for UnmerkleizedBatch<F, H, update::Unordered<K, V>, N>
     where
-        F: Graftable,
+        F: Graftable + qmdb::RootSpec,
         K: Key,
         V: ValueEncoding + 'static,
         H: Hasher,
@@ -1020,20 +1022,29 @@ mod trait_impls {
             Self::write(self, key, value)
         }
 
-        fn merkleize(
+        async fn merkleize(
             self,
             db: &CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N>,
             metadata: Option<V::Value>,
-        ) -> impl Future<Output = Result<Self::Merkleized, crate::qmdb::Error<F>>> {
-            self.merkleize(db, metadata)
+        ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
+            let Self {
+                inner,
+                grafted_parent,
+                bitmap_parent,
+            } = self;
+            let scan = BitmapScan::new(&bitmap_parent);
+            let inner = inner
+                .merkleize_with_floor_scan(&db.any, metadata, scan)
+                .await?;
+            compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
         }
     }
 
     impl<F, K, V, H, E, C, I, const N: usize>
         UnmerkleizedBatchTrait<CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N>>
         for UnmerkleizedBatch<F, H, update::Ordered<K, V>, N>
     where
-        F: Graftable,
+        F: Graftable + qmdb::RootSpec,
         K: Key,
         V: ValueEncoding + 'static,
         H: Hasher,
@@ -1053,12 +1064,21 @@ mod trait_impls {
             Self::write(self, key, value)
         }
 
-        fn merkleize(
+        async fn merkleize(
             self,
             db: &CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N>,
             metadata: Option<V::Value>,
-        ) -> impl Future<Output = Result<Self::Merkleized, crate::qmdb::Error<F>>> {
-            self.merkleize(db, metadata)
+        ) -> Result<Self::Merkleized, crate::qmdb::Error<F>> {
+            let Self {
+                inner,
+                grafted_parent,
+                bitmap_parent,
+            } = self;
+            let scan = BitmapScan::new(&bitmap_parent);
+            let inner = inner
+                .merkleize_with_floor_scan(&db.any, metadata, scan)
+                .await?;
+            compute_current_layer(inner, db, &grafted_parent, &bitmap_parent).await
         }
     }
 
@@ -1077,7 +1097,7 @@ mod trait_impls {
     impl<F, E, K, V, C, I, H, const N: usize> BatchableDb
         for CurrentDb<F, E, C, I, H, update::Unordered<K, V>, N>
     where
-        F: Graftable,
+        F: Graftable + qmdb::RootSpec,
         E: Context,
         K: Key,
         V: ValueEncoding + 'static,
@@ -1094,22 +1114,22 @@ mod trait_impls {
         type Batch = UnmerkleizedBatch<F, H, update::Unordered<K, V>, N>;
 
         fn new_batch(&self) -> Self::Batch {
-            self.new_batch()
+            Self::new_batch(self)
         }
 
         fn apply_batch(
             &mut self,
             batch: Self::Merkleized,
         ) -> impl Future<Output = Result<core::ops::Range<Location<F>>, crate::qmdb::Error<F>>>
         {
-            self.apply_batch(batch)
+            Self::apply_batch(self, batch)
         }
     }
 
     impl<F, E, K, V, C, I, H, const N: usize> BatchableDb
         for CurrentDb<F, E, C, I, H, update::Ordered<K, V>, N>
     where
-        F: Graftable,
+        F: Graftable + qmdb::RootSpec,
         E: Context,
         K: Key,
         V: ValueEncoding + 'static,
@@ -1126,15 +1146,15 @@ mod trait_impls {
         type Batch = UnmerkleizedBatch<F, H, update::Ordered<K, V>, N>;
 
         fn new_batch(&self) -> Self::Batch {
-            self.new_batch()
+            Self::new_batch(self)
         }
 
         fn apply_batch(
             &mut self,
             batch: Self::Merkleized,
         ) -> impl Future<Output = Result<core::ops::Range<Location<F>>, crate::qmdb::Error<F>>>
         {
-            self.apply_batch(batch)
+            Self::apply_batch(self, batch)
         }
     }
 }