tiny_trie/
fixed_len_nibble_trie.rs

1//! Fixed-Length Nibble Trie — a fixed-fanout radix trie with fixed-length key slots.
2//!
3//! Like [`NibbleTrie`], each node has 16 child slots (one per nibble value 0–15).
4//! The key difference is that keys are stored in fixed-length slots in `buf`,
5//! eliminating the need for a separate index vector and the `offset` field per node.
6//!
7//! # Key Storage
8//!
9//! Each key occupies exactly `max_len` bytes in `buf`, zero-padded on the right.
10//! The key index `i` maps directly to `buf[i * max_len .. (i + 1) * max_len]`.
11//! Actual key lengths are stored in a `lens: Vec<u16>` alongside `buf`, giving
12//! O(1) key retrieval without an index vector. This adds only 2 bytes/key overhead
13//! (vs NibbleTrie's ~10 bytes/key for `(usize, LEN)` on 64-bit).
14//!
15//! # Key Index
16//!
17//! Leaf key indices are 0-based (unlike NibbleTrie's 1-based scheme). The sentinel
18//! value for empty child slots is `PTR::max_value()` (instead of 0). This gives
19//! a max entry count of `PTR::max_value() - 1`.
20//!
21//! # Offset Elimination
22//!
23//! The `offset` field in NibbleTrie's `Node` is replaced by computing
24//! `leaf.as_usize() * max_len` on demand. The terminal flag is stored in a
25//! `flags: u8` field (bit 0).
26//!
27//! # Optimization
28//!
29//! [`optimize()`] rewrites `buf` and `values` in DFS-sorted order and remaps leaf
30//! indices. Called automatically after each insert when `values.len()` is a power
31//! of two (amortized O(1) per insert).
32
33use std::simd::{Simd, cmp::SimdPartialEq};
34
35// ---------------------------------------------------------------------------
36// Constants
37// ---------------------------------------------------------------------------
38
39/// Bit 0 of `FixedLenNode.flags` marks a terminal node (key ends here).
40const TERMINAL_FLAG: u8 = 1;
41
42// ---------------------------------------------------------------------------
43// FixedLenNode
44// ---------------------------------------------------------------------------
45
46/// A single node in the fixed-length nibble trie arena.
47///
48/// Generic over `PTR` (pointer/index type for children and arena references).
49/// Unlike `Node`, there is no `LEN` parameter (`prefix_len` is always `u16`)
50/// and no `offset` field (computed from `leaf * max_len`).
51///
52/// Layout with PTR=u16: 40 bytes (saves 8 vs NibbleTrie's 48).
53/// Layout with PTR=u32: 76 bytes (saves 4 vs NibbleTrie's 80).
54#[derive(Copy, Clone)]
55pub struct FixedLenNode<PTR: TrieIndex> {
56    pub children: [PTR; 16],
57    pub prefix_len: u16,
58    pub leaf_mask: u16,
59    pub leaf: PTR,
60    pub flags: u8,
61    // compiler pads to align(PTR)
62}
63
64impl<PTR: TrieIndex> FixedLenNode<PTR> {
65    fn new() -> Self {
66        FixedLenNode {
67            children: [PTR::max_value_sentinel(); 16],
68            prefix_len: 0,
69            leaf_mask: 0,
70            leaf: PTR::max_value_sentinel(),
71            flags: 0,
72        }
73    }
74
75    #[inline]
76    pub fn is_terminal(&self) -> bool {
77        self.flags & TERMINAL_FLAG != 0
78    }
79
80    #[inline]
81    fn set_terminal(&mut self, val: bool) {
82        if val {
83            self.flags |= TERMINAL_FLAG;
84        } else {
85            self.flags &= !TERMINAL_FLAG;
86        }
87    }
88
89    #[inline]
90    pub fn is_leaf(&self, nib: usize) -> bool {
91        debug_assert!(nib < 16);
92        (self.leaf_mask >> nib) & 1 == 1
93    }
94
95    #[inline]
96    fn set_leaf(&mut self, nib: usize) {
97        debug_assert!(nib < 16);
98        self.leaf_mask |= 1 << nib;
99    }
100
101    #[inline]
102    fn clear_leaf(&mut self, nib: usize) {
103        debug_assert!(nib < 16);
104        self.leaf_mask &= !(1 << nib);
105    }
106
107    /// Store a leaf key index at `nib`. Key index must not equal the sentinel.
108    #[inline]
109    fn set_leaf_child(&mut self, nib: usize, ki: PTR) {
110        debug_assert!(nib < 16);
111        debug_assert!(ki != PTR::max_value_sentinel(), "key index max_value is sentinel");
112        self.set_leaf(nib);
113        self.children[nib] = ki;
114    }
115
116    /// Store an arena index at `nib` (internal node reference).
117    /// Arena index must not equal the sentinel.
118    #[inline]
119    fn set_internal_child(&mut self, nib: usize, arena_idx: PTR) {
120        debug_assert!(nib < 16);
121        debug_assert!(arena_idx != PTR::max_value_sentinel(), "arena index max_value is sentinel");
122        self.clear_leaf(nib);
123        self.children[nib] = arena_idx;
124    }
125
126    /// Decode a leaf child at `nib` into a key index.
127    /// Returns `None` if the slot is empty or not a leaf.
128    #[inline]
129    fn leaf_key_index(&self, nib: usize) -> Option<PTR> {
130        debug_assert!(nib < 16);
131        if self.is_leaf(nib) && self.children[nib] != PTR::max_value_sentinel() {
132            Some(self.children[nib])
133        } else {
134            None
135        }
136    }
137
138    /// Compute a 16-bit mask where bit N is set if `children[N]` is occupied
139    /// (either leaf or internal — not the empty sentinel).
140    #[inline]
141    pub fn children_mask(&self) -> u16 {
142        // Invert children, then check for nonzero: XOR with max_value turns
143        // sentinel into 0 and real values into nonzero.
144        let mut mask = 0u16;
145        for i in 0..16 {
146            if self.children[i] != PTR::max_value_sentinel() {
147                mask |= 1 << i;
148            }
149        }
150        mask
151    }
152}
153
154impl<PTR: TrieIndex> std::fmt::Debug for FixedLenNode<PTR> {
155    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156        let active: Vec<(usize, &str, PTR)> = (0..16)
157            .filter(|&n| self.children[n] != PTR::max_value_sentinel())
158            .map(|n| {
159                let tag = if self.is_leaf(n) { "L" } else { "I" };
160                (n, tag, self.children[n])
161            })
162            .collect();
163        f.debug_struct("FixedLenNode")
164            .field("prefix_len", &self.prefix_len)
165            .field("leaf_mask", &format_args!("0x{:04x}", self.leaf_mask))
166            .field("leaf", &self.leaf)
167            .field("terminal", &self.is_terminal())
168            .field("children", &active)
169            .finish()
170    }
171}
172
173// ---------------------------------------------------------------------------
174// TrieIndex extension
175// ---------------------------------------------------------------------------
176
177/// TrieIndex now provides `max_value_sentinel()` directly.
178use crate::nibble_trie::TrieIndex;
179
180// ---------------------------------------------------------------------------
181// Nibble helpers (reuse from nibble_trie)
182// ---------------------------------------------------------------------------
183
184#[inline]
185fn key_nibble_at(key: &[u8], idx: usize) -> u8 {
186    let byte_idx = idx / 2;
187    if byte_idx < key.len() {
188        if idx % 2 == 0 {
189            key[byte_idx] >> 4
190        } else {
191            key[byte_idx] & 0x0F
192        }
193    } else {
194        0
195    }
196}
197
198#[inline]
199fn nibble_count(key: &[u8]) -> usize {
200    key.len() * 2
201}
202
203#[inline]
204fn diverging_nibble(xor: u8, byte_idx: usize) -> usize {
205    byte_idx * 2 + ((xor >> 4 == 0) as usize)
206}
207
208// ---------------------------------------------------------------------------
209// Divergence result
210// ---------------------------------------------------------------------------
211
212enum DivergeResult {
213    Duplicate,
214    At(usize),
215}
216
217#[inline]
218fn find_divergence(key_a: &[u8], key_b: &[u8], from: usize) -> DivergeResult {
219    let total_a = nibble_count(key_a);
220    let total_b = nibble_count(key_b);
221    let min = total_a.min(total_b);
222    let mut d = from;
223    while d < min {
224        if key_nibble_at(key_a, d) != key_nibble_at(key_b, d) {
225            return DivergeResult::At(d);
226        }
227        d += 1;
228    }
229    if total_a == total_b {
230        DivergeResult::Duplicate
231    } else {
232        DivergeResult::At(d)
233    }
234}
235
236fn simd_find_divergence<const N: usize>(key_a: &[u8], key_b: &[u8], from: usize) -> DivergeResult
237{
238    let minlen = key_a.len().min(key_b.len());
239    let mut i = from / 2;
240
241    while i + N <= minlen {
242        let a = Simd::<u8, N>::from_slice(unsafe { key_a.get_unchecked(i..i + N) });
243        let b = Simd::<u8, N>::from_slice(unsafe { key_b.get_unchecked(i..i + N) });
244        let mask = a.simd_ne(b);
245        if mask.any() {
246            let diff_byte_idx = i + mask.first_set().unwrap();
247            let xor = unsafe { *key_a.get_unchecked(diff_byte_idx) ^ *key_b.get_unchecked(diff_byte_idx) };
248            return DivergeResult::At(diverging_nibble(xor, diff_byte_idx));
249        }
250        i += N;
251    }
252
253    find_divergence(key_a, key_b, i * 2)
254}
255
256// ---------------------------------------------------------------------------
257// FixedLenNibbleTrie
258// ---------------------------------------------------------------------------
259
260/// A fixed-length nibble trie map.
261///
262/// Keys are stored in fixed-size slots of `max_len` bytes, zero-padded on the
263/// right. Key length is recovered by scanning backward from the slot boundary
264/// for the first nonzero byte. **Keys must not contain trailing zero bytes** —
265/// a key like `b"a\0"` would be retrieved as `b"a"`.
266///
267/// The `PTR` type parameter controls the width of arena/key indices and thus
268/// the maximum number of entries (~`PTR::max_value() - 1`).
269#[derive(Clone)]
270pub struct FixedLenNibbleTrie<T, PTR: TrieIndex = u32> {
271    pub arena: Vec<FixedLenNode<PTR>>,
272    pub buf: Vec<u8>,
273    pub values: Vec<T>,
274    pub lens: Vec<u16>,    // actual byte length of each key
275    pub max_len: usize,
276}
277
278impl<T, PTR: TrieIndex> FixedLenNibbleTrie<T, PTR> {
279    // -------------------------------------------------------------------
280    // Construction & basic accessors
281    // -------------------------------------------------------------------
282
283    /// Create a new empty trie with the given maximum key length.
284    ///
285    /// Keys longer than `max_len` will be rejected by `insert`.
286    pub fn new(max_len: usize) -> Self {
287        FixedLenNibbleTrie {
288            arena: Vec::new(),
289            buf: Vec::new(),
290            values: Vec::new(),
291            lens: Vec::new(),
292            max_len,
293        }
294    }
295
296    pub fn len(&self) -> usize {
297        self.values.len()
298    }
299
300    pub fn is_empty(&self) -> bool {
301        self.values.is_empty()
302    }
303
304    // -------------------------------------------------------------------
305    // Key retrieval
306    // -------------------------------------------------------------------
307
308    /// Return the actual (unpadded) key slice for key index `ki`.
309    /// Uses the stored length for O(1) retrieval.
310    #[inline]
311    pub fn key_slice(&self, ki: PTR) -> &[u8] {
312        let idx = ki.as_usize();
313        let start = idx * self.max_len;
314        let len = self.lens[idx] as usize;
315        &self.buf[start..start + len]
316    }
317
318    /// Check whether the key at index `ki` matches `key`.
319    /// Compares the stored key length first, then the bytes.
320    #[inline]
321    fn key_matches(&self, ki: PTR, key: &[u8]) -> bool {
322        let idx = ki.as_usize();
323        let len = self.lens[idx] as usize;
324        if len != key.len() {
325            return false;
326        }
327        let start = idx * self.max_len;
328        self.buf[start..start + len] == *key
329    }
330
331    // -------------------------------------------------------------------
332    // Lookup
333    // -------------------------------------------------------------------
334
335    pub fn get_index(&self, key: &[u8]) -> Option<usize> {
336        if key.len() > self.max_len || self.arena.is_empty() {
337            return None;
338        }
339        let mut node_idx: PTR = PTR::zero();
340        let max_nib = key.len() * 2;
341        loop {
342            let node = &self.arena[node_idx.as_usize()];
343            if node.prefix_len as usize >= max_nib {
344                if node.is_terminal() {
345                    if self.key_matches(node.leaf, key) {
346                        return Some(node.leaf.as_usize());
347                    }
348                }
349                return None;
350            }
351            let nib = key_nibble_at(key, node.prefix_len as usize) as usize;
352            let slot = node.children[nib];
353            if slot == PTR::max_value_sentinel() {
354                return None;
355            }
356            if node.is_leaf(nib) {
357                let key_index = slot;
358                return if self.key_matches(key_index, key) {
359                    Some(key_index.as_usize())
360                } else {
361                    None
362                };
363            }
364            node_idx = slot;
365        }
366    }
367
368    /// Unchecked lookup — assumes the key is present in the trie.
369    ///
370    /// # Safety
371    ///
372    /// The key **must** have been inserted into this trie. If the key is not
373    /// present, the result is unspecified.
374    #[cfg(feature = "unchecked")]
375    unsafe fn get_index_unchecked(&self, key: &[u8]) -> Option<usize> {
376        if self.arena.is_empty() {
377            return None;
378        }
379        let mut node_idx: PTR = PTR::zero();
380        let max_nib = key.len() * 2;
381        loop {
382            let node = unsafe { self.arena.get_unchecked(node_idx.as_usize()) };
383            let prefix_len = node.prefix_len as usize;
384            if prefix_len >= max_nib {
385                debug_assert!(node.is_terminal(), "get_unchecked: key not in set");
386                return Some(node.leaf.as_usize());
387            }
388            let nib = key_nibble_at(key, prefix_len) as usize;
389            let slot = unsafe { *node.children.get_unchecked(nib) };
390            if slot == PTR::max_value_sentinel() {
391                return None;
392            }
393            if node.is_leaf(nib) {
394                return Some(slot.as_usize());
395            }
396            node_idx = slot;
397        }
398    }
399
400    pub fn get(&self, key: &[u8]) -> Option<&T> {
401        self.get_index(key).map(|ki| &self.values[ki])
402    }
403
404    pub fn get_mut(&mut self, key: &[u8]) -> Option<&mut T> {
405        self.get_index(key).map(|ki| &mut self.values[ki])
406    }
407
408    /// Unchecked value lookup — assumes the key is present in the trie.
409    ///
410    /// # Safety
411    ///
412    /// The key **must** have been inserted into this trie. If the key is not
413    /// present, the result is unspecified.
414    #[cfg(feature = "unchecked")]
415    pub unsafe fn get_unchecked(&self, key: &[u8]) -> Option<&T> {
416        unsafe { self.get_index_unchecked(key).map(|ki| &self.values[ki]) }
417    }
418
419    // -------------------------------------------------------------------
420    // Insertion
421    // -------------------------------------------------------------------
422
423    pub fn insert(&mut self, key: Vec<u8>, value: T) -> Result<usize, ()> {
424        if key.len() > self.max_len {
425            return Err(());
426        }
427        // Overflow: arena/key indices must fit in PTR. max_value() is sentinel.
428        if self.arena.len() >= PTR::max_value() {
429            return Err(());
430        }
431        if self.values.len() + 1 >= PTR::max_value() {
432            return Err(());
433        }
434
435        // Allocate slot
436        let ki = self.values.len();
437        let start = self.buf.len();
438        self.buf.resize(start + self.max_len, 0);
439        self.buf[start..start + key.len()].copy_from_slice(&key);
440        self.values.push(value);
441        self.lens.push(key.len() as u16);
442        let new_ki = PTR::from_usize(ki);
443        let max_nib = key.len() * 2;
444
445        if self.arena.is_empty() {
446            if max_nib == 0 {
447                // Empty key — root is terminal
448                let mut root = FixedLenNode::new();
449                root.set_terminal(true);
450                root.leaf = new_ki;
451                self.arena.push(root);
452                return Ok(new_ki.as_usize());
453            }
454            let first_nib = key_nibble_at(&key, 0) as usize;
455            let mut root = FixedLenNode::new();
456            root.set_leaf_child(first_nib, new_ki);
457            root.leaf = new_ki;
458            self.arena.push(root);
459            return Ok(new_ki.as_usize());
460        }
461
462        let mut node_idx: PTR = PTR::zero();
463        let mut confirmed: usize = 0;
464
465        loop {
466            let node = &self.arena[node_idx.as_usize()];
467            let ref_key = self.key_slice(node.leaf);
468            let prefix_len = node.prefix_len as usize;
469
470            match simd_find_divergence::<8>(&key, ref_key, confirmed) {
471                DivergeResult::Duplicate => {
472                    // Roll back allocation
473                    self.buf.truncate(start);
474                    self.values.pop();
475                    self.lens.pop();
476                    return Err(());
477                }
478                DivergeResult::At(diverge) if diverge < prefix_len => {
479                    // Divergence before discriminating nibble — split this node
480                    let new_nib = key_nibble_at(&key, diverge) as usize;
481                    let ref_nib = key_nibble_at(ref_key, diverge) as usize;
482
483                    let mut new_parent = FixedLenNode::new();
484                    new_parent.prefix_len = diverge as u16;
485
486                    if diverge >= max_nib {
487                        // New key ends at the split point — terminal
488                        new_parent.set_terminal(true);
489                        new_parent.leaf = new_ki;
490                    } else {
491                        new_parent.set_leaf_child(new_nib, new_ki);
492                        new_parent.leaf = new_ki;
493                    }
494
495                    let old_node = std::mem::replace(
496                        &mut self.arena[node_idx.as_usize()],
497                        new_parent,
498                    );
499                    let old_idx = PTR::from_usize(self.arena.len());
500                    self.arena.push(old_node);
501
502                    self.arena[node_idx.as_usize()].set_internal_child(ref_nib, old_idx);
503                    self.sort_internal_children(node_idx);
504
505                    return Ok(new_ki.as_usize());
506                }
507                DivergeResult::At(_) => {
508                    // Divergence at or after prefix_len — follow the child.
509                    if max_nib <= prefix_len {
510                        // Key exhausted at this node — mark terminal
511                        self.arena[node_idx.as_usize()].set_terminal(true);
512                        self.arena[node_idx.as_usize()].leaf = new_ki;
513                        return Ok(new_ki.as_usize());
514                    }
515
516                    confirmed = prefix_len + 1;
517                    let nib = key_nibble_at(&key, prefix_len) as usize;
518                    let slot = node.children[nib];
519
520                    if slot == PTR::max_value_sentinel() {
521                        // Empty slot — new key diverges here
522                        self.arena[node_idx.as_usize()].set_leaf_child(nib, new_ki);
523                        return Ok(new_ki.as_usize());
524                    }
525
526                    if node.is_leaf(nib) {
527                        let existing_ki = slot;
528                        let existing_key = self.key_slice(existing_ki);
529
530                        match simd_find_divergence::<8>(&key, existing_key, confirmed) {
531                            DivergeResult::Duplicate => {
532                                // Roll back
533                                self.buf.truncate(start);
534                                self.values.pop();
535                                self.lens.pop();
536                                return Err(());
537                            }
538                            DivergeResult::At(d) => {
539                                let mut split_node = FixedLenNode::new();
540                                split_node.prefix_len = d as u16;
541
542                                if d >= max_nib {
543                                    // New key ends at split — terminal
544                                    let exist_nib = key_nibble_at(existing_key, d) as usize;
545                                    split_node.set_terminal(true);
546                                    split_node.leaf = new_ki;
547                                    split_node.set_leaf_child(exist_nib, existing_ki);
548                                } else if d >= existing_key.len() * 2 {
549                                    // Existing key ends at split — terminal
550                                    let new_nib = key_nibble_at(&key, d) as usize;
551                                    split_node.set_terminal(true);
552                                    split_node.leaf = existing_ki;
553                                    split_node.set_leaf_child(new_nib, new_ki);
554                                } else {
555                                    // Neither key ends — they diverge
556                                    let new_nib = key_nibble_at(&key, d) as usize;
557                                    let exist_nib = key_nibble_at(existing_key, d) as usize;
558                                    debug_assert_ne!(new_nib, exist_nib);
559                                    split_node.set_leaf_child(new_nib, new_ki);
560                                    split_node.set_leaf_child(exist_nib, existing_ki);
561                                    split_node.leaf = existing_ki;
562                                }
563
564                                let split_idx = PTR::from_usize(self.arena.len());
565                                self.arena.push(split_node);
566                                self.arena[node_idx.as_usize()].set_internal_child(nib, split_idx);
567                                self.sort_internal_children(node_idx);
568
569                                return Ok(new_ki.as_usize());
570                            }
571                        }
572                    }
573
574                    // Internal node — descend
575                    node_idx = slot;
576                }
577            }
578        }
579    }
580
581    // -------------------------------------------------------------------
582    // Optimization (DFS key-sorted buf + values rewrite)
583    // -------------------------------------------------------------------
584
585    /// Rewrite `buf` and `values` so that keys appear in sorted order, with
586    /// contiguous layout for sequential access during iteration.
587    ///
588    /// After `optimize()`, a forward iteration visits keys in ascending memory
589    /// order within `buf`. Leaf indices in arena nodes are remapped to reflect
590    /// the new positions.
591    ///
592    /// No-op for empty tries.
593    pub fn optimize(&mut self) {
594        if self.arena.is_empty() {
595            return;
596        }
597
598        let n = self.values.len();
599        let mut new_buf = vec![0u8; n * self.max_len];
600        let mut new_values = Vec::with_capacity(n);
601        let mut new_lens = Vec::with_capacity(n);
602        // Safety: we'll fill all n slots, so this is safe
603        unsafe { new_values.set_len(n); }
604        unsafe { new_lens.set_len(n); }
605
606        let mut cursor: usize = 0;
607        let mut remap: Vec<usize> = vec![0; n]; // old ki -> new ki
608
609        self.walk_optimize(PTR::zero(), &mut new_buf, &mut new_values, &mut new_lens, &mut remap, &mut cursor);
610
611        // Remap all leaf indices in arena nodes
612        for node in &mut self.arena {
613            if node.leaf != PTR::max_value_sentinel() {
614                node.leaf = PTR::from_usize(remap[node.leaf.as_usize()]);
615            }
616            for nib in 0..16 {
617                if node.is_leaf(nib) && node.children[nib] != PTR::max_value_sentinel() {
618                    node.children[nib] = PTR::from_usize(remap[node.children[nib].as_usize()]);
619                }
620            }
621        }
622
623        self.buf = new_buf;
624        self.values = new_values;
625        self.lens = new_lens;
626    }
627
628    fn walk_optimize(
629        &mut self,
630        node_idx: PTR,
631        new_buf: &mut [u8],
632        new_values: &mut [T],
633        new_lens: &mut [u16],
634        remap: &mut [usize],
635        cursor: &mut usize,
636    ) {
637        let node = self.arena[node_idx.as_usize()]; // copy to avoid borrow conflicts
638
639        if node.is_terminal() {
640            let old_ki = node.leaf.as_usize();
641            let new_ki = *cursor;
642            let old_start = old_ki * self.max_len;
643            let new_start = new_ki * self.max_len;
644            new_buf[new_start..new_start + self.max_len]
645                .copy_from_slice(&self.buf[old_start..old_start + self.max_len]);
646            // SAFETY: we're writing into a set_len'd Vec; all slots will be written
647            unsafe {
648                std::ptr::write(new_values.as_mut_ptr().add(new_ki), std::ptr::read(self.values.as_ptr().add(old_ki)));
649            }
650            new_lens[new_ki] = self.lens[old_ki];
651            remap[old_ki] = new_ki;
652            *cursor += 1;
653        }
654
655        for nib in 0..16 {
656            if node.children[nib] == PTR::max_value_sentinel() {
657                continue;
658            }
659            if node.is_leaf(nib) {
660                let old_ki = node.children[nib].as_usize();
661                let new_ki = *cursor;
662                let old_start = old_ki * self.max_len;
663                let new_start = new_ki * self.max_len;
664                new_buf[new_start..new_start + self.max_len]
665                    .copy_from_slice(&self.buf[old_start..old_start + self.max_len]);
666                unsafe {
667                    std::ptr::write(new_values.as_mut_ptr().add(new_ki), std::ptr::read(self.values.as_ptr().add(old_ki)));
668                }
669                new_lens[new_ki] = self.lens[old_ki];
670                remap[old_ki] = new_ki;
671                *cursor += 1;
672            } else {
673                self.walk_optimize(node.children[nib], new_buf, new_values, new_lens, remap, cursor);
674            }
675        }
676    }
677
678    /// Optimize after insert if `values.len()` is a power of two.
679    fn maybe_optimize(&mut self) {
680        let n = self.values.len();
681        if n > 0 && n.is_power_of_two() {
682            self.optimize();
683        }
684    }
685
686    /// Insert and auto-optimize on power-of-two sizes.
687    pub fn insert_auto(&mut self, key: Vec<u8>, value: T) -> Result<usize, ()> {
688        let result = self.insert(key, value)?;
689        self.maybe_optimize();
690        Ok(result)
691    }
692
693    // -------------------------------------------------------------------
694    // Iteration
695    // -------------------------------------------------------------------
696
697    pub fn iter(&self) -> FixedLenIter<'_, T, PTR> {
698        FixedLenIter::new(self)
699    }
700
701    pub fn iter_last(&self) -> FixedLenIter<'_, T, PTR> {
702        FixedLenIter::new_last(self)
703    }
704
705    pub fn into_keys_values(self) -> (Vec<Vec<u8>>, Vec<T>) {
706        let keys: Vec<Vec<u8>> = (0..self.values.len())
707            .map(|i| self.key_slice(PTR::from_usize(i)).to_vec())
708            .collect();
709        (keys, self.values)
710    }
711
712    // -------------------------------------------------------------------
713    // Arena maintenance
714    // -------------------------------------------------------------------
715
716    fn swap_arena(&mut self, a: PTR, b: PTR) {
717        if a == b {
718            return;
719        }
720        self.arena.swap(a.as_usize(), b.as_usize());
721        for node in &mut self.arena {
722            for nib in 0..16 {
723                let child = node.children[nib];
724                if child != PTR::max_value_sentinel() && !node.is_leaf(nib) {
725                    if child == a {
726                        node.children[nib] = b;
727                    } else if child == b {
728                        node.children[nib] = a;
729                    }
730                }
731            }
732        }
733    }
734
735    fn sort_internal_children(&mut self, node_idx: PTR) {
736        let mut internals: [u8; 16] = [0; 16];
737        let mut arena_ids: [PTR; 16] = [PTR::max_value_sentinel(); 16];
738        let mut count = 0usize;
739        {
740            let node = &self.arena[node_idx.as_usize()];
741            for nib in 0u8..16 {
742                if node.children[nib as usize] != PTR::max_value_sentinel()
743                    && !node.is_leaf(nib as usize)
744                {
745                    internals[count] = nib;
746                    arena_ids[count] = node.children[nib as usize];
747                    count += 1;
748                }
749            }
750        }
751        if count <= 1 {
752            return;
753        }
754        let max_arena_idx = (0..count).fold(PTR::zero(), |m, i| {
755            if arena_ids[i].as_usize() > m.as_usize() { arena_ids[i] } else { m }
756        });
757        let insert_pos = (0..count).find(|&i| arena_ids[i] == max_arena_idx).unwrap();
758        for i in insert_pos..count - 1 {
759            self.swap_arena(arena_ids[i], arena_ids[i + 1]);
760            let tmp = arena_ids[i];
761            arena_ids[i] = arena_ids[i + 1];
762            arena_ids[i + 1] = tmp;
763        }
764    }
765
766    // -------------------------------------------------------------------
767    // Capacity
768    // -------------------------------------------------------------------
769
770    pub fn near_capacity(&self) -> bool {
771        self.arena.len() >= PTR::max_value() || self.values.len() + 1 >= PTR::max_value()
772    }
773}
774
775impl<T, PTR: TrieIndex> Default for FixedLenNibbleTrie<T, PTR> {
776    fn default() -> Self {
777        // Default max_len of 256 — reasonable for most string keys.
778        // Users should call new() with an appropriate max_len.
779        Self::new(256)
780    }
781}
782
783// ---------------------------------------------------------------------------
784// Iterator
785// ---------------------------------------------------------------------------
786
787const TERMINAL_NIB: usize = 16;
788
789pub struct FixedLenIter<'a, T, PTR: TrieIndex> {
790    trie: &'a FixedLenNibbleTrie<T, PTR>,
791    stack: Vec<(PTR, u16, usize)>,
792}
793
794impl<'a, T, PTR: TrieIndex> FixedLenIter<'a, T, PTR> {
795    fn new(trie: &'a FixedLenNibbleTrie<T, PTR>) -> Self {
796        if trie.arena.is_empty() {
797            return FixedLenIter { trie, stack: Vec::new() };
798        }
799        let mask = trie.arena[0].children_mask();
800        let nib = if trie.arena[0].is_terminal() { TERMINAL_NIB } else { usize::MAX };
801        FixedLenIter { trie, stack: vec![(PTR::zero(), mask, nib)] }
802    }
803
804    fn new_last(trie: &'a FixedLenNibbleTrie<T, PTR>) -> Self {
805        if trie.arena.is_empty() {
806            return FixedLenIter { trie, stack: Vec::new() };
807        }
808        let mut stack = Vec::new();
809        let mut idx: PTR = PTR::zero();
810        loop {
811            let node = &trie.arena[idx.as_usize()];
812            let mask = node.children_mask();
813            if mask != 0 {
814                let nib = 15 - mask.leading_zeros() as usize;
815                stack.push((idx, mask, nib));
816                if node.is_leaf(nib) {
817                    break;
818                } else {
819                    idx = node.children[nib];
820                }
821            } else if node.is_terminal() {
822                stack.push((idx, mask, TERMINAL_NIB));
823                break;
824            } else {
825                break;
826            }
827        }
828        FixedLenIter { trie, stack }
829    }
830
831    #[inline]
832    fn descend_first(&mut self, mut idx: PTR) {
833        loop {
834            let node = &self.trie.arena[idx.as_usize()];
835            if node.is_terminal() {
836                let mask = node.children_mask();
837                self.stack.push((idx, mask, TERMINAL_NIB));
838                return;
839            }
840            let mask = node.children_mask();
841            debug_assert!(mask != 0, "descend_first: non-terminal node with no children");
842            let nib = mask.trailing_zeros() as usize;
843            self.stack.push((idx, mask, nib));
844            if node.is_leaf(nib) {
845                return;
846            } else {
847                idx = node.children[nib];
848            }
849        }
850    }
851
852    #[inline]
853    fn descend_last(&mut self, mut idx: PTR) {
854        loop {
855            let node = &self.trie.arena[idx.as_usize()];
856            if node.is_terminal() {
857                let mask = node.children_mask();
858                if mask == 0 {
859                    self.stack.push((idx, mask, TERMINAL_NIB));
860                    return;
861                }
862            }
863            let mask = node.children_mask();
864            if mask == 0 {
865                if node.is_terminal() {
866                    self.stack.push((idx, mask, TERMINAL_NIB));
867                }
868                return;
869            }
870            let nib = 15 - mask.leading_zeros() as usize;
871            self.stack.push((idx, mask, nib));
872            if node.is_leaf(nib) {
873                return;
874            } else {
875                idx = node.children[nib];
876            }
877        }
878    }
879
880    #[inline]
881    fn push_next_child(&mut self, arena_idx: PTR, mask: u16, start_nib: usize) -> bool {
882        let shifted = if start_nib >= 16 { 0u16 } else { mask >> start_nib };
883        if shifted == 0 {
884            return false;
885        }
886        let nib = start_nib + shifted.trailing_zeros() as usize;
887        debug_assert!(nib < 16);
888        debug_assert!(mask & (1 << nib) != 0);
889        self.stack.push((arena_idx, mask, nib));
890        if !self.trie.arena[arena_idx.as_usize()].is_leaf(nib) {
891            self.descend_first(self.trie.arena[arena_idx.as_usize()].children[nib]);
892        }
893        true
894    }
895
896    #[inline]
897    fn backtrack_to_next(&mut self) -> Option<(&[u8], &T)> {
898        loop {
899            let (parent_idx, parent_mask, child_nib) = self.stack.pop()?;
900            if self.push_next_child(parent_idx, parent_mask, child_nib + 1) {
901                return self.current();
902            }
903        }
904    }
905
906    pub fn current(&self) -> Option<(&[u8], &T)> {
907        let &(arena_idx, _mask, nib) = self.stack.last()?;
908        if nib == usize::MAX {
909            return None;
910        }
911        let node = &self.trie.arena[arena_idx.as_usize()];
912        if nib == TERMINAL_NIB {
913            let key = self.trie.key_slice(node.leaf);
914            let value = &self.trie.values[node.leaf.as_usize()];
915            Some((key, value))
916        } else if let Some(key_index) = node.leaf_key_index(nib) {
917            let key = self.trie.key_slice(key_index);
918            let value = &self.trie.values[key_index.as_usize()];
919            Some((key, value))
920        } else {
921            None
922        }
923    }
924
925    pub fn current_index(&self) -> Option<usize> {
926        let &(arena_idx, _mask, nib) = self.stack.last()?;
927        if nib == usize::MAX {
928            return None;
929        }
930        let node = &self.trie.arena[arena_idx.as_usize()];
931        if nib == TERMINAL_NIB {
932            Some(node.leaf.as_usize())
933        } else {
934            node.leaf_key_index(nib).map(|ki| ki.as_usize())
935        }
936    }
937
938    #[inline]
939    fn advance_next(&mut self) -> bool {
940        loop {
941            let (arena_idx, mask, nib) = match self.stack.pop() {
942                Some(v) => v,
943                None => return false,
944            };
945
946            if nib == TERMINAL_NIB {
947                if self.push_next_child(arena_idx, mask, 0) {
948                    return true;
949                }
950                continue;
951            }
952
953            let search_start = if nib == usize::MAX { 0 } else { nib + 1 };
954            if self.push_next_child(arena_idx, mask, search_start) {
955                return true;
956            }
957        }
958    }
959
960    #[inline]
961    fn advance_prev(&mut self) -> bool {
962        loop {
963            let (arena_idx, mask, nib) = match self.stack.pop() {
964                Some(v) => v,
965                None => return false,
966            };
967
968            if nib == TERMINAL_NIB {
969                continue;
970            }
971
972            if nib == 0 || nib == usize::MAX {
973                if self.trie.arena[arena_idx.as_usize()].is_terminal() {
974                    self.stack.push((arena_idx, mask, TERMINAL_NIB));
975                    return true;
976                }
977                continue;
978            }
979
980            let mask_below = mask & ((1 << nib) - 1);
981            if mask_below != 0 {
982                let prev_nib = 15 - mask_below.leading_zeros() as usize;
983                self.stack.push((arena_idx, mask, prev_nib));
984                if !self.trie.arena[arena_idx.as_usize()].is_leaf(prev_nib) {
985                    self.descend_last(self.trie.arena[arena_idx.as_usize()].children[prev_nib]);
986                }
987                return true;
988            }
989
990            if self.trie.arena[arena_idx.as_usize()].is_terminal() {
991                self.stack.push((arena_idx, mask, TERMINAL_NIB));
992                return true;
993            }
994        }
995    }
996
997    #[inline]
998    pub fn next_index(&mut self) -> Option<usize> {
999        if self.advance_next() { self.current_index() } else { None }
1000    }
1001
1002    #[inline]
1003    pub fn prev_index(&mut self) -> Option<usize> {
1004        if self.advance_prev() { self.current_index() } else { None }
1005    }
1006
1007    #[inline]
1008    pub fn next(&mut self) -> Option<(&[u8], &T)> {
1009        if self.advance_next() { self.current() } else { None }
1010    }
1011
1012    #[inline]
1013    pub fn prev(&mut self) -> Option<(&[u8], &T)> {
1014        if self.advance_prev() { self.current() } else { None }
1015    }
1016
1017    pub fn seek(&mut self, key: &[u8]) -> Option<(&[u8], &T)> {
1018        if self.trie.arena.is_empty() {
1019            self.stack.clear();
1020            return None;
1021        }
1022
1023        self.stack.clear();
1024        let mut node_idx: PTR = PTR::zero();
1025        let max_nib = key.len() * 2;
1026
1027        loop {
1028            let node = &self.trie.arena[node_idx.as_usize()];
1029            let mask = node.children_mask();
1030
1031            if node.is_terminal() && node.prefix_len as usize >= max_nib {
1032                let node_key = self.trie.key_slice(node.leaf);
1033                if node_key >= key {
1034                    self.stack.push((node_idx, mask, TERMINAL_NIB));
1035                    return self.current();
1036                }
1037            }
1038
1039            if node.prefix_len as usize >= max_nib {
1040                if self.push_next_child(node_idx, mask, 0) {
1041                    return self.current();
1042                }
1043                return self.backtrack_to_next();
1044            }
1045
1046            let nib = key_nibble_at(key, node.prefix_len as usize) as usize;
1047            let slot = node.children[nib];
1048            if slot != PTR::max_value_sentinel() {
1049                self.stack.push((node_idx, mask, nib));
1050                if node.is_leaf(nib) {
1051                    let leaf_key = self.trie.key_slice(slot);
1052                    if leaf_key >= key {
1053                        return self.current();
1054                    }
1055                    return self.next();
1056                } else {
1057                    node_idx = slot;
1058                    continue;
1059                }
1060            }
1061
1062            if self.push_next_child(node_idx, mask, nib + 1) {
1063                return self.current();
1064            }
1065            return self.backtrack_to_next();
1066        }
1067    }
1068}
1069
1070// ---------------------------------------------------------------------------
1071// Tests
1072// ---------------------------------------------------------------------------
1073
1074#[cfg(test)]
1075#[path = "tests/fixed_len_nibble_trie.rs"]
1076mod tests;
tiny_trie/fixed_len_nibble_trie.rs

tiny_trie/
fixed_len_nibble_trie.rs
