Move range-proof-builder (#361)

firewood/src/merkle.rs

@@ -3,8 +3,9 @@
 use crate::shale::{self, disk_address::DiskAddress, ObjWriteError, ShaleError, ShaleStore};
 use crate::v2::api;
 use crate::{nibbles::Nibbles, v2::api::Proof};
-use futures::Stream;
+use futures::{Stream, StreamExt, TryStreamExt};
 use sha3::Digest;
+use std::future::ready;
 use std::{
     cmp::Ordering,
     collections::HashMap,
@@ -17,7 +18,6 @@ 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};
@@ -1190,6 +1190,79 @@ impl<S: ShaleStore<Node> + Send + Sync> Merkle<S> {
             merkle: self,
         })
     }
+
+    pub(super) 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,
+        }))
+    }
 }
 
 enum IteratorState<'a> {