Rkuris/range proofs (#345)

firewood/src/db.rs

@@ -305,13 +305,16 @@ impl<S: ShaleStore<Node> + Send + Sync> api::DbView for DbRev<S> {
             .map_err(|e| api::Error::IO(std::io::Error::new(ErrorKind::Other, e)))
     }
 
-    async fn range_proof<K: api::KeyType, V, N>(
+    async fn range_proof<K: api::KeyType, V>(
         &self,
-        _first_key: Option<K>,
-        _last_key: Option<K>,
-        _limit: usize,
-    ) -> Result<Option<api::RangeProof<K, V, N>>, api::Error> {
-        todo!()
+        first_key: Option<K>,
+        last_key: Option<K>,
+        limit: Option<usize>,
+    ) -> Result<Option<api::RangeProof<Vec<u8>, Vec<u8>>>, api::Error> {
+        self.merkle
+            .range_proof(self.header.kv_root, first_key, last_key, limit)
+            .await
+            .map_err(|e| api::Error::InternalError(Box::new(e)))
     }
 }
 
@@ -322,7 +325,7 @@ impl<S: ShaleStore<Node> + Send + Sync> DbRev<S> {
     ) -> Result<merkle::MerkleKeyValueStream<'_, S>, api::Error> {
         self.merkle
             .get_iter(start_key, self.header.kv_root)
-            .map_err(|e| api::Error::InternalError(e.into()))
+            .map_err(|e| api::Error::InternalError(Box::new(e)))
     }
 
     fn flush_dirty(&mut self) -> Option<()> {
@@ -454,7 +457,7 @@ impl Db {
         }
 
         block_in_place(|| Db::new_internal(db_path, cfg.clone()))
-            .map_err(|e| api::Error::InternalError(e.into()))
+            .map_err(|e| api::Error::InternalError(Box::new(e)))
     }
 
     /// Open a database.

firewood/src/db/proposal.rs

@@ -295,12 +295,12 @@ impl api::DbView for Proposal {
             .map_err(|e| api::Error::IO(std::io::Error::new(ErrorKind::Other, e)))
     }
 
-    async fn range_proof<K, V, N>(
+    async fn range_proof<K, V>(
         &self,
         _first_key: Option<K>,
         _last_key: Option<K>,
-        _limit: usize,
-    ) -> Result<Option<api::RangeProof<K, V, N>>, api::Error>
+        _limit: Option<usize>,
+    ) -> Result<Option<api::RangeProof<Vec<u8>, Vec<u8>>>, api::Error>
     where
         K: api::KeyType,
     {

firewood/src/merkle.rs

@@ -1,6 +1,5 @@
 // Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE.md for licensing terms.
-
 use crate::shale::{self, disk_address::DiskAddress, ObjWriteError, ShaleError, ShaleStore};
 use crate::v2::api;
 use crate::{nibbles::Nibbles, v2::api::Proof};
@@ -18,6 +17,7 @@ use thiserror::Error;
 
 mod node;
 mod partial_path;
+pub(super) mod range_proof;
 mod trie_hash;
 
 pub use node::{BranchNode, Data, ExtNode, LeafNode, Node, NodeType, NBRANCH};
@@ -1243,7 +1243,7 @@ impl<'a, S: shale::ShaleStore<node::Node> + Send + Sync> Stream for MerkleKeyVal
                 let root_node = self
                     .merkle
                     .get_node(self.merkle_root)
-                    .map_err(|e| api::Error::InternalError(e.into()))?;
+                    .map_err(|e| api::Error::InternalError(Box::new(e)))?;
                 let mut last_node = root_node;
                 let mut parents = vec![];
                 let leaf = loop {
@@ -1272,7 +1272,7 @@ impl<'a, S: shale::ShaleStore<node::Node> + Send + Sync> Stream for MerkleKeyVal
                             let next = self
                                 .merkle
                                 .get_node(leftmost_address)
-                                .map_err(|e| api::Error::InternalError(e.into()))?;
+                                .map_err(|e| api::Error::InternalError(Box::new(e)))?;
 
                             parents.push((last_node, leftmost_position as u8));
 
@@ -1295,12 +1295,12 @@ impl<'a, S: shale::ShaleStore<node::Node> + Send + Sync> Stream for MerkleKeyVal
                 let root_node = self
                     .merkle
                     .get_node(self.merkle_root)
-                    .map_err(|e| api::Error::InternalError(e.into()))?;
+                    .map_err(|e| api::Error::InternalError(Box::new(e)))?;
 
                 let (found_node, parents) = self
                     .merkle
                     .get_node_and_parents_by_key(root_node, &key)
-                    .map_err(|e| api::Error::InternalError(e.into()))?;
+                    .map_err(|e| api::Error::InternalError(Box::new(e)))?;
 
                 let Some(last_node) = found_node else {
                     return Poll::Ready(None);
@@ -1343,7 +1343,7 @@ impl<'a, S: shale::ShaleStore<node::Node> + Send + Sync> Stream for MerkleKeyVal
                         let current_node = self
                             .merkle
                             .get_node(child_address)
-                            .map_err(|e| api::Error::InternalError(e.into()))?;
+                            .map_err(|e| api::Error::InternalError(Box::new(e)))?;
 
                         let found_key = key_from_parents(&parents);
 
@@ -1389,7 +1389,7 @@ impl<'a, S: shale::ShaleStore<node::Node> + Send + Sync> Stream for MerkleKeyVal
                                         .merkle
                                         .get_node(found_address)
                                         .map(|node| (node, None))
-                                        .map_err(|e| api::Error::InternalError(e.into()))?;
+                                        .map_err(|e| api::Error::InternalError(Box::new(e)))?;
 
                                     // stop_descending if:
                                     //  - on a branch and it has a value; OR
@@ -1569,9 +1569,9 @@ mod tests {
     use std::sync::Arc;
     use test_case::test_case;
 
-    #[test_case(vec![0x12, 0x34, 0x56], vec![0x1, 0x2, 0x3, 0x4, 0x5, 0x6])]
-    #[test_case(vec![0xc0, 0xff], vec![0xc, 0x0, 0xf, 0xf])]
-    fn to_nibbles(bytes: Vec<u8>, nibbles: Vec<u8>) {
+    #[test_case(vec![0x12, 0x34, 0x56], &[0x1, 0x2, 0x3, 0x4, 0x5, 0x6])]
+    #[test_case(vec![0xc0, 0xff], &[0xc, 0x0, 0xf, 0xf])]
+    fn to_nibbles(bytes: Vec<u8>, nibbles: &[u8]) {
         let n: Vec<_> = bytes.into_iter().flat_map(to_nibble_array).collect();
         assert_eq!(n, nibbles);
     }
@@ -1688,7 +1688,7 @@ mod tests {
         let root = merkle.init_root().unwrap();
         let mut it = merkle.get_iter(Some(b"x"), root).unwrap();
         let next = it.next().await;
-        assert!(next.is_none())
+        assert!(next.is_none());
     }
 
     #[test_case(Some(&[u8::MIN]); "Starting at first key")]
@@ -1747,7 +1747,7 @@ mod tests {
 
     #[test]
     fn remove_many() {
-        let mut merkle = create_test_merkle();
+        let mut merkle: Merkle<CompactSpace<Node, DynamicMem>> = create_test_merkle();
         let root = merkle.init_root().unwrap();
 
         // insert values
@@ -1775,4 +1775,66 @@ mod tests {
             assert!(fetched_val.is_none());
         }
     }
+
+    #[tokio::test]
+    async fn empty_range_proof() {
+        let merkle: Merkle<CompactSpace<Node, DynamicMem>> = create_test_merkle();
+        let root = merkle.init_root().unwrap();
+
+        assert!(merkle
+            .range_proof::<&[u8]>(root, None, None, None)
+            .await
+            .unwrap()
+            .is_none());
+    }
+
+    #[tokio::test]
+    async fn full_range_proof() {
+        let mut merkle: Merkle<CompactSpace<Node, DynamicMem>> = create_test_merkle();
+        let root = merkle.init_root().unwrap();
+        // insert values
+        for key_val in u8::MIN..=u8::MAX {
+            let key = &[key_val];
+            let val = &[key_val];
+
+            merkle.insert(key, val.to_vec(), root).unwrap();
+        }
+        merkle.flush_dirty();
+
+        let rangeproof = merkle
+            .range_proof::<&[u8]>(root, None, None, None)
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(rangeproof.middle.len(), (u8::MAX - 1).into());
+        assert_ne!(rangeproof.first_key_proof.0, rangeproof.last_key_proof.0);
+        let left_proof = merkle.prove([u8::MIN], root).unwrap();
+        let right_proof = merkle.prove([u8::MAX], root).unwrap();
+        assert_eq!(rangeproof.first_key_proof.0, left_proof.0);
+        assert_eq!(rangeproof.last_key_proof.0, right_proof.0);
+    }
+
+    #[tokio::test]
+    async fn single_value_range_proof() {
+        const RANDOM_KEY: u8 = 42;
+
+        let mut merkle: Merkle<CompactSpace<Node, DynamicMem>> = create_test_merkle();
+        let root = merkle.init_root().unwrap();
+        // insert values
+        for key_val in u8::MIN..=u8::MAX {
+            let key = &[key_val];
+            let val = &[key_val];
+
+            merkle.insert(key, val.to_vec(), root).unwrap();
+        }
+        merkle.flush_dirty();
+
+        let rangeproof = merkle
+            .range_proof(root, Some([RANDOM_KEY]), None, Some(1))
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(rangeproof.first_key_proof.0, rangeproof.last_key_proof.0);
+        assert_eq!(rangeproof.middle.len(), 0);
+    }
 }

firewood/src/merkle/range_proof.rs

@@ -0,0 +1,88 @@
+// Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE.md for licensing terms.
+
+use std::future::ready;
+
+use futures::{StreamExt as _, TryStreamExt as _};
+
+use crate::{
+    shale::{disk_address::DiskAddress, ShaleStore},
+    v2::api,
+};
+
+use super::{Merkle, Node};
+
+impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
+    pub(crate) async fn range_proof<K: api::KeyType + Send + Sync>(
+        &self,
+        root: DiskAddress,
+        first_key: Option<K>,
+        last_key: Option<K>,
+        limit: Option<usize>,
+    ) -> Result<Option<api::RangeProof<Vec<u8>, Vec<u8>>>, api::Error> {
+        // limit of 0 is always an empty RangeProof
+        if let Some(0) = limit {
+            return Ok(None);
+        }
+
+        let mut stream = self
+            .get_iter(first_key, root)
+            .map_err(|e| api::Error::InternalError(Box::new(e)))?;
+
+        // fetch the first key from the stream
+        let first_result = stream.next().await;
+
+        // transpose the Option<Result<T, E>> to Result<Option<T>, E>
+        // If this is an error, the ? operator will return it
+        let Some((key, _)) = first_result.transpose()? else {
+            // nothing returned, either the trie is empty or the key wasn't found
+            return Ok(None);
+        };
+
+        let first_key_proof = self
+            .prove(key, root)
+            .map_err(|e| api::Error::InternalError(Box::new(e)))?;
+        let limit = limit.map(|old_limit| old_limit - 1);
+
+        // we stop streaming if either we hit the limit or the key returned was larger
+        // than the largest key requested
+        let mut middle = stream
+            .take(limit.unwrap_or(usize::MAX))
+            .take_while(|kv_result| {
+                // no last key asked for, so keep going
+                let Some(last_key) = last_key.as_ref() else {
+                    return ready(true);
+                };
+
+                // return the error if there was one
+                let Ok(kv) = kv_result else {
+                    return ready(true);
+                };
+
+                // keep going if the key returned is less than the last key requested
+                ready(kv.0.as_slice() <= last_key.as_ref())
+            })
+            .try_collect::<Vec<(Vec<u8>, Vec<u8>)>>()
+            .await?;
+
+        // remove the last key from middle and do a proof on it
+        let last_key_proof = match middle.pop() {
+            None => {
+                return Ok(Some(api::RangeProof {
+                    first_key_proof: first_key_proof.clone(),
+                    middle: vec![],
+                    last_key_proof: first_key_proof,
+                }))
+            }
+            Some((last_key, _)) => self
+                .prove(last_key, root)
+                .map_err(|e| api::Error::InternalError(Box::new(e)))?,
+        };
+
+        Ok(Some(api::RangeProof {
+            first_key_proof,
+            middle,
+            last_key_proof,
+        }))
+    }
+}

firewood/src/v2/api.rs

@@ -71,22 +71,25 @@ pub enum Error {
     InvalidProposal,
 
     #[error("Internal error")]
-    InternalError(Box<dyn std::error::Error>),
+    InternalError(Box<dyn std::error::Error + Send>),
+
+    #[error("Range too small")]
+    RangeTooSmall,
 }
 
 /// A range proof, consisting of a proof of the first key and the last key,
 /// and a vector of all key/value pairs
 #[derive(Debug)]
-pub struct RangeProof<K, V, N> {
-    pub first_key: Proof<N>,
-    pub last_key: Proof<N>,
+pub struct RangeProof<K, V> {
+    pub first_key_proof: Proof<Vec<u8>>,
+    pub last_key_proof: Proof<Vec<u8>>,
     pub middle: Vec<(K, V)>,
 }
 
 /// A proof that a single key is present
 ///
 /// The generic N represents the storage for the node data
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct Proof<N>(pub HashMap<HashKey, N>);
 
 /// The database interface, which includes a type for a static view of
@@ -152,12 +155,12 @@ pub trait DbView {
     /// * `last_key` - If None, continue to the end of the database
     /// * `limit` - The maximum number of keys in the range proof
     ///
-    async fn range_proof<K: KeyType, V, N>(
+    async fn range_proof<K: KeyType, V: Send + Sync>(
         &self,
         first_key: Option<K>,
         last_key: Option<K>,
-        limit: usize,
-    ) -> Result<Option<RangeProof<K, V, N>>, Error>;
+        limit: Option<usize>,
+    ) -> Result<Option<RangeProof<Vec<u8>, Vec<u8>>>, Error>;
 }
 
 /// A proposal for a new revision of the database.

firewood/src/v2/db.rs

@@ -32,11 +32,13 @@ pub struct Db<T> {
 impl From<DbError> for api::Error {
     fn from(value: DbError) -> Self {
         match value {
-            DbError::InvalidParams => api::Error::InternalError(value.into()),
-            DbError::Merkle(e) => api::Error::InternalError(e.into()),
+            DbError::InvalidParams => api::Error::InternalError(Box::new(value)),
+            DbError::Merkle(e) => api::Error::InternalError(Box::new(e)),
             DbError::System(e) => api::Error::IO(e.into()),
-            DbError::KeyNotFound | DbError::CreateError => api::Error::InternalError(value.into()),
-            DbError::Shale(e) => api::Error::InternalError(e.into()),
+            DbError::KeyNotFound | DbError::CreateError => {
+                api::Error::InternalError(Box::new(value))
+            }
+            DbError::Shale(e) => api::Error::InternalError(Box::new(e)),
             DbError::IO(e) => api::Error::IO(e),
             DbError::InvalidProposal => api::Error::InvalidProposal,
         }

firewood/src/v2/emptydb.rs

@@ -66,12 +66,12 @@ impl DbView for HistoricalImpl {
         Ok(None)
     }
 
-    async fn range_proof<K: KeyType, V, N>(
+    async fn range_proof<K: KeyType, V>(
         &self,
         _first_key: Option<K>,
         _last_key: Option<K>,
-        _limit: usize,
-    ) -> Result<Option<RangeProof<K, V, N>>, Error> {
+        _limit: Option<usize>,
+    ) -> Result<Option<RangeProof<Vec<u8>, Vec<u8>>>, Error> {
         Ok(None)
     }
 }