Use v2::api::Proof<V> over proof::Proof (#215)

firewood/examples/rev.rs

@@ -6,8 +6,8 @@ use std::{collections::VecDeque, path::Path};
 use firewood::{
     db::{BatchOp, Db, DbConfig, Proposal, Revision, WalConfig},
     merkle::{Node, TrieHash},
-    proof::Proof,
     storage::StoreRevShared,
+    v2::api::Proof,
 };
 use shale::compact::CompactSpace;
 
@@ -166,7 +166,7 @@ impl RevisionTracker {
     }
 }
 
-fn build_proof(revision: &Revision<SharedStore>, items: &[(&str, &str)]) -> Proof {
+fn build_proof(revision: &Revision<SharedStore>, items: &[(&str, &str)]) -> Proof<Vec<u8>> {
     let mut proof = revision.prove(items[0].0).unwrap();
     let end = revision.prove(items.last().unwrap().0).unwrap();
     proof.concat_proofs(end);

firewood/src/db.rs

@@ -8,12 +8,13 @@ pub use crate::{
 use crate::{
     file,
     merkle::{Merkle, MerkleError, Node, TrieHash, TRIE_HASH_LEN},
-    proof::{Proof, ProofError},
+    proof::ProofError,
     storage::{
         buffer::{BufferWrite, DiskBuffer, DiskBufferRequester},
         AshRecord, CachedSpace, MemStoreR, SpaceWrite, StoreConfig, StoreDelta, StoreRevMut,
         StoreRevShared, ZeroStore, PAGE_SIZE_NBIT,
     },
+    v2::api::Proof,
 };
 use bytemuck::{cast_slice, AnyBitPattern};
 use metered::{metered, HitCount};
@@ -313,15 +314,14 @@ impl<S: ShaleStore<Node> + Send + Sync> DbRev<S> {
             .map_err(DbError::Merkle)
     }
 
-    /// Provides a proof that a key is in the Trie.
-    pub fn prove<K: AsRef<[u8]>>(&self, key: K) -> Result<Proof, MerkleError> {
-        self.merkle.prove(key, self.header.kv_root)
+    pub fn prove<K: AsRef<[u8]>>(&self, key: K) -> Result<Proof<Vec<u8>>, MerkleError> {
+        self.merkle.prove::<K>(key, self.header.kv_root)
     }
 
     /// Verifies a range proof is valid for a set of keys.
-    pub fn verify_range_proof<K: AsRef<[u8]>, V: AsRef<[u8]>>(
+    pub fn verify_range_proof<N: AsRef<[u8]> + Send, K: AsRef<[u8]>, V: AsRef<[u8]>>(
         &self,
-        proof: Proof,
+        proof: Proof<N>,
         first_key: K,
         last_key: K,
         keys: Vec<K>,

firewood/src/merkle.rs

@@ -1,7 +1,7 @@
 // Copyright (C) 2023, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE.md for licensing terms.
 
-use crate::{nibbles::Nibbles, proof::Proof};
+use crate::{nibbles::Nibbles, v2::api::Proof};
 use sha3::Digest;
 use shale::{disk_address::DiskAddress, ObjRef, ShaleError, ShaleStore};
 use std::{
@@ -1031,13 +1031,13 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
     /// If the trie does not contain a value for key, the returned proof contains
     /// all nodes of the longest existing prefix of the key, ending with the node
     /// that proves the absence of the key (at least the root node).
-    pub fn prove<K: AsRef<[u8]>>(&self, key: K, root: DiskAddress) -> Result<Proof, MerkleError>
+    pub fn prove<K>(&self, key: K, root: DiskAddress) -> Result<Proof<Vec<u8>>, MerkleError>
     where
         K: AsRef<[u8]>,
     {
         let key_nibbles = Nibbles::<0>::new(key.as_ref());
 
-        let mut proofs: HashMap<[u8; TRIE_HASH_LEN], Vec<u8>> = HashMap::new();
+        let mut proofs = HashMap::new();
         if root.is_null() {
             return Ok(Proof(proofs));
         }

firewood/src/merkle_util.rs

@@ -3,7 +3,8 @@
 
 use crate::{
     merkle::{Merkle, Node, Ref, RefMut, TrieHash},
-    proof::{Proof, ProofError},
+    proof::ProofError,
+    v2::api::Proof,
 };
 use shale::{
     cached::DynamicMem, compact::CompactSpace, disk_address::DiskAddress, CachedStore, ShaleStore,
@@ -86,26 +87,26 @@ impl<S: ShaleStore<Node> + Send + Sync> MerkleSetup<S> {
         String::from_utf8(s).map_err(|_err| DataStoreError::UTF8Error)
     }
 
-    pub fn prove<K: AsRef<[u8]>>(&self, key: K) -> Result<Proof, DataStoreError> {
+    pub fn prove<K: AsRef<[u8]>>(&self, key: K) -> Result<Proof<Vec<u8>>, DataStoreError> {
         self.merkle
             .prove(key, self.root)
             .map_err(|_err| DataStoreError::ProofError)
     }
 
-    pub fn verify_proof<K: AsRef<[u8]>>(
+    pub fn verify_proof<N: AsRef<[u8]> + Send, K: AsRef<[u8]>>(
         &self,
         key: K,
-        proof: &Proof,
+        proof: &Proof<N>,
     ) -> Result<Option<Vec<u8>>, DataStoreError> {
         let hash: [u8; 32] = *self.root_hash()?;
         proof
             .verify_proof(key, hash)
             .map_err(|_err| DataStoreError::ProofVerificationError)
     }
 
-    pub fn verify_range_proof<K: AsRef<[u8]>, V: AsRef<[u8]>>(
+    pub fn verify_range_proof<N: AsRef<[u8]> + Send, K: AsRef<[u8]>, V: AsRef<[u8]>>(
         &self,
-        proof: &Proof,
+        proof: &Proof<N>,
         first_key: K,
         last_key: K,
         keys: Vec<K>,

firewood/src/proof.rs

@@ -10,14 +10,17 @@ use crate::{
     merkle_util::{new_merkle, DataStoreError, MerkleSetup},
 };
 use nix::errno::Errno;
-use serde::{Deserialize, Serialize};
 use sha3::Digest;
-use shale::{disk_address::DiskAddress, ShaleError, ShaleStore};
-use std::{cmp::Ordering, collections::HashMap, error::Error, fmt, ops::Deref};
+use shale::disk_address::DiskAddress;
+use shale::ShaleError;
+use shale::ShaleStore;
 
-/// Hash -> RLP encoding map
-#[derive(Debug, Serialize, Deserialize)]
-pub struct Proof(pub HashMap<[u8; 32], Vec<u8>>);
+use std::cmp::Ordering;
+use std::error::Error;
+use std::fmt;
+use std::ops::Deref;
+
+use crate::v2::api::Proof;
 
 #[derive(Debug)]
 pub enum ProofError {
@@ -116,10 +119,12 @@ pub struct SubProof {
     hash: Option<[u8; 32]>,
 }
 
-impl Proof {
+impl<N: AsRef<[u8]> + Send> Proof<N> {
     /// verify_proof checks merkle proofs. The given proof must contain the value for
     /// key in a trie with the given root hash. VerifyProof returns an error if the
     /// proof contains invalid trie nodes or the wrong value.
+    ///
+    /// The generic N represents the storage for the node data
     pub fn verify_proof<K: AsRef<[u8]>>(
         &self,
         key: K,
@@ -137,7 +142,8 @@ impl Proof {
             let cur_proof = proofs_map
                 .get(&cur_hash)
                 .ok_or(ProofError::ProofNodeMissing)?;
-            let (sub_proof, size) = self.locate_subproof(remaining_key_nibbles, cur_proof)?;
+            let (sub_proof, size) =
+                self.locate_subproof(remaining_key_nibbles, cur_proof.as_ref())?;
             index += size;
 
             match sub_proof {
@@ -242,7 +248,7 @@ impl Proof {
         }
     }
 
-    pub fn concat_proofs(&mut self, other: Proof) {
+    pub fn concat_proofs(&mut self, other: Proof<N>) {
         self.0.extend(other.0)
     }
 
@@ -290,7 +296,7 @@ impl Proof {
         // ensure there are no more accounts / slots in the trie.
         if keys.is_empty() {
             let proof_to_path =
-                self.proof_to_path(first_key.as_ref(), root_hash, &mut merkle_setup, true)?;
+                self.proof_to_path(first_key, root_hash, &mut merkle_setup, true)?;
             return match proof_to_path {
                 Some(_) => Err(ProofError::InvalidData),
                 None => Ok(false),
@@ -365,9 +371,9 @@ impl Proof {
     /// necessary nodes will be resolved and leave the remaining as hashnode.
     ///
     /// The given edge proof is allowed to be an existent or non-existent proof.
-    fn proof_to_path<K: AsRef<[u8]>, S: ShaleStore<Node> + Send + Sync>(
+    fn proof_to_path<KV: AsRef<[u8]>, S: ShaleStore<Node> + Send + Sync>(
         &self,
-        key: K,
+        key: KV,
         root_hash: [u8; 32],
         merkle_setup: &mut MerkleSetup<S>,
         allow_non_existent_node: bool,
@@ -392,7 +398,7 @@ impl Proof {
                 .ok_or(ProofError::ProofNodeMissing)?;
             // TODO(Hao): (Optimization) If a node is alreay decode we don't need to decode again.
             let (mut chd_ptr, sub_proof, size) =
-                self.decode_node(merkle, cur_key, cur_proof, false)?;
+                self.decode_node(merkle, cur_key, cur_proof.as_ref(), false)?;
 
             // Link the child to the parent based on the node type.
             match &u_ref.inner() {
@@ -464,7 +470,7 @@ impl Proof {
                             let proof =
                                 proofs_map.get(p_hash).ok_or(ProofError::ProofNodeMissing)?;
 
-                            chd_ptr = self.decode_node(merkle, cur_key, proof, true)?.0;
+                            chd_ptr = self.decode_node(merkle, cur_key, proof.as_ref(), true)?.0;
 
                             // Link the child to the parent based on the node type.
                             match &u_ref.inner() {

firewood/src/v2/api.rs

@@ -76,8 +76,10 @@ pub struct RangeProof<K: KeyType, V: ValueType> {
 }
 
 /// A proof that a single key is present
+///
+/// The generic N represents the storage for the node data
 #[derive(Debug)]
-pub struct Proof<V>(pub HashMap<HashKey, V>);
+pub struct Proof<N: Send>(pub HashMap<HashKey, N>);
 
 /// The database interface, which includes a type for a static view of
 /// the database (the DbView). The most common implementation of the DbView

firewood/tests/db.rs

@@ -209,7 +209,7 @@ fn create_db_issue_proof() {
     let key = "doe".as_bytes();
     let root_hash = rev.kv_root_hash();
 
-    match rev.prove(key) {
+    match rev.prove::<&[u8]>(key) {
         Ok(proof) => {
             let verification = proof.verify_proof(key, *root_hash.unwrap()).unwrap();
             assert!(verification.is_some());

firewood/tests/merkle.rs

@@ -1,7 +1,8 @@
 use firewood::{
     merkle::Node,
     merkle_util::{new_merkle, DataStoreError, MerkleSetup},
-    proof::{Proof, ProofError},
+    proof::ProofError,
+    v2::api::Proof,
 };
 use rand::Rng;
 use shale::{cached::DynamicMem, compact::CompactSpace};
@@ -679,7 +680,7 @@ fn test_all_elements_proof() -> Result<(), ProofError> {
     let keys: Vec<&[u8; 32]> = item_iter.clone().map(|item| item.0).collect();
     let vals: Vec<&[u8; 20]> = item_iter.map(|item| item.1).collect();
 
-    let empty_proof = Proof(HashMap::new());
+    let empty_proof = Proof(HashMap::<[u8; 32], Vec<u8>>::new());
     let empty_key: [u8; 32] = [0; 32];
     merkle.verify_range_proof(
         &empty_proof,