Strongly type disktree Node type

src/disktree/mod.rs

@@ -5,6 +5,7 @@ pub use value::ReadVal;
 
 mod dptr;
 mod iter;
+mod node;
 mod tree;
 mod value;
 mod writer;

src/disktree/node.rs

@@ -0,0 +1,43 @@
+use crate::{disktree::dptr::Dptr, error::Result};
+use byteorder::ReadBytesExt;
+use std::{
+    io::{Read, Seek},
+    mem::size_of,
+};
+
+// Enough bytes to read node tag and 7 child dptrs.
+const NODE_BUF_SZ: usize = size_of::<u8>() + 7 * Dptr::size() as usize;
+
+pub(crate) enum Node {
+    // File position for the fist byte of value data.
+    Leaf(Dptr),
+    // (H3 Cell digit, file position of child's node tag)
+    Parent([Option<Dptr>; 7]),
+}
+
+impl Node {
+    pub(crate) fn read<R>(rdr: &mut R) -> Result<Node>
+    where
+        R: Seek + Read,
+    {
+        // dptr to either leaf value of first child dptr
+        let base_pos = Dptr::from(rdr.stream_position()? + size_of::<u8>() as u64);
+        let mut buf = [0u8; NODE_BUF_SZ];
+        let bytes_read = rdr.read(&mut buf)?;
+        let buf_rdr = &mut &buf[..bytes_read];
+        let node_tag = buf_rdr.read_u8()?;
+        assert!(node_tag == 0 || node_tag > 0b1000_0000);
+        if node_tag == 0 {
+            Ok(Node::Leaf(base_pos))
+        } else {
+            let mut children: [Option<Dptr>; 7] = [None, None, None, None, None, None, None];
+            for (_digit, child) in (0..7)
+                .zip(children.iter_mut())
+                .filter(|(digit, _)| node_tag & (1 << digit) != 0)
+            {
+                *child = Some(Dptr::read(buf_rdr)?);
+            }
+            Ok(Node::Parent(children))
+        }
+    }
+}

src/disktree/tree.rs

@@ -1,6 +1,6 @@
 use crate::{
     digits::Digits,
-    disktree::{dptr::Dptr, iter::Iter, ReadVal},
+    disktree::{dptr::Dptr, iter::Iter, node::Node, ReadVal},
     error::{Error, Result},
     Cell,
 };
@@ -55,7 +55,8 @@ impl<R: Read + Seek> DiskTree<R> {
             return Ok(None);
         }
         let digits = Digits::new(cell);
-        if let Some((cell, _)) = self._get(0, node_dptr, cell, digits)? {
+        if let Some((cell, dptr)) = self._get(0, node_dptr, cell, digits)? {
+            self.seek_to_pos(dptr.into())?;
             Ok(Some((cell, &mut self.0)))
         } else {
             Ok(None)
@@ -98,26 +99,19 @@ impl<R: Read + Seek> DiskTree<R> {
         node_dptr: Dptr,
         cell: Cell,
         mut digits: Digits,
-    ) -> Result<Option<(Cell, u64)>> {
+    ) -> Result<Option<(Cell, Dptr)>> {
         self.seek_to_pos(node_dptr)?;
-        let node_tag = self.0.read_u8()?;
-        assert!(node_tag == 0 || node_tag > 0b1000_0000);
-        match (digits.next(), node_tag) {
-            (None, 0) => Ok(Some((cell, self.0.stream_position()?))),
-            (Some(_), 0) => Ok(Some((
+        let node = Node::read(&mut self.0)?;
+        match (digits.next(), node) {
+            (None, Node::Leaf(dptr)) => Ok(Some((cell, dptr))),
+            (Some(_), Node::Leaf(dptr)) => Ok(Some((
                 cell.to_parent(res).expect("invalid condition"),
-                self.0.stream_position()?,
+                dptr,
             ))),
-            (Some(digit), _) => {
-                let bit_cnt = (((node_tag as u16) << (8 - digit)) & 0xFF).count_ones();
-                self.seek_forward(u64::from(bit_cnt) * Dptr::size())?;
-                let child_dptr = Dptr::read(&mut self.0)?;
-                if child_dptr.is_null() {
-                    Ok(None)
-                } else {
-                    self._get(res + 1, child_dptr, cell, digits)
-                }
-            }
+            (Some(digit), Node::Parent(children)) => match children[digit as usize] {
+                None => Ok(None),
+                Some(dptr) => self._get(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),
@@ -129,11 +123,6 @@ impl<R: Read + Seek> DiskTree<R> {
         Ok(())
     }
 
-    fn seek_forward(&mut self, offset: u64) -> Result {
-        self.0.seek(SeekFrom::Current(offset as i64))?;
-        Ok(())
-    }
-
     /// Returns the DPtr to a base (res0) cell dptr.
     fn base_cell_dptr(cell: Cell) -> Dptr {
         Dptr::from(HDR_SZ + Dptr::size() * (cell.base() as u64))

src/disktree/writer.rs

@@ -88,7 +88,7 @@ impl<W: Write + Seek> DiskTreeWriter<W> {
                             node_fixups.push((self.pos()?, node));
                             Dptr::null().write(&mut self.0)?;
                         }
-                    }
+                    };
                 }
                 self.seek_to(tag_pos)?;
                 // Make the top bit 1 as a sentinel.