diff --git a/src/disktree/dptr.rs b/src/disktree/dptr.rs index 7500591..0ee907e 100644 --- a/src/disktree/dptr.rs +++ b/src/disktree/dptr.rs @@ -22,6 +22,25 @@ where Ok(dptr) } +/// Read 5 * `n` bytes from disk, for up to n=7, and parses them as +/// litte-endien `u64`s. +pub(crate) fn read_n(src: &mut R, n: usize) -> Result> +where + R: Read + Seek, +{ + assert!(n <= 7); + let mut buf = [0; DPTR_SZ * 7]; + src.read_exact(&mut buf[..(DPTR_SZ * n)])?; + Ok(buf[..(DPTR_SZ * n)] + .chunks(DPTR_SZ) + .map(|chunk| { + let mut buf = [0u8; size_of::()]; + buf[..DPTR_SZ].copy_from_slice(chunk); + u64::from_le_bytes(buf) + }) + .collect()) +} + /// Writes the 5 lower bytes of a `u64` to disk. pub(crate) fn write(dst: &mut W, dptr: u64) -> Result where diff --git a/src/disktree/iter.rs b/src/disktree/iter.rs index c8de2a7..65f55d8 100644 --- a/src/disktree/iter.rs +++ b/src/disktree/iter.rs @@ -48,24 +48,22 @@ where fn read_node(&mut self, dptr: u64) -> Result { self.seek_to(dptr)?; let node_tag = self.rdr.read_u8()?; - let base_pos = self.rdr.stream_position()?; + let base_pos = dptr + std::mem::size_of_val(&node_tag) as u64; + debug_assert_eq!(base_pos, self.rdr.stream_position().unwrap()); assert!(node_tag == 0 || node_tag > 0b1000_0000); - let node = if node_tag == 0 { - Node::Leaf(base_pos) + if node_tag == 0 { + Ok(Node::Leaf(base_pos)) } else { let mut buf = self.node_buf(); + let n_children = (node_tag & 0b0111_1111).count_ones() as usize; + let mut child_dptrs = dptr::read_n(&mut self.rdr, n_children)?; for digit in (0..7).rev() { if node_tag & (1 << digit) != 0 { - let bit_cnt = (((node_tag as u16) << (8 - digit)) & 0xFF).count_ones(); - let child_dptr_pos = base_pos + (bit_cnt as u64 * dptr::DPTR_SZ as u64); - self.seek_to(child_dptr_pos)?; - let child_dptr = dptr::read(&mut self.rdr)?; - buf.push((digit, child_dptr)); + buf.push((digit, child_dptrs.pop().unwrap())); } } - Node::Parent(buf) - }; - Ok(node) + Ok(Node::Parent(buf)) + } } fn node_buf(&mut self) -> Vec<(u8, u64)> {