From 12b82adfad4e19f25a8725d3d8d0551cd57d5484 Mon Sep 17 00:00:00 2001
From: Jay Kickliter <jay@kickliter.com>
Date: Wed, 28 Feb 2024 12:19:30 -0700
Subject: [PATCH] wip

---
 src/disktree/iter.rs | 23 ++++++++++++++++++++++-
 src/disktree/tree.rs | 19 +++++++++++++++++--
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/disktree/iter.rs b/src/disktree/iter.rs
index d66ccd5..83bc63e 100644
--- a/src/disktree/iter.rs
+++ b/src/disktree/iter.rs
@@ -1,6 +1,6 @@
 use crate::{
     cell::CellStack,
-    disktree::{dptr::Dp, dtseek::DtSeek, tree::HDR_SZ, varint},
+    disktree::{dptr::Dp, dtseek::DtSeek, tree::HDR_SZ, varint, DiskTreeMap},
     error::{Error, Result},
     Cell,
 };
@@ -106,6 +106,27 @@ impl<'a> Iter<'a> {
             node_stack,
         })
     }
+
+    pub(crate) fn empty(disktree_buf: &'a [u8]) -> Iter<'a> {
+        let disktree_csr = Cursor::new(disktree_buf);
+        let cell_stack = CellStack::new();
+        let node_stack = Vec::new();
+        let recycle_bin = Vec::new();
+        let curr_node = None;
+        Self {
+            cell_stack,
+            curr_node,
+            disktree_buf,
+            disktree_csr,
+            recycle_bin,
+            node_stack,
+        }
+    }
+
+    /// Creates a new `Iter` over `cell` and its descendants.
+    pub(crate) fn subtree(_disktree: &'a DiskTreeMap, _cell: Cell) -> Result<Iter<'a>> {
+        unimplemented!()
+    }
 }
 
 impl<'a> Iterator for Iter<'a> {
diff --git a/src/disktree/tree.rs b/src/disktree/tree.rs
index b38ca1c..bdac5f0 100755
--- a/src/disktree/tree.rs
+++ b/src/disktree/tree.rs
@@ -116,8 +116,23 @@ impl DiskTreeMap {
     }
 
     /// Returns an iterator visiting the specified `cell` or its descendants.
-    pub fn subtree_iter(&self, _cell: Cell) -> Result<impl Iterator<Item = Result<(Cell, &[u8])>>> {
-        Ok(None.into_iter())
+    pub fn subtree_iter(&self, cell: Cell) -> Result<impl Iterator<Item = Result<(Cell, &[u8])>>> {
+        let iter = match self.get_raw(cell)? {
+            None => {
+                let iter = crate::disktree::iter::Iter::empty((*self.0).as_ref());
+                None.into_iter().chain(iter)
+            }
+            Some((cell, Node::Leaf(range))) => {
+                let iter = crate::disktree::iter::Iter::empty((*self.0).as_ref());
+                let val_bytes = &(*self.0).as_ref()[range];
+                Some(Ok((cell, val_bytes))).into_iter().chain(iter)
+            }
+            Some((cell, Node::Parent(_children))) => {
+                let iter = crate::disktree::iter::Iter::subtree(self, cell)?;
+                None.into_iter().chain(iter)
+            }
+        };
+        Ok(iter)
     }
 
     /// Returns the DPtr to a base (res0) cell dptr.