Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Feb 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions src/disktree/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
1 change: 1 addition & 0 deletions src/disktree/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>),
Expand Down
57 changes: 29 additions & 28 deletions src/disktree/tree.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ use std::{
fs::File,
io::{Cursor, Read, Seek, SeekFrom},
marker::Send,
ops::Range,
path::Path,
};

Expand Down Expand Up @@ -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()))?;
Expand All @@ -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)
Expand Down
22 changes: 20 additions & 2 deletions src/hex_tree_map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand All @@ -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) => {
Expand Down Expand Up @@ -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()),
Expand All @@ -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()),
Expand Down
58 changes: 6 additions & 52 deletions src/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down Expand Up @@ -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,
Expand All @@ -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,
}
}
}
Loading