Merge pull request #42 from JayKickliter/jsk/refactor-get

Implement get[_mut] on top of get_raw[_mut]

src/disktree/mod.rs

@@ -57,8 +57,22 @@ mod tests {
             .unwrap();
         let monaco_disktree = DiskTreeMap::open(path).unwrap();
 
-        assert_eq!(monaco.get(point_2).unzip().1, None);
-        assert_eq!(monaco.get(point_1).unzip().1, Some(&Region::Monaco));
+        assert_eq!(monaco.get(point_2), None);
+        assert_eq!(
+            monaco.get(point_1),
+            Some((point_1.to_parent(9).unwrap(), &Region::Monaco))
+        );
+
+        let point_1_res8 = point_1.to_parent(8).unwrap();
+        assert!(matches!(
+            monaco.get_raw(point_1_res8),
+            Some((cell, crate::node::Node::Parent(_))) if cell == point_1_res8
+        ));
+
+        assert!(matches!(
+            monaco_disktree.get_raw(point_1_res8).unwrap(),
+            Some((cell, crate::disktree::node::Node::Parent(_))) if cell == point_1_res8
+        ));
 
         for (ht_cell, &ht_val) in monaco.iter() {
             let now = std::time::Instant::now();

src/disktree/node.rs

@@ -8,6 +8,7 @@ use std::{io::Read, mem::size_of, ops::Range};
 // Enough bytes to read node tag and 7 child dptrs.
 const NODE_BUF_SZ: usize = size_of::<u8>() + 7 * Dp::size();
 
+#[derive(Debug)]
 pub(crate) enum Node {
     // value_begin..value_end
     Leaf(Range<usize>),

src/disktree/tree.rs

@@ -10,7 +10,6 @@ use std::{
     fs::File,
     io::{Cursor, Read, Seek, SeekFrom},
     marker::Send,
-    ops::Range,
     path::Path,
 };
 
@@ -62,6 +61,16 @@ impl DiskTreeMap {
 
     /// Returns `(Cell, &[u8])`, if present.
     pub fn get(&self, cell: Cell) -> Result<Option<(Cell, &[u8])>> {
+        if let Some((cell, Node::Leaf(range))) = self.get_raw(cell)? {
+            let val_bytes = &(*self.0).as_ref()[range];
+            Ok(Some((cell, val_bytes)))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Returns `(Cell, Node)`, if present.
+    pub(crate) fn get_raw(&self, cell: Cell) -> Result<Option<(Cell, Node)>> {
         let base_cell_pos = Self::base_cell_dptr(cell);
         let mut csr = Cursor::new((*self.0).as_ref());
         csr.seek(SeekFrom::Start(base_cell_pos.into()))?;
@@ -70,50 +79,42 @@ impl DiskTreeMap {
             return Ok(None);
         }
         let digits = Digits::new(cell);
-        if let Some((cell, range)) = Self::_get(&mut csr, 0, node_dptr, cell, digits)? {
-            let val_bytes = &(*self.0).as_ref()[range];
-            Ok(Some((cell, val_bytes)))
-        } else {
-            Ok(None)
-        }
-    }
-
-    /// Returns `true` if the tree fully contains `cell`.
-    pub fn contains(&self, cell: Cell) -> Result<bool> {
-        self.get(cell).map(|opt| opt.is_some())
+        Self::_get_raw(&mut csr, 0, node_dptr, cell, digits)
     }
 
-    /// Returns an iterator visiting all `(Cell, &[u8])` pairs in
-    /// arbitrary order.
-    pub fn iter(&self) -> Result<impl Iterator<Item = Result<(Cell, &[u8])>>> {
-        Iter::new((*self.0).as_ref())
-    }
-
-    fn _get(
+    fn _get_raw(
         csr: &mut Cursor<&[u8]>,
         res: u8,
         node_dptr: Dp,
         cell: Cell,
         mut digits: Digits,
-    ) -> Result<Option<(Cell, Range<usize>)>> {
+    ) -> Result<Option<(Cell, Node)>> {
         csr.seek(SeekFrom::Start(node_dptr.into()))?;
         let node = Node::read(csr)?;
-        match (digits.next(), node) {
-            (None, Node::Leaf(range)) => Ok(Some((cell, range))),
-            (Some(_), Node::Leaf(range)) => Ok(Some((
+        match (digits.next(), &node) {
+            (None, _) => Ok(Some((cell, node))),
+            (Some(_), Node::Leaf(_)) => Ok(Some((
                 cell.to_parent(res).expect("invalid condition"),
-                range,
+                node,
             ))),
             (Some(digit), Node::Parent(children)) => match children[digit as usize] {
                 None => Ok(None),
-                Some(dptr) => Self::_get(csr, res + 1, dptr, cell, digits),
+                Some(dptr) => Self::_get_raw(csr, res + 1, dptr, cell, digits),
             },
-            // No digits left, but `self` isn't full, so this cell
-            // can't fully contain the target.
-            (None, _) => Ok(None),
         }
     }
 
+    /// Returns `true` if the tree fully contains `cell`.
+    pub fn contains(&self, cell: Cell) -> Result<bool> {
+        self.get(cell).map(|opt| opt.is_some())
+    }
+
+    /// Returns an iterator visiting all `(Cell, &[u8])` pairs in
+    /// arbitrary order.
+    pub fn iter(&self) -> Result<impl Iterator<Item = Result<(Cell, &[u8])>>> {
+        Iter::new((*self.0).as_ref())
+    }
+
     /// 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)

src/hex_tree_map.rs

@@ -172,7 +172,16 @@ impl<V, C> HexTreeMap<V, C> {
     ///
     /// Note that this method also returns a Cell, which may be a
     /// parent of the target cell provided.
+    #[inline]
     pub fn get(&self, cell: Cell) -> Option<(Cell, &V)> {
+        match self.get_raw(cell) {
+            Some((cell, Node::Leaf(val))) => Some((cell, val)),
+            _ => None,
+        }
+    }
+
+    #[inline]
+    pub(crate) fn get_raw(&self, cell: Cell) -> Option<(Cell, &Node<V>)> {
         let base_cell = cell.base();
         match self.nodes[base_cell as usize].as_ref() {
             Some(node) => {
@@ -188,7 +197,16 @@ impl<V, C> HexTreeMap<V, C> {
     ///
     /// Note that this method also returns a Cell, which may be a
     /// parent of the target cell provided.
+    #[inline]
     pub fn get_mut(&mut self, cell: Cell) -> Option<(Cell, &mut V)> {
+        match self.get_raw_mut(cell) {
+            Some((cell, &mut Node::Leaf(ref mut val))) => Some((cell, val)),
+            _ => None,
+        }
+    }
+
+    #[inline]
+    pub(crate) fn get_raw_mut(&mut self, cell: Cell) -> Option<(Cell, &mut Node<V>)> {
         let base_cell = cell.base();
         match self.nodes[base_cell as usize].as_mut() {
             Some(node) => {
@@ -231,7 +249,7 @@ impl<V, C> HexTreeMap<V, C> {
         match self.nodes[base_cell as usize].as_ref() {
             Some(node) => {
                 let digits = Digits::new(cell);
-                match node.get_raw(0, cell, digits) {
+                match node.get(0, cell, digits) {
                     Some((cell, Node::Leaf(val))) => Some((cell, val))
                         .into_iter()
                         .chain(crate::iteration::Iter::empty()),
@@ -252,7 +270,7 @@ impl<V, C> HexTreeMap<V, C> {
         match self.nodes[base_cell as usize].as_mut() {
             Some(node) => {
                 let digits = Digits::new(cell);
-                match node.get_raw_mut(0, cell, digits) {
+                match node.get_mut(0, cell, digits) {
                     Some((cell, Node::Leaf(val))) => Some((cell, val))
                         .into_iter()
                         .chain(crate::iteration::IterMut::empty()),

src/node.rs

@@ -7,7 +7,7 @@ use crate::{compaction::Compactor, digits::Digits, Cell};
 //
 // The benefit of storing indices is vastly simpler Cell+Value
 // iteration of a tree.
-#[derive(Clone, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[repr(align(64))]
 pub(crate) enum Node<V> {
@@ -103,45 +103,22 @@ impl<V> Node<V> {
         }
     }
 
-    pub(crate) fn get(&self, res: u8, cell: Cell, mut digits: Digits) -> Option<(Cell, &V)> {
-        match (digits.next(), self) {
-            (None, Self::Leaf(val)) => Some((cell, val)),
-            (Some(_), Self::Leaf(val)) => {
-                Some((cell.to_parent(res).expect("invalid condition"), val))
-            }
-            (Some(digit), Self::Parent(children)) => match &children.as_slice()[digit as usize] {
-                Some(node) => node.get(res + 1, cell, digits),
-                None => None,
-            },
-            // No digits left, but `self` isn't full, so this cell
-            // can't fully contain the target.
-            (None, Self::Parent(_)) => None,
-        }
-    }
-
-    /// Returns a raw [`Node`] for `cell`, if present.
-    ///
-    /// Unlike the public [`Node::get`], this function returns [`Some`] for
-    /// parent nodes.
-    pub(crate) fn get_raw(
-        &self,
-        res: u8,
-        cell: Cell,
-        mut digits: Digits,
-    ) -> Option<(Cell, &Node<V>)> {
+    #[inline]
+    pub(crate) fn get(&self, res: u8, cell: Cell, mut digits: Digits) -> Option<(Cell, &Node<V>)> {
         match (digits.next(), self) {
             (None, _) => Some((cell, self)),
             (Some(_), Self::Leaf(_)) => {
                 Some((cell.to_parent(res).expect("invalid condition"), self))
             }
             (Some(digit), Self::Parent(children)) => match &children.as_slice()[digit as usize] {
-                Some(node) => node.get_raw(res + 1, cell, digits),
+                Some(node) => node.get(res + 1, cell, digits),
                 None => None,
             },
         }
     }
 
-    pub(crate) fn get_raw_mut(
+    #[inline]
+    pub(crate) fn get_mut(
         &mut self,
         res: u8,
         cell: Cell,
@@ -154,33 +131,10 @@ impl<V> Node<V> {
             }
             (Some(digit), Self::Parent(ref mut children)) => {
                 match children.as_mut_slice()[digit as usize].as_deref_mut() {
-                    Some(node) => node.get_raw_mut(res + 1, cell, digits),
-                    None => None,
-                }
-            }
-        }
-    }
-
-    pub(crate) fn get_mut(
-        &mut self,
-        res: u8,
-        cell: Cell,
-        mut digits: Digits,
-    ) -> Option<(Cell, &mut V)> {
-        match (digits.next(), self) {
-            (None, Self::Leaf(val)) => Some((cell, val)),
-            (Some(_), Self::Leaf(val)) => {
-                Some((cell.to_parent(res).expect("invalid condition"), val))
-            }
-            (Some(digit), Self::Parent(children)) => {
-                match &mut children.as_mut_slice()[digit as usize] {
                     Some(node) => node.get_mut(res + 1, cell, digits),
                     None => None,
                 }
             }
-            // No digits left, but `self` isn't full, so this cell
-            // can't fully contain the target.
-            (None, Self::Parent(_)) => None,
         }
     }
 }