Use u64 trie iteration tokens

src/bridge_node.rs

@@ -398,11 +398,11 @@ impl<V: Clone + Send + Sync> TrieNode<V> for BridgeNode<V> {
         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<V: Clone + Send + Sync> TrieNode<V> for BridgeNode<V> {
         (NODE_ITER_FINISHED, &[])
     }
     #[inline(always)]
-    fn next_items(&self, token: u128) -> (u128, &[u8], Option<&TrieNodeODRc<V>>, Option<&V>) {
+    fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc<V>>, Option<&V>) {
         if token == 0 {
             let node_key = self.key();
             if self.is_used_child() {

src/dense_byte_node.rs

@@ -303,6 +303,28 @@ impl<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> ByteNode<Cf, A>
         unsafe{ self.values.get_unchecked_mut(ix) }
     }
 
+    #[inline(always)]
+    fn next_iter_byte_from(&self, start: IterToken) -> Option<u8> {
+        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<V: Clone + Send + Sync, A: Allocator, Cf: CoFree<V=V, A=A>> TrieNode<V, A>
         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<V, A>>, 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<V, A>>, 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<u64, usize>) -> 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]);
+}

src/empty_node.rs

@@ -58,13 +58,13 @@ impl<V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> 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<V, A>>, Option<&V>) {
+    fn next_items(&self, _token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc<V, A>>, Option<&V>) {
         (NODE_ITER_FINISHED, &[], None, None)
     }
     fn node_val_count(&self, _cache: &mut HashMap<u64, usize>) -> usize {

src/line_list_node.rs

@@ -1893,7 +1893,7 @@ impl<V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> for LineListNode<V, A>
     // *   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<V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> for LineListNode<V, A>
     /// - == 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<V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> for LineListNode<V, A>
         NODE_ITER_FINISHED
     }
     #[inline(always)]
-    fn next_items(&self, token: u128) -> (u128, &[u8], Option<&TrieNodeODRc<V, A>>, Option<&V>) {
+    fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc<V, A>>, Option<&V>) {
         match token {
             0 => {
                 if !self.is_used::<0>() {

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<u8>,
-    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> {

src/tiny_node.rs

@@ -214,9 +214,9 @@ impl<'a, V: Clone + Send + Sync, A: Allocator> TrieNode<V, A> 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<V, A>>, 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<V, A>>, Option<&V>) { unreachable!() }
     fn node_val_count(&self, cache: &mut HashMap<u64, usize>) -> usize {
         let temp_node = self.into_full().unwrap();
         temp_node.node_val_count(cache)

src/trie_node.rs

@@ -185,28 +185,16 @@ pub(crate) trait TrieNode<V: Clone + Send + Sync, A: Allocator>: 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<V: Clone + Send + Sync, A: Allocator>: 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<V, A>>, Option<&V>);
+    fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&TrieNodeODRc<V, A>>, 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<V: Clone + Send + Sync, A: Allocator> {
     fn convert_to_cell_node(&mut self) -> TrieNodeODRc<V, A>;
 }
 
+/// 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<V, A>>, Option<&'a V>) {
+        pub fn next_items(&self, token: IterToken) -> (IterToken, &'a[u8], Option<&'a TrieNodeODRc<V, A>>, 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<V, A>>, Option<&'a V>) {
+        pub fn next_items(&self, token: IterToken) -> (IterToken, &[u8], Option<&'a TrieNodeODRc<V, A>>, Option<&'a V>) {
             let (ptr, tag) = self.ptr.get_raw_parts();
             match tag {
                 EMPTY_NODE_TAG => (NODE_ITER_FINISHED, &[], None, None),

src/zipper.rs

@@ -1311,12 +1311,12 @@ pub(crate) mod read_zipper_core {
         focus_node: MiriWrapper<TaggedNodeRef<'a, V, A>>,
         /// 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<u8>,
         /// 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,
     }