diff --git a/Cargo.toml b/Cargo.toml index 29a2eb8..3d9d5bb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -129,6 +129,7 @@ name = "graft_masked_branches" harness = false [[bench]] +name = "byte_mask" name = "zipper_head_owned" harness = false diff --git a/benches/byte_mask.rs b/benches/byte_mask.rs new file mode 100644 index 0000000..d881fd7 --- /dev/null +++ b/benches/byte_mask.rs @@ -0,0 +1,81 @@ +use divan::{black_box, Bencher, Divan}; + +use pathmap::utils::ByteMask; + +fn main() { + let divan = Divan::from_args().sample_count(4000); + + divan.main(); +} + +fn spread_mask(on_bits: usize) -> ByteMask { + debug_assert!(on_bits <= 256); + + (0..on_bits) + .map(|idx| ((idx * 73 + 19) & 0xFF) as u8) + .collect() +} + +fn iter_mask(mask: ByteMask) -> u64 { + let mut acc = 0u64; + let mut count = 0u64; + for byte in black_box(mask).iter() { + let byte = black_box(byte) as u64; + acc = acc.wrapping_mul(257).wrapping_add(byte + count); + count += 1; + } + black_box(acc ^ count) +} + +#[divan::bench(args = [0, 1, 2, 4, 8, 16, 150, 200, 256])] +fn bytemask_iter(bencher: Bencher, on_bits: usize) { + let mask = spread_mask(on_bits); + let mut sink = 0u64; + + bencher.bench_local(|| { + sink = sink.wrapping_add(iter_mask(mask)); + black_box(sink); + }); +} + +fn recursive_masks(depth: usize, on_bits: usize) -> Vec { + debug_assert!(on_bits <= 2); + + (0..depth) + .map(|level| { + (0..on_bits) + .map(|idx| ((level * 37 + idx * 131 + 11) & 0xFF) as u8) + .collect() + }) + .collect() +} + +fn recursive_iter(masks: &[ByteMask], level: usize, acc: u64) -> u64 { + if level == masks.len() { + return black_box(acc); + } + + let mut out = acc; + let mut iter = black_box(masks[level]).iter(); + if let Some(byte) = iter.next() { + out = out.wrapping_add(recursive_iter(masks, level + 1, acc.wrapping_mul(257).wrapping_add(black_box(byte) as u64))); + } + if let Some(byte) = iter.next() { + out = out.wrapping_add(black_box(byte) as u64); + } + black_box(&mut iter); + black_box(out) +} + +#[divan::bench(args = [1, 2])] +fn bytemask_iter_recursive_stack(bencher: Bencher, on_bits: usize) { + const DEPTH: usize = 50; + + let masks = recursive_masks(DEPTH, on_bits); + let mut sink = 0u64; + + bencher.bench_local(|| { + sink = sink.wrapping_add(recursive_iter(&masks, 0, 0)); + black_box(sink); + }); +} diff --git a/src/bridge_node.rs b/src/bridge_node.rs index 56135a6..15ca534 100644 --- a/src/bridge_node.rs +++ b/src/bridge_node.rs @@ -398,11 +398,11 @@ impl TrieNode for BridgeNode { self.is_empty() } #[inline(always)] - fn new_iter_token(&self) -> u128 { + fn new_iter_token(&self) -> IterToken { 0 } #[inline(always)] - fn iter_token_for_path(&self, key: &[u8]) -> (u128, &[u8]) { + fn iter_token_for_path(&self, key: &[u8]) -> (IterToken, &[u8]) { let node_key = self.key(); if key.len() <= node_key.len() { let short_key = &node_key[..key.len()]; @@ -416,7 +416,7 @@ impl TrieNode for BridgeNode { (NODE_ITER_FINISHED, &[]) } #[inline(always)] - fn next_items(&self, token: u128) -> (u128, &[u8], Option<&TrieNodeODRc>, Option<&V>) { + fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc>, Option<&V>) { if token == 0 { let node_key = self.key(); if self.is_used_child() { diff --git a/src/dense_byte_node.rs b/src/dense_byte_node.rs index ca253c3..22cf8f2 100644 --- a/src/dense_byte_node.rs +++ b/src/dense_byte_node.rs @@ -303,6 +303,52 @@ impl> ByteNode unsafe{ self.values.get_unchecked_mut(ix) } } + // ------ IterToken format for ByteNode ------ + // Iter tokens pack the next byte position to inspect with the corresponding index into `values`. + // The low 9 bits store one more than the last returned byte, or 0 before iteration starts. The + // upper bits store the index of the next CoFree in `values`, avoiding a `mask.index_of` recompute + // on every item. `iter_token_for_path` computes the values index once for the requested path. + const ITER_TOKEN_BYTE_BITS: u64 = 9; + const ITER_TOKEN_BYTE_MASK: IterToken = (1 << Self::ITER_TOKEN_BYTE_BITS) - 1; + + #[inline(always)] + fn iter_token(next_byte: u16, values_idx: u16) -> IterToken { + (values_idx as IterToken) << Self::ITER_TOKEN_BYTE_BITS | next_byte as IterToken + } + + #[inline(always)] + fn iter_token_next_byte(token: IterToken) -> u16 { + (token & Self::ITER_TOKEN_BYTE_MASK) as u16 + } + + #[inline(always)] + fn iter_token_values_idx(token: IterToken) -> usize { + (token >> Self::ITER_TOKEN_BYTE_BITS) as usize + } + + #[inline(always)] + fn next_iter_item_from(&self, token: IterToken) -> Option<(u8, usize)> { + let start = Self::iter_token_next_byte(token); + if start >= 256 { + return None; + } + + let mut word_idx = (start >> 6) as usize; + let mut bit_idx = (start & 0x3F) as u32; + loop { + let word = unsafe { *self.mask.0.get_unchecked(word_idx) } & (!0u64 << bit_idx); + if word != 0 { + return Some(((word_idx as u8) * 64 + word.trailing_zeros() as u8, Self::iter_token_values_idx(token))); + } + + word_idx += 1; + if word_idx == 4 { + return None; + } + bit_idx = 0; + } + } + #[inline] fn is_empty(&self) -> bool { self.mask.is_empty_mask() @@ -923,49 +969,32 @@ impl> TrieNode self.values.len() == 0 } #[inline(always)] - fn new_iter_token(&self) -> u128 { - self.mask.0[0] as u128 + fn new_iter_token(&self) -> IterToken { + Self::iter_token(0, 0) } #[inline(always)] - fn iter_token_for_path(&self, key: &[u8]) -> u128 { + fn iter_token_for_path(&self, key: &[u8]) -> IterToken { if key.len() != 1 { self.new_iter_token() } else { - let k = *unsafe{ key.get_unchecked(0) } as usize; - let idx = (k & 0b11000000) >> 6; - let bit_i = k & 0b00111111; - debug_assert!(idx < 4); - let mask: u64 = if bit_i+1 < 64 { - (0xFFFFFFFFFFFFFFFF << bit_i+1) & unsafe{ self.mask.0.get_unchecked(idx) } - } else { - 0 - }; - ((idx as u128) << 64) | (mask as u128) + let key_byte = unsafe{ *key.get_unchecked(0) }; + let mut values_idx = self.mask.index_of(key_byte); + if self.mask.test_bit(key_byte) { + values_idx += 1; + } + Self::iter_token(key_byte as u16 + 1, values_idx as u16) } } #[inline(always)] - fn next_items(&self, token: u128) -> (u128, &[u8], Option<&TrieNodeODRc>, Option<&V>) { - let mut i = (token >> 64) as u8; - let mut w = token as u64; - loop { - if w != 0 { - let wi = w.trailing_zeros() as u8; - w ^= 1u64 << wi; - let k = i*64 + wi; - - let new_token = ((i as u128) << 64) | (w as u128); - let cf = unsafe{ self.get_unchecked(k) }; - let k = k as usize; - return (new_token, &ALL_BYTES[k..=k], cf.rec(), cf.val()) - - } else if i < 3 { - i += 1; + fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc>, Option<&V>) { + let Some((k, values_idx)) = self.next_iter_item_from(token) else { + return (NODE_ITER_FINISHED, &[], None, None); + }; - w = unsafe { *self.mask.0.get_unchecked(i as usize) }; - } else { - return (NODE_ITER_FINISHED, &[], None, None) - } - } + let next_token = Self::iter_token(k as u16 + 1, values_idx as u16 + 1); + let cf = unsafe{ self.values.get_unchecked(values_idx) }; + let k = k as usize; + (next_token, &ALL_BYTES[k..=k], cf.rec(), cf.val()) } fn node_val_count(&self, cache: &mut HashMap) -> usize { //Discussion: These two implementations do the same thing but with a slightly different ordering of @@ -2378,3 +2407,48 @@ fn bit_siblings() { assert_eq!(63, bit_sibling(63, 1u64 << 63, false)); assert_eq!(63, bit_sibling(63, 1u64 << 63, true)); } + +#[test] +fn byte_node_iter_token_crosses_mask_word_boundaries() { + let mut node = DenseByteNode::new_in(crate::alloc::global_alloc()); + for byte in [0, 63, 64, 127, 128, 191, 192, 255] { + node.set_val(byte, byte); + } + + let mut token = node.new_iter_token(); + let mut visited = Vec::new(); + while token != NODE_ITER_FINISHED { + let (next_token, path, _child, value) = node.next_items(token); + token = next_token; + if token != NODE_ITER_FINISHED { + assert_eq!(path.len(), 1); + assert_eq!(Some(&path[0]), value); + visited.push(path[0]); + } + } + assert_eq!(visited, [0, 63, 64, 127, 128, 191, 192, 255]); + + let mut token = node.iter_token_for_path(&[127]); + let mut visited_after_127 = Vec::new(); + while token != NODE_ITER_FINISHED { + let (next_token, path, _child, value) = node.next_items(token); + token = next_token; + if token != NODE_ITER_FINISHED { + assert_eq!(Some(&path[0]), value); + visited_after_127.push(path[0]); + } + } + assert_eq!(visited_after_127, [128, 191, 192, 255]); + + let mut token = node.iter_token_for_path(&[65]); + let mut visited_after_missing_65 = Vec::new(); + while token != NODE_ITER_FINISHED { + let (next_token, path, _child, value) = node.next_items(token); + token = next_token; + if token != NODE_ITER_FINISHED { + assert_eq!(Some(&path[0]), value); + visited_after_missing_65.push(path[0]); + } + } + assert_eq!(visited_after_missing_65, [127, 128, 191, 192, 255]); +} diff --git a/src/empty_node.rs b/src/empty_node.rs index 98159ea..722952a 100644 --- a/src/empty_node.rs +++ b/src/empty_node.rs @@ -58,13 +58,13 @@ impl TrieNode for EmptyNode { } fn node_remove_unmasked_branches(&mut self, _key: &[u8], _mask: ByteMask, _prune: bool) {} fn node_is_empty(&self) -> bool { true } - fn new_iter_token(&self) -> u128 { + fn new_iter_token(&self) -> IterToken { 0 } - fn iter_token_for_path(&self, _key: &[u8]) -> u128 { + fn iter_token_for_path(&self, _key: &[u8]) -> IterToken { 0 } - fn next_items(&self, _token: u128) -> (u128, &[u8], Option<&TrieNodeODRc>, Option<&V>) { + fn next_items(&self, _token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc>, Option<&V>) { (NODE_ITER_FINISHED, &[], None, None) } fn node_val_count(&self, _cache: &mut HashMap) -> usize { diff --git a/src/line_list_node.rs b/src/line_list_node.rs index 92c3133..bfc650b 100644 --- a/src/line_list_node.rs +++ b/src/line_list_node.rs @@ -1893,7 +1893,7 @@ impl TrieNode for LineListNode // * NODE_ITER_FINISHED // *==--==**==--==**==--==**==--==**==--==**==--==**==--==**==--==**==--==**==--==**==--==**==--==* #[inline(always)] - fn new_iter_token(&self) -> u128 { + fn new_iter_token(&self) -> IterToken { 0 } /// Explanation of logic: The ListNode contains a sorted list of keys (up to 2 of them), and the @@ -1904,7 +1904,7 @@ impl TrieNode for LineListNode /// - == key1, we should return (2, key1) /// - > key1, (NODE_ITER_FINISHED, &[]) #[inline(always)] - fn iter_token_for_path(&self, key: &[u8]) -> u128 { + fn iter_token_for_path(&self, key: &[u8]) -> IterToken { if key.len() == 0 { return 0 } @@ -1921,7 +1921,7 @@ impl TrieNode for LineListNode NODE_ITER_FINISHED } #[inline(always)] - fn next_items(&self, token: u128) -> (u128, &[u8], Option<&TrieNodeODRc>, Option<&V>) { + fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc>, Option<&V>) { match token { 0 => { if !self.is_used::<0>() { diff --git a/src/old_cursor.rs b/src/old_cursor.rs index bbacc07..d1c9f36 100644 --- a/src/old_cursor.rs +++ b/src/old_cursor.rs @@ -249,7 +249,7 @@ impl <'a, V : Clone + Send + Sync> Iterator for ByteTrieNodeIter<'a, V> { pub struct PathMapCursor<'a, V: Clone + Send + Sync> { prefix_buf: Vec, - btnis: Vec<(TaggedNodeRef<'a, V, GlobalAlloc>, u128, usize)>, + btnis: Vec<(TaggedNodeRef<'a, V, GlobalAlloc>, IterToken, usize)>, } impl <'a, V : Clone + Send + Sync + Unpin> PathMapCursor<'a, V> { diff --git a/src/tiny_node.rs b/src/tiny_node.rs index 08d61fc..5cdb357 100644 --- a/src/tiny_node.rs +++ b/src/tiny_node.rs @@ -214,9 +214,9 @@ impl<'a, V: Clone + Send + Sync, A: Allocator> TrieNode for TinyRefNode<'a fn node_is_empty(&self) -> bool { self.header & (1 << 7) == 0 } - fn new_iter_token(&self) -> u128 { unreachable!() } - fn iter_token_for_path(&self, _key: &[u8]) -> u128 { unreachable!() } - fn next_items(&self, _token: u128) -> (u128, &'a[u8], Option<&TrieNodeODRc>, Option<&V>) { unreachable!() } + fn new_iter_token(&self) -> IterToken { unreachable!() } + fn iter_token_for_path(&self, _key: &[u8]) -> IterToken { unreachable!() } + fn next_items(&self, _token: IterToken) -> (IterToken, &'a[u8], Option<&TrieNodeODRc>, Option<&V>) { unreachable!() } fn node_val_count(&self, cache: &mut HashMap) -> usize { let temp_node = self.into_full().unwrap(); temp_node.node_val_count(cache) diff --git a/src/trie_node.rs b/src/trie_node.rs index 82a8b16..570c38c 100644 --- a/src/trie_node.rs +++ b/src/trie_node.rs @@ -187,26 +187,22 @@ pub(crate) trait TrieNode: TrieNodeDowncas /// Generates a new iter token, to iterate the children and values contained within this node /// - /// GOAT: Do we really *need* 128 bits for the iter token? Or could we use 64? The idea is that an - /// iter token can represent any position in any arbitrary node type, and involve a minimum of computation - /// to advance to the next position. Currently the only node type that uses more than 64 bits is the - /// ByteNode, and that is because it represents the mask in the first 64 bits of the token. However it - /// seems just as efficient (I think actually slightly more efficient) to store both one more than the path - /// byte returned by the last call (or 0 if iteration is starting) and the index in the values vec. This means - /// we actually only need 16 bits for the byte node. + /// The iter token is a node-local cursor. It must represent any position within a node type, and + /// should involve a minimum of computation to advance to the next position, but it does not need to + /// encode an arbitrary path through the trie. /// - /// To the more general question of whether it will be enough for any possible future node structure, that - /// is a more difficult consideration. Currently MAX_NODE_KEY_BYTES is limited to 48, but there is no limit - /// on the branching factor within that node. So even 128 bits is insufficient to encode all paths in theory. - /// However a fixed-size node structure has a physical limit on its complexity. If we assume we will limit a - /// node to 4KB, 12 bits is enough to address any byte within that physical structure, so there is probably some - /// clever encoding that can address any path that it could contain, using 64 bits, with a reasonable time and - /// memory-fetch overhead. - fn new_iter_token(&self) -> u128; + /// To the more general question of whether 64 bits will be enough for any possible future node + /// structure, currently MAX_NODE_KEY_BYTES is limited to 48, but there is no limit on the branching + /// factor within that node. So even 128 bits would be insufficient to encode all paths in theory. + /// However a fixed-size node structure has a physical limit on its complexity. If we assume we will + /// limit a node to 4KB, 12 bits is enough to address any byte within that physical structure, so there + /// is probably some clever encoding that can address any position that it could contain, using 64 bits, + /// with a reasonable time and memory-fetch overhead. + fn new_iter_token(&self) -> IterToken; /// Generates an iter token that can be passed to [Self::next_items] to continue iteration from the /// specified path - fn iter_token_for_path(&self, key: &[u8]) -> u128; + fn iter_token_for_path(&self, key: &[u8]) -> IterToken; /// Steps to the next existing path within the node, in a depth-first order /// @@ -216,7 +212,7 @@ pub(crate) trait TrieNode: TrieNodeDowncas /// - `path` is relative to the start of `node` /// - `child_node` an onward node link, of `None` /// - `value` that exists at the path, or `None` - fn next_items(&self, token: u128) -> (u128, &[u8], Option<&TrieNodeODRc>, Option<&V>); + fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc>, Option<&V>); /// Returns the total number of leaves contained within the whole subtree defined by the node /// GOAT, this should be deprecated @@ -368,11 +364,14 @@ pub trait TrieNodeDowncast { fn convert_to_cell_node(&mut self) -> TrieNodeODRc; } +/// Node-local cursor used by the trie-node iteration interface +pub type IterToken = u64; + /// Special sentinel token value indicating iteration of a node has not been initialized -pub const NODE_ITER_INVALID: u128 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF; +pub const NODE_ITER_INVALID: IterToken = IterToken::MAX; /// Special sentinel token value indicating iteration of a node has concluded -pub const NODE_ITER_FINISHED: u128 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE; +pub const NODE_ITER_FINISHED: IterToken = IterToken::MAX - 1; /// Internal. A pointer to an onward link or a value contained within a node pub(crate) enum PayloadRef<'a, V: Clone + Send + Sync, A: Allocator> { @@ -1206,7 +1205,7 @@ mod tagged_node_ref { } #[inline(always)] - pub fn new_iter_token(&self) -> u128 { + pub fn new_iter_token(&self) -> IterToken { match self { Self::DenseByteNode(node) => node.new_iter_token(), Self::LineListNode(node) => node.new_iter_token(), @@ -1218,7 +1217,7 @@ mod tagged_node_ref { } } #[inline(always)] - pub fn iter_token_for_path(&self, key: &[u8]) -> u128 { + pub fn iter_token_for_path(&self, key: &[u8]) -> IterToken { match self { Self::DenseByteNode(node) => node.iter_token_for_path(key), Self::LineListNode(node) => node.iter_token_for_path(key), @@ -1230,7 +1229,7 @@ mod tagged_node_ref { } } #[inline(always)] - pub fn next_items(&self, token: u128) -> (u128, &'a[u8], Option<&'a TrieNodeODRc>, Option<&'a V>) { + pub fn next_items(&self, token: IterToken) -> (IterToken, &'a[u8], Option<&'a TrieNodeODRc>, Option<&'a V>) { match self { Self::DenseByteNode(node) => node.next_items(token), Self::LineListNode(node) => node.next_items(token), @@ -1821,7 +1820,7 @@ mod tagged_node_ref { } #[inline] - pub fn new_iter_token(&self) -> u128 { + pub fn new_iter_token(&self) -> IterToken { let (ptr, tag) = self.ptr.get_raw_parts(); match tag { EMPTY_NODE_TAG => 0, @@ -1833,7 +1832,7 @@ mod tagged_node_ref { } } #[inline] - pub fn iter_token_for_path(&self, key: &[u8]) -> u128 { + pub fn iter_token_for_path(&self, key: &[u8]) -> IterToken { let (ptr, tag) = self.ptr.get_raw_parts(); match tag { EMPTY_NODE_TAG => 0, @@ -1845,7 +1844,7 @@ mod tagged_node_ref { } } #[inline] - pub fn next_items(&self, token: u128) -> (u128, &[u8], Option<&'a TrieNodeODRc>, Option<&'a V>) { + pub fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&'a TrieNodeODRc>, Option<&'a V>) { let (ptr, tag) = self.ptr.get_raw_parts(); match tag { EMPTY_NODE_TAG => (NODE_ITER_FINISHED, &[], None, None), diff --git a/src/utils/mod.rs b/src/utils/mod.rs index d8b66e7..c100e1f 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -134,6 +134,13 @@ impl ByteMask { self.0 } /// Create an iterator over every byte, in ascending order + /// + /// DEVELOPER NOTE: This iterator owns a copy of the 256-bit mask and clears bits as it advances. + /// A cursor design that borrows the `ByteMask` is possible, reducing iterator state from the 32-byte mask plus word + /// cursor down to a mask reference and the cursor padded to a word. + /// However, because the borrowed cursor cannot clear the source mask, each `next` call has to rebuild + /// a shifted word mask to hide already-visited bits. Benchmarks showed that fixed per-item cost more + /// than doubled iteration overhead when the current owned iterator fits in registers. #[inline] pub fn iter(&self) -> ByteMaskIter { ByteMaskIter::from(self.0) diff --git a/src/zipper.rs b/src/zipper.rs index 9d60215..2a77f9d 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -1283,7 +1283,6 @@ pub(crate) const EXPECTED_DEPTH: usize = 16; pub(crate) const EXPECTED_PATH_LEN: usize = 64; pub(crate) mod read_zipper_core { - use crate::trie_node::*; use crate::PathMap; use crate::zipper::*; @@ -1311,12 +1310,12 @@ pub(crate) mod read_zipper_core { focus_node: MiriWrapper>, /// An iter token corresponding to the location of the `node_key` within the `focus_node`, or NODE_ITER_INVALID /// if iteration is not in-process - focus_iter_token: u128, + focus_iter_token: IterToken, /// Stores the entire path from the root node, including the bytes from `root_key` prefix_buf: Vec, /// Stores a stack of parent node references. Does not include the focus_node /// The tuple contains: `(node_ref, iter_token, key_offset_in_prefix_buf)` - ancestors: Vec<(TaggedNodeRef<'a, V, A>, u128, usize)>, + ancestors: Vec<(TaggedNodeRef<'a, V, A>, IterToken, usize)>, pub(crate) alloc: A, }