From 44a31abf9a7e81a98592f842da4e06f45d9c94f9 Mon Sep 17 00:00:00 2001 From: MesTTo Date: Sun, 5 Jul 2026 04:35:21 +1000 Subject: [PATCH] Use u64 trie iteration tokens Implements the answer the iter-token note in the TrieNode trait already sketched: the token is a node-local cursor, so it never needs to encode a whole path, and 64 bits are enough. The ByteNode stops caching the remaining mask word inside the token (the source of the 128-bit requirement); it keeps only one more than the last byte returned and recomputes the next set bit from its own mask (next_iter_byte_from, four masked trailing_zeros probes, still O(1) per step). All other node types already used small counters. iter_token_for_path semantics are unchanged (resume strictly after the given byte; the 63/127/191/255 word-edge cases land identically through the word advance), pinned by a new mask-word-boundary test. Consumers treat tokens as opaque, so the swap is type-level for them; ReadZipperCore shrinks from 144 to 128 bytes (measured on both trees), and each ancestor stack entry loses 8 bytes of padding. --- src/bridge_node.rs | 6 +-- src/dense_byte_node.rs | 101 +++++++++++++++++++++++++++-------------- src/empty_node.rs | 6 +-- src/line_list_node.rs | 6 +-- src/old_cursor.rs | 2 +- src/tiny_node.rs | 6 +-- src/trie_node.rs | 47 ++++++++----------- src/zipper.rs | 4 +- 8 files changed, 101 insertions(+), 77 deletions(-) diff --git a/src/bridge_node.rs b/src/bridge_node.rs index 56135a68..15ca5343 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 ca253c30..dbae42d6 100644 --- a/src/dense_byte_node.rs +++ b/src/dense_byte_node.rs @@ -303,6 +303,28 @@ impl> ByteNode unsafe{ self.values.get_unchecked_mut(ix) } } + #[inline(always)] + fn next_iter_byte_from(&self, start: IterToken) -> Option { + 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); + } + + 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 +945,27 @@ 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 { + 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) + unsafe { *key.get_unchecked(0) as IterToken + 1 } } } #[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) = self.next_iter_byte_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 = k as IterToken + 1; + let cf = unsafe { self.get_unchecked(k) }; + 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 +2378,36 @@ 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]); +} diff --git a/src/empty_node.rs b/src/empty_node.rs index 98159eac..722952a3 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 92c3133e..bfc650b5 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 bbacc076..d1c9f360 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 08d61fc3..5cdb3573 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 82a8b165..59224591 100644 --- a/src/trie_node.rs +++ b/src/trie_node.rs @@ -185,28 +185,16 @@ pub(crate) trait TrieNode: TrieNodeDowncas /// Returns `true` if the node contains no children nor values, otherwise false fn node_is_empty(&self) -> bool; - /// 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. - /// - /// 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; + /// Generates a new iter token, to iterate the children and values contained within this node. + /// The token is a node-local cursor. It only needs to represent the next position to inspect within + /// the current node, not an arbitrary path through the trie, so 64 bits are enough: the ByteNode keeps + /// one more than the last path byte returned (or 0 at the start) and recomputes the remaining mask from + /// the node itself, rather than caching the mask word inside the token. + 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 +204,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 +356,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 +1197,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 +1209,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 +1221,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 +1812,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 +1824,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 +1836,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/zipper.rs b/src/zipper.rs index 9d602151..8ca9628b 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -1311,12 +1311,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, }