From 9f50399539220a89ae84f3bd123026c33bb5a28d Mon Sep 17 00:00:00 2001 From: Jay Kickliter Date: Tue, 27 Feb 2024 09:48:52 -0700 Subject: [PATCH] Implement get[_mut] on top of get_raw[_mut] --- src/disktree/mod.rs | 18 ++++++++++++-- src/disktree/node.rs | 1 + src/disktree/tree.rs | 57 ++++++++++++++++++++++--------------------- src/hex_tree_map.rs | 22 +++++++++++++++-- src/node.rs | 58 +++++--------------------------------------- 5 files changed, 72 insertions(+), 84 deletions(-) diff --git a/src/disktree/mod.rs b/src/disktree/mod.rs index c6f95a6..ec416ad 100644 --- a/src/disktree/mod.rs +++ b/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(); diff --git a/src/disktree/node.rs b/src/disktree/node.rs index b9a0186..6981083 100644 --- a/src/disktree/node.rs +++ b/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::() + 7 * Dp::size(); +#[derive(Debug)] pub(crate) enum Node { // value_begin..value_end Leaf(Range), diff --git a/src/disktree/tree.rs b/src/disktree/tree.rs index 6b521bb..ac7897e 100755 --- a/src/disktree/tree.rs +++ b/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> { + 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> { 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 { - 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>> { - 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)>> { + ) -> Result> { 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 { + self.get(cell).map(|opt| opt.is_some()) + } + + /// Returns an iterator visiting all `(Cell, &[u8])` pairs in + /// arbitrary order. + pub fn iter(&self) -> Result>> { + 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) diff --git a/src/hex_tree_map.rs b/src/hex_tree_map.rs index d5e63fe..4ffc601 100644 --- a/src/hex_tree_map.rs +++ b/src/hex_tree_map.rs @@ -172,7 +172,16 @@ impl HexTreeMap { /// /// 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)> { let base_cell = cell.base(); match self.nodes[base_cell as usize].as_ref() { Some(node) => { @@ -188,7 +197,16 @@ impl HexTreeMap { /// /// 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)> { let base_cell = cell.base(); match self.nodes[base_cell as usize].as_mut() { Some(node) => { @@ -231,7 +249,7 @@ impl HexTreeMap { 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 HexTreeMap { 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()), diff --git a/src/node.rs b/src/node.rs index 26274f8..7542c3f 100644 --- a/src/node.rs +++ b/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 { @@ -103,45 +103,22 @@ impl Node { } } - 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)> { + #[inline] + pub(crate) fn get(&self, res: u8, cell: Cell, mut digits: Digits) -> Option<(Cell, &Node)> { 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 Node { } (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, } } }