Impl DiskTreeMap::subtree_iter

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> {

src/disktree/mod.rs

@@ -237,4 +237,63 @@ mod tests {
         let disktree = DiskTreeMap::with_buf(wtr).unwrap();
         assert_eq!(0, disktree.iter().unwrap().count());
     }
+
+    #[test]
+    fn test_subtree_iter() {
+        use crate::{compaction::NullCompactor, Cell, HexTreeMap};
+        use h3o::{CellIndex, Resolution};
+        use std::{convert::TryFrom, io::Cursor};
+
+        // https://wolf-h3-viewer.glitch.me/?h3=863969a47ffffff
+        let monaco_res6_ci = CellIndex::try_from(0x863969a47ffffff).unwrap();
+        let monaco_res6_cell = Cell::try_from(u64::from(monaco_res6_ci)).unwrap();
+        // https://wolf-h3-viewer.glitch.me/?h3=863969a6fffffff
+        let not_monaco_res6_ci = CellIndex::try_from(0x863969a6fffffff).unwrap();
+
+        let monaco_res10_cells = monaco_res6_ci
+            .children(Resolution::Ten)
+            .map(|ci| Cell::try_from(u64::from(ci)).unwrap())
+            .collect::<Vec<_>>();
+
+        let not_monaco_res10_cells = not_monaco_res6_ci
+            .children(Resolution::Ten)
+            .map(|ci| Cell::try_from(u64::from(ci)).unwrap())
+            .collect::<Vec<_>>();
+
+        let monaco_hextree: HexTreeMap<(), NullCompactor> = monaco_res10_cells
+            .iter()
+            .copied()
+            .map(|cell| (cell, ()))
+            .collect();
+
+        let combined_hextree: HexTreeMap<(), NullCompactor> = monaco_res10_cells
+            .iter()
+            .chain(not_monaco_res10_cells.iter())
+            .copied()
+            .map(|cell| (cell, ()))
+            .collect();
+
+        let combined_disktree = {
+            let mut combined_disktree_buf = vec![];
+            combined_hextree
+                .to_disktree(Cursor::new(&mut combined_disktree_buf), |wtr, ()| {
+                    wtr.write_all(&[])
+                })
+                .unwrap();
+            DiskTreeMap::with_buf(combined_disktree_buf).unwrap()
+        };
+
+        assert_eq!(
+            combined_hextree.len(),
+            combined_disktree.iter().unwrap().count()
+        );
+
+        let combined_disktree_subtree_collect = combined_disktree
+            .subtree_iter(monaco_res6_cell)
+            .unwrap()
+            .map(|item| item.unwrap().0)
+            .collect::<Vec<_>>();
+        let monaco_hextree_collect = monaco_hextree.iter().map(|item| item.0).collect::<Vec<_>>();
+        assert_eq!(combined_disktree_subtree_collect, monaco_hextree_collect);
+    }
 }

src/disktree/tree.rs

@@ -115,6 +115,26 @@ impl DiskTreeMap {
         Iter::new((*self.0).as_ref())
     }
 
+    /// Returns an iterator visiting the specified `cell` or its descendants.
+    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.
     fn base_cell_dptr(cell: Cell) -> Dp {
         Dp::from(HDR_SZ + Dp::size() * cell.base() as usize)