From f2ded024552e6c11ac5ad35cdd80dffe76a28a57 Mon Sep 17 00:00:00 2001
From: "Dr. Chat" <arkolbed@gmail.com>
Date: Wed, 1 Jan 2025 13:14:01 -0600
Subject: [PATCH] Outline `algos::compute_depth`

---
 atrium-repo/src/mst.rs | 49 ++++++++++++++++++++++++++----------------
 1 file changed, 30 insertions(+), 19 deletions(-)

diff --git a/atrium-repo/src/mst.rs b/atrium-repo/src/mst.rs
index 76c6b005..d1c70658 100644
--- a/atrium-repo/src/mst.rs
+++ b/atrium-repo/src/mst.rs
@@ -69,6 +69,31 @@ mod algos {
         Stop(R),
     }
 
+    /// Compute the depth of the specified node.
+    ///
+    /// If both the node and its nested subtrees do not contain leaves, this will return `None`.
+    pub async fn compute_depth(
+        mut bs: impl AsyncBlockStoreRead,
+        node: Cid,
+    ) -> Result<Option<usize>, Error> {
+        // Recursively iterate through the tree until we encounter a leaf node, and then
+        // use that to calculate the depth of the entire tree.
+        let mut subtrees = vec![(node, 0usize)];
+
+        loop {
+            if let Some((subtree, depth)) = subtrees.pop() {
+                let node = Node::read_from(&mut bs, subtree).await?;
+                if let Some(layer) = node.layer() {
+                    return Ok(Some(depth + layer));
+                }
+
+                subtrees.extend(node.trees().cloned().zip(std::iter::repeat(depth + 1)));
+            } else {
+                return Ok(None);
+            }
+        }
+    }
+
     /// Traverse a merkle search tree.
     ///
     /// This executes the closure provided in `f` and takes the action
@@ -385,7 +410,7 @@ impl<S: AsyncBlockStoreRead + AsyncBlockStoreWrite> Tree<S> {
         // 2) The target layer is contained within the tree (and we will traverse to find it).
         // 3) The tree is currently empty (trivial).
         let mut node = match self.depth(None).await {
-            Ok(layer) => {
+            Ok(Some(layer)) => {
                 match layer.cmp(&target_layer) {
                     // The new key can be inserted into the root node.
                     Ordering::Equal => Node::read_from(&mut self.storage, node_cid).await?,
@@ -458,7 +483,7 @@ impl<S: AsyncBlockStoreRead + AsyncBlockStoreWrite> Tree<S> {
                     }
                 }
             }
-            Err(Error::EmptyTree) => {
+            Ok(None) => {
                 // The tree is currently empty.
                 Node { entries: vec![] }
             }
@@ -632,23 +657,8 @@ impl<S: AsyncBlockStoreRead> Tree<S> {
     }
 
     /// Compute the depth of the merkle search tree from either the specified node or the root
-    pub async fn depth(&mut self, node: Option<Cid>) -> Result<usize, Error> {
-        // Recursively iterate through the tree until we encounter a leaf node, and then
-        // use that to calculate the depth of the entire tree.
-        let mut subtrees = vec![(node.unwrap_or(self.root), 0usize)];
-
-        loop {
-            if let Some((subtree, depth)) = subtrees.pop() {
-                let node = Node::read_from(&mut self.storage, subtree).await?;
-                if let Some(layer) = node.layer() {
-                    return Ok(depth + layer);
-                }
-
-                subtrees.extend(node.trees().cloned().zip(std::iter::repeat(depth + 1)));
-            } else {
-                return Err(Error::EmptyTree);
-            }
-        }
+    pub async fn depth(&mut self, node: Option<Cid>) -> Result<Option<usize>, Error> {
+        algos::compute_depth(&mut self.storage, node.unwrap_or(self.root)).await
     }
 
     /// Returns a stream of all entries in this tree, in lexicographic order.
@@ -994,6 +1004,7 @@ pub enum Error {
 mod test {
     use std::str::FromStr;
 
+    use futures::TryStreamExt;
     use ipld_core::cid::multihash::Multihash;
 
     use crate::blockstore::{MemoryBlockStore, SHA2_256};