use crate::ByteKey;
use crate::tiny_array::TinyArray;
use std::{fmt, marker::PhantomData, num::NonZero, ops::{Bound, RangeBounds}, simd::{Simd, cmp::SimdPartialEq}};
pub type Slot<LEN, T> = (NonZero<usize>, LEN, T);
pub trait TrieIndex: Copy + Clone + Default + PartialEq + Eq + fmt::Debug + 'static {
fn as_usize(self) -> usize;
fn max_value() -> usize;
fn zero() -> Self;
fn max_value_sentinel() -> Self;
fn from_usize(n: usize) -> Self;
fn children_mask(children: &[Self; 16]) -> u16;
}
impl TrieIndex for u8 {
#[inline] fn as_usize(self) -> usize { self as usize }
#[inline] fn max_value() -> usize { u8::MAX as usize }
#[inline] fn zero() -> Self { 0 }
#[inline] fn max_value_sentinel() -> Self { u8::MAX }
#[inline] fn from_usize(n: usize) -> Self {
debug_assert!(n <= u8::MAX as usize, "u8 overflow: {n}");
n as u8
}
#[inline] fn children_mask(children: &[Self; 16]) -> u16 {
crate::simd::children_mask_u8(children)
}
}
impl TrieIndex for u16 {
#[inline] fn as_usize(self) -> usize { self as usize }
#[inline] fn max_value() -> usize { u16::MAX as usize }
#[inline] fn zero() -> Self { 0 }
#[inline] fn max_value_sentinel() -> Self { u16::MAX }
#[inline] fn from_usize(n: usize) -> Self {
debug_assert!(n <= u16::MAX as usize, "u16 overflow: {n}");
n as u16
}
#[inline] fn children_mask(children: &[Self; 16]) -> u16 {
crate::simd::children_mask_u16(children)
}
}
impl TrieIndex for u32 {
#[inline] fn as_usize(self) -> usize { self as usize }
#[inline] fn max_value() -> usize { u32::MAX as usize }
#[inline] fn zero() -> Self { 0 }
#[inline] fn max_value_sentinel() -> Self { u32::MAX }
#[inline] fn from_usize(n: usize) -> Self {
debug_assert!(n <= u32::MAX as usize, "u32 overflow: {n}");
n as u32
}
#[inline] fn children_mask(children: &[Self; 16]) -> u16 {
crate::simd::children_mask(children)
}
}
impl TrieIndex for u64 {
#[inline] fn as_usize(self) -> usize { self as usize }
#[inline] fn max_value() -> usize { u64::MAX as usize }
#[inline] fn zero() -> Self { 0 }
#[inline] fn max_value_sentinel() -> Self { u64::MAX }
#[inline] fn from_usize(n: usize) -> Self { n as u64 }
#[inline] fn children_mask(children: &[Self; 16]) -> u16 {
crate::simd::children_mask_u64(children)
}
}
#[repr(transparent)]
#[derive(Copy, Clone, PartialEq, Eq)]
pub(crate) struct OptNz<PTR: TrieIndex>(PTR);
impl<PTR: TrieIndex> OptNz<PTR> {
#[inline]
pub(crate) fn empty() -> Self { Self(PTR::zero()) }
#[allow(dead_code)]
#[inline]
pub(crate) fn new(v: PTR) -> Option<Self> {
if v == PTR::zero() { None } else { Some(Self(v)) }
}
#[inline]
pub(crate) fn from_index(v: PTR) -> Self {
debug_assert!(v != PTR::zero(), "OptNz::from_index: zero value");
Self(v)
}
#[inline]
pub(crate) fn get(self) -> PTR { self.0 }
#[inline]
pub(crate) fn is_some(self) -> bool { self.0 != PTR::zero() }
#[inline]
pub(crate) fn is_none(self) -> bool { self.0 == PTR::zero() }
}
impl<PTR: TrieIndex> Default for OptNz<PTR> {
fn default() -> Self { Self(PTR::zero()) }
}
impl<PTR: TrieIndex> fmt::Debug for OptNz<PTR> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.is_none() { write!(f, "-") } else { write!(f, "{:?}", self.0) }
}
}
#[derive(Copy, Clone)]
pub(crate) struct Node<PTR: TrieIndex, LEN: TrieIndex> {
pub(crate) children: [OptNz<PTR>; 16], pub(crate) prefix_len: LEN, pub(crate) leaf_mask: u16, pub(crate) leaf: OptNz<PTR>, pub(crate) terminal: bool, }
impl<PTR: TrieIndex, LEN: TrieIndex> Node<PTR, LEN> {
pub(crate) fn new() -> Self {
Node {
children: [OptNz::empty(); 16],
prefix_len: LEN::zero(),
leaf_mask: 0,
leaf: OptNz::empty(),
terminal: false,
}
}
#[inline]
pub(crate) fn is_terminal(&self) -> bool {
self.terminal
}
#[inline]
fn set_terminal(&mut self, val: bool) {
self.terminal = val;
}
#[inline]
pub(crate) fn is_leaf(&self, nib: usize) -> bool {
debug_assert!(nib < 16);
(self.leaf_mask >> nib) & 1 == 1
}
#[inline]
fn set_leaf(&mut self, nib: usize) {
debug_assert!(nib < 16);
self.leaf_mask |= 1 << nib;
}
#[inline]
fn clear_leaf(&mut self, nib: usize) {
debug_assert!(nib < 16);
self.leaf_mask &= !(1 << nib);
}
#[inline]
pub(crate) fn is_occupied(&self, nib: usize) -> bool {
debug_assert!(nib < 16);
self.children[nib].is_some()
}
#[inline]
fn set_leaf_child(&mut self, nib: usize, key_index: PTR) {
debug_assert!(nib < 16);
debug_assert!(key_index != PTR::zero(), "zero key index");
self.set_leaf(nib);
self.children[nib] = OptNz::from_index(key_index);
}
#[inline]
fn set_internal_child(&mut self, nib: usize, arena_idx: PTR) {
debug_assert!(nib < 16);
debug_assert!(arena_idx != PTR::zero(), "zero arena index");
self.clear_leaf(nib);
self.children[nib] = OptNz::from_index(arena_idx);
}
#[inline]
fn leaf_key_index(&self, nib: usize) -> Option<PTR> {
debug_assert!(nib < 16);
if self.is_leaf(nib) && self.children[nib].is_some() {
Some(self.children[nib].get())
} else {
None
}
}
#[inline]
pub(crate) fn children_mask(&self) -> u16 {
let raw: &[PTR; 16] = unsafe { &*(&self.children as *const [OptNz<PTR>; 16] as *const [PTR; 16]) };
PTR::children_mask(raw)
}
pub(crate) fn promote<NewPTR: TrieIndex>(self) -> Node<NewPTR, LEN> {
let mut children = [OptNz::empty(); 16];
for i in 0..16 {
if self.children[i].is_some() {
children[i] = OptNz::from_index(NewPTR::from_usize(self.children[i].get().as_usize()));
}
}
Node {
children,
prefix_len: self.prefix_len,
leaf_mask: self.leaf_mask,
leaf: if self.leaf.is_some() {
OptNz::from_index(NewPTR::from_usize(self.leaf.get().as_usize()))
} else {
OptNz::empty()
},
terminal: self.terminal,
}
}
pub(crate) fn demote<NewPTR: TrieIndex>(self) -> Result<Node<NewPTR, LEN>, Self> {
for i in 0..16 {
if self.children[i].is_some() && self.children[i].get().as_usize() > NewPTR::max_value() {
return Err(self);
}
}
if self.leaf.is_some() && self.leaf.get().as_usize() > NewPTR::max_value() {
return Err(self);
}
let mut children = [OptNz::empty(); 16];
for i in 0..16 {
if self.children[i].is_some() {
children[i] = OptNz::from_index(NewPTR::from_usize(self.children[i].get().as_usize()));
}
}
Ok(Node {
children,
prefix_len: self.prefix_len,
leaf_mask: self.leaf_mask,
leaf: if self.leaf.is_some() {
OptNz::from_index(NewPTR::from_usize(self.leaf.get().as_usize()))
} else {
OptNz::empty()
},
terminal: self.terminal,
})
}
}
impl<PTR: TrieIndex, LEN: TrieIndex> fmt::Debug for Node<PTR, LEN> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let active: Vec<(usize, &str, PTR)> = (0..16)
.filter(|&n| self.children[n].is_some())
.map(|n| {
let tag = if self.is_leaf(n) { "L" } else { "I" };
(n, tag, self.children[n].get())
})
.collect();
f.debug_struct("Node")
.field("prefix_len", &self.prefix_len)
.field("leaf_mask", &format_args!("{:016b}", self.leaf_mask))
.field("terminal", &self.terminal)
.field("leaf", &self.leaf)
.field("children", &active)
.finish()
}
}
pub(crate) const FNODE_CAP: usize = 16;
pub(crate) const FNODE_SLOTS: usize = 15;
pub(crate) const FNODE_OFFSET_NULL: u8 = 0xFF;
#[derive(Copy, Clone, Debug)]
pub(crate) struct FlatNode<PTR: TrieIndex, LEN: TrieIndex> {
pub(crate) nibbles: u64, pub(crate) base: PTR, pub(crate) terminal: bool, pub(crate) slots: TinyArray<(LEN, u8), FNODE_SLOTS>, }
impl<PTR: TrieIndex, LEN: TrieIndex> FlatNode<PTR, LEN> {
pub(crate) fn new() -> Self {
FlatNode {
nibbles: 0,
base: PTR::zero(),
terminal: false,
slots: TinyArray::new(),
}
}
#[allow(dead_code)]
#[inline]
pub(crate) fn slot_key_index(&self, i: usize) -> Option<PTR> {
let (_plen, offset) = self.slots.as_slice()[i];
if offset == FNODE_OFFSET_NULL {
None
} else {
Some(PTR::from_usize(self.base.as_usize() + offset as usize))
}
}
#[inline]
pub(crate) fn slot_nibble(&self, i: usize) -> u8 {
((self.nibbles >> (4 * i)) & 0xF) as u8
}
#[inline]
pub(crate) fn pos_key_index(&self, pos: usize) -> Option<PTR> {
if pos == 0 {
if self.terminal { Some(self.base) } else { None }
} else {
let i = pos - 1;
let (_plen, offset) = self.slots.as_slice()[i];
if offset == FNODE_OFFSET_NULL {
None
} else {
Some(PTR::from_usize(self.base.as_usize() + offset as usize))
}
}
}
#[inline]
pub(crate) fn first_terminal_pos(&self) -> Option<usize> {
if self.terminal {
Some(0)
} else {
self.next_terminal_pos(0)
}
}
#[inline]
pub(crate) fn next_terminal_pos(&self, pos: usize) -> Option<usize> {
let slots = self.slots.as_slice();
for i in pos..slots.len() {
let (_plen, offset) = slots[i];
if offset != FNODE_OFFSET_NULL {
return Some(i + 1);
}
}
None
}
pub(crate) fn key_count(&self) -> usize {
let mut n = if self.terminal { 1 } else { 0 };
for (_, offset) in self.slots.as_slice() {
if *offset != FNODE_OFFSET_NULL {
n += 1;
}
}
n
}
fn promote<NewPTR: TrieIndex>(self) -> FlatNode<NewPTR, LEN> {
FlatNode {
nibbles: self.nibbles,
base: NewPTR::from_usize(self.base.as_usize()),
terminal: self.terminal,
slots: self.slots,
}
}
fn demote<NewPTR: TrieIndex>(self) -> Result<FlatNode<NewPTR, LEN>, Self> {
if self.base.as_usize() > NewPTR::max_value() {
return Err(self);
}
Ok(FlatNode {
nibbles: self.nibbles,
base: NewPTR::from_usize(self.base.as_usize()),
terminal: self.terminal,
slots: self.slots,
})
}
}
impl<PTR: TrieIndex, LEN: TrieIndex> Default for FlatNode<PTR, LEN> {
fn default() -> Self { Self::new() }
}
#[derive(Copy, Clone, Debug)]
pub(crate) enum ArenaNode<PTR: TrieIndex, LEN: TrieIndex> {
Inode(Node<PTR, LEN>),
Fnode(FlatNode<PTR, LEN>),
}
impl<PTR: TrieIndex, LEN: TrieIndex> ArenaNode<PTR, LEN> {
fn promote<NewPTR: TrieIndex>(self) -> ArenaNode<NewPTR, LEN> {
match self {
ArenaNode::Inode(n) => ArenaNode::Inode(n.promote()),
ArenaNode::Fnode(f) => ArenaNode::Fnode(f.promote()),
}
}
fn demote<NewPTR: TrieIndex>(self) -> Result<ArenaNode<NewPTR, LEN>, Self> {
match self {
ArenaNode::Inode(n) => n.demote().map(ArenaNode::Inode).map_err(ArenaNode::Inode),
ArenaNode::Fnode(f) => f.demote().map(ArenaNode::Fnode).map_err(ArenaNode::Fnode),
}
}
}
#[derive(Clone)]
pub struct NibbleTrie<K, T, PTR: TrieIndex = u32, LEN: TrieIndex = u16>
where
K: ByteKey,
{
pub(crate) arena: Vec<ArenaNode<PTR, LEN>>,
pub(crate) buf: Vec<u8>, pub(crate) index: Vec<Option<Slot<LEN, T>>>, pub(crate) n_keys: usize, _key: PhantomData<K>,
}
enum DivergeResult {
Duplicate,
At(usize),
}
enum PrefixCheck {
Matches,
Diverges(usize),
}
#[inline]
fn find_divergence(key_a: &[u8], key_b: &[u8], from: usize) -> DivergeResult {
let total_a = nibble_count(key_a);
let total_b = nibble_count(key_b);
let min = total_a.min(total_b);
let mut d = from;
while d < min {
if key_nibble_at(key_a, d) != key_nibble_at(key_b, d) {
return DivergeResult::At(d);
}
d += 1;
}
if total_a == total_b {
DivergeResult::Duplicate
} else {
DivergeResult::At(d)
}
}
#[inline]
fn diverging_nibble(xor: u8, byte_idx: usize) -> usize {
byte_idx * 2 + ((xor >> 4 == 0) as usize)
}
#[inline]
fn simd_eq(a: &[u8], b: &[u8]) -> bool {
if a.len() != b.len() {
return false;
}
let len = a.len();
let mut i = 0;
while i + 16 <= len {
let va = Simd::<u8, 16>::from_slice(unsafe { a.get_unchecked(i..i + 16) });
let vb = Simd::<u8, 16>::from_slice(unsafe { b.get_unchecked(i..i + 16) });
if va.simd_ne(vb).any() {
return false;
}
i += 16;
}
while i < len {
if unsafe { *a.get_unchecked(i) != *b.get_unchecked(i) } {
return false;
}
i += 1;
}
true
}
fn simd_find_divergence<const N: usize>(key_a: &[u8], key_b: &[u8], from: usize) -> DivergeResult
{
let minlen = key_a.len().min(key_b.len());
let mut i = from / 2;
while i + N <= minlen {
let a = Simd::<u8, N>::from_slice(unsafe { key_a.get_unchecked(i..i + N) });
let b = Simd::<u8, N>::from_slice(unsafe { key_b.get_unchecked(i..i + N) });
let mask = a.simd_ne(b);
if mask.any() {
let diff_byte_idx = i + mask.first_set().unwrap();
let xor = unsafe { *key_a.get_unchecked(diff_byte_idx) ^ *key_b.get_unchecked(diff_byte_idx) };
return DivergeResult::At(diverging_nibble(xor, diff_byte_idx));
}
i += N;
}
find_divergence(key_a, key_b, i * 2)
}
#[inline]
fn check_prefix(key_a: &[u8], key_b: &[u8], from: usize, to: usize) -> PrefixCheck {
for nib in from..to {
if key_nibble_at(key_a, nib) != key_nibble_at(key_b, nib) {
return PrefixCheck::Diverges(nib);
}
}
PrefixCheck::Matches
}
fn simd_check_prefix<const N: usize>(key_a: &[u8], key_b: &[u8], from: usize, to: usize) -> PrefixCheck
{
if from >= to {
return PrefixCheck::Matches;
}
let from_byte = from / 2;
let to_byte = (to + 1) / 2; let minlen = key_a.len().min(key_b.len()).min(to_byte);
let mut i = from_byte;
while i + N <= minlen {
let a = Simd::<u8, N>::from_slice(unsafe { key_a.get_unchecked(i..i + N) });
let b = Simd::<u8, N>::from_slice(unsafe { key_b.get_unchecked(i..i + N) });
let mask = a.simd_ne(b);
if mask.any() {
let diff_byte_idx = i + mask.first_set().unwrap();
let xor = unsafe { *key_a.get_unchecked(diff_byte_idx) ^ *key_b.get_unchecked(diff_byte_idx) };
let nib = diverging_nibble(xor, diff_byte_idx);
if nib < to {
return PrefixCheck::Diverges(nib);
}
return PrefixCheck::Matches;
}
i += N;
}
check_prefix(key_a, key_b, i * 2, to)
}
#[inline]
fn key_nibble_at(key: &[u8], idx: usize) -> u8 {
let byte_idx = idx / 2;
if byte_idx < key.len() {
if idx % 2 == 0 {
key[byte_idx] >> 4
} else {
key[byte_idx] & 0x0F
}
} else {
0
}
}
#[allow(dead_code)]
#[inline]
unsafe fn key_nibble_at_unchecked(key: &[u8], idx: usize) -> u8 {
let byte_idx = idx / 2;
debug_assert!(byte_idx < key.len(), "nibble {idx} out of bounds for key len {}", key.len());
if idx % 2 == 0 {
unsafe { *key.get_unchecked(byte_idx) >> 4 }
} else {
unsafe { *key.get_unchecked(byte_idx) & 0x0F }
}
}
#[inline]
fn nibble_count(key: &[u8]) -> usize {
key.len() * 2
}
impl<K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> NibbleTrie<K, T, PTR, LEN> {
#[inline]
fn key_slice(&self, key_index: PTR) -> &[u8] {
let (off, len, _) = self.index[key_index.as_usize()].as_ref().unwrap();
&self.buf[off.get()..off.get() + len.as_usize()]
}
#[inline]
fn inode(&self, i: usize) -> &Node<PTR, LEN> {
match &self.arena[i] {
ArenaNode::Inode(n) => n,
ArenaNode::Fnode(_) => panic!("inode(): arena[{i}] is an Fnode (Inode-only path)"),
}
}
#[inline]
fn inode_mut(&mut self, i: usize) -> &mut Node<PTR, LEN> {
match &mut self.arena[i] {
ArenaNode::Inode(n) => n,
ArenaNode::Fnode(_) => panic!("inode_mut(): arena[{i}] is an Fnode (Inode-only path)"),
}
}
pub fn new() -> Self {
NibbleTrie {
arena: Vec::new(),
buf: vec![0], index: vec![None], n_keys: 0,
_key: PhantomData,
}
}
pub fn len(&self) -> usize {
self.n_keys
}
pub fn is_empty(&self) -> bool {
self.n_keys == 0
}
fn flat_get(&self, node: &FlatNode<PTR, LEN>, key: &[u8]) -> Option<usize> {
let slots = node.slots.as_slice();
let max_nib = key.len() * 2;
if !slots.is_empty() {
let mut depth = slots[0].0.as_usize(); if depth < max_nib {
let mut i = 0;
while i < slots.len() {
let d = slots[i].0.as_usize();
if d < depth {
break;
}
if d > depth {
i += 1;
continue;
}
let nib = node.slot_nibble(i);
if key_nibble_at(key, d) != nib {
i += 1;
continue;
}
let can_descend = i + 1 < slots.len()
&& slots[i + 1].0.as_usize() > d
&& slots[i + 1].0.as_usize() < max_nib;
if can_descend {
depth = slots[i + 1].0.as_usize();
i += 1;
} else {
let offset = slots[i].1;
return if offset != FNODE_OFFSET_NULL {
let ki = node.base.as_usize() + offset as usize;
if simd_eq(self.key_slice(PTR::from_usize(ki)), key) {
Some(ki)
} else {
None
}
} else {
None };
}
}
}
}
if node.terminal {
let ki = node.base;
if simd_eq(self.key_slice(ki), key) {
return Some(ki.as_usize());
}
}
None
}
pub fn get_index(&self, key: &[u8]) -> Option<usize> {
if self.arena.is_empty() {
return None;
}
let mut phys_idx: usize = 0;
let max_nib = key.len() * 2;
loop {
let node = self.inode(phys_idx);
let prefix_len = node.prefix_len.as_usize();
if prefix_len >= max_nib {
if node.is_terminal() {
let ki = node.leaf.get();
let (off, len, _) = self.index[ki.as_usize()].as_ref().unwrap();
let off = off.get();
let key_in_buf = &self.buf[off..off + len.as_usize()];
if key.len() == len.as_usize() && simd_eq(&key_in_buf[..key.len()], key) {
return Some(ki.as_usize());
}
}
return None;
}
let nib = key_nibble_at(key, prefix_len) as usize;
if !node.is_occupied(nib) {
return None;
}
if node.is_leaf(nib) {
let key_index = node.children[nib].get();
return if simd_eq(self.key_slice(key_index), key) {
Some(key_index.as_usize())
} else {
None
};
}
let child = node.children[nib].get().as_usize();
match &self.arena[child] {
ArenaNode::Inode(_) => phys_idx = child,
ArenaNode::Fnode(f) => return self.flat_get(f, key),
}
}
}
#[cfg(feature = "unchecked")]
unsafe fn get_index_unchecked(&self, key: &[u8]) -> Option<usize> {
if self.arena.is_empty() {
return None;
}
let mut phys_idx: usize = 0;
let max_nib = key.len() * 2;
loop {
let node = match unsafe { self.arena.get_unchecked(phys_idx) } {
ArenaNode::Inode(n) => n,
ArenaNode::Fnode(_) => panic!("get_unchecked: phys_idx {phys_idx} is an Fnode (dispatcher missed it)"),
};
let prefix_len = node.prefix_len.as_usize();
if prefix_len >= max_nib {
debug_assert!(node.is_terminal(), "get_unchecked: key not in set");
return Some(node.leaf.get().as_usize());
}
let nib = unsafe { key_nibble_at_unchecked(key, prefix_len) } as usize;
let slot = unsafe { node.children.get_unchecked(nib) };
if slot.is_none() {
return None;
}
if node.is_leaf(nib) {
return Some(slot.get().as_usize());
}
let child = slot.get().as_usize();
match unsafe { self.arena.get_unchecked(child) } {
ArenaNode::Inode(_) => phys_idx = child,
ArenaNode::Fnode(f) => return self.flat_get(f, key),
}
}
}
pub fn get(&self, key: &[u8]) -> Option<&T> {
self.get_index(key).map(|idx| &self.index[idx].as_ref().unwrap().2)
}
pub fn get_mut(&mut self, key: &[u8]) -> Option<&mut T> {
self.get_index(key).map(|idx| &mut self.index[idx].as_mut().unwrap().2)
}
#[cfg(feature = "unchecked")]
pub unsafe fn get_unchecked(&self, key: &[u8]) -> Option<&T> {
unsafe {self.get_index_unchecked(key).map(|idx| &self.index[idx].as_ref().unwrap().2) }
}
pub(crate) fn walk_iter(&self) -> NibbleIter<'_, K, T, PTR, LEN> {
NibbleIter::new(self)
}
pub fn iter(&self) -> Cursor<'_, K, T, PTR, LEN> {
Cursor::new(self)
}
pub fn iter_last(&self) -> Cursor<'_, K, T, PTR, LEN> {
Cursor::new_last(self)
}
pub fn iter_mut(&mut self) -> CursorMut<'_, K, T, PTR, LEN> {
CursorMut::new(self)
}
pub fn iter_mut_last(&mut self) -> CursorMut<'_, K, T, PTR, LEN> {
CursorMut::new_last(self)
}
pub fn range<'q>(&self, bounds: impl RangeBounds<&'q [u8]>) -> Range<'_, K, T, PTR, LEN> {
let start = bounds.start_bound().map(|b| *b);
let end = bounds.end_bound().map(|b| *b);
Range::new(self, start, end)
}
pub fn range_bounds(
&self,
start: Bound<&[u8]>,
end: Bound<&[u8]>,
) -> Range<'_, K, T, PTR, LEN> {
Range::new(self, start, end)
}
pub fn into_keys_values(self) -> (Vec<K>, Vec<T>) {
let buf = self.buf;
let mut keys: Vec<K> = Vec::with_capacity(self.n_keys);
let mut values: Vec<T> = Vec::with_capacity(self.n_keys);
for (i, slot) in self.index.into_iter().enumerate() {
if i == 0 { continue; } if let Some((off, len, val)) = slot {
keys.push(K::from_bytes(&buf[off.get()..off.get() + len.as_usize()]));
values.push(val);
}
}
(keys, values)
}
pub fn near_capacity(&self) -> bool {
self.arena.len() >= PTR::max_value() || self.index.len() >= PTR::max_value()
}
pub fn optimize(&mut self) {
if self.arena.is_empty() {
return;
}
let n = self.n_keys;
let cap = 2 * n + 1;
let mut new_index: Vec<Option<Slot<LEN, T>>> = (0..cap).map(|_| None).collect();
let mut new_buf: Vec<u8> = Vec::with_capacity(self.buf.len());
new_buf.push(0); let mut cursor: usize = 1;
let mut i: usize = 0;
self.walk_optimize(0, &mut new_index, &mut new_buf, &mut cursor, &mut i);
new_buf.truncate(cursor);
self.buf = new_buf;
self.index = new_index;
}
fn walk_optimize(
&mut self,
phys_idx: usize,
new_index: &mut Vec<Option<Slot<LEN, T>>>,
new_buf: &mut Vec<u8>,
cursor: &mut usize,
i: &mut usize,
) -> usize {
let node = *self.inode(phys_idx);
let mut first: Option<usize> = None;
if node.is_terminal() {
let slot = self.place_key(node.leaf.get().as_usize(), new_index, new_buf, cursor, i);
first = Some(slot);
}
for nib in 0..16 {
if !node.is_occupied(nib) {
continue;
}
let child_phys = node.children[nib].get().as_usize();
if node.is_leaf(nib) {
let slot = self.place_key(child_phys, new_index, new_buf, cursor, i);
self.inode_mut(phys_idx).children[nib] = OptNz::from_index(PTR::from_usize(slot));
if first.is_none() {
first = Some(slot);
}
} else {
let child_first = match self.arena[child_phys] {
ArenaNode::Inode(_) => {
self.walk_optimize(child_phys, new_index, new_buf, cursor, i)
}
ArenaNode::Fnode(f) => {
let old_base = f.base.as_usize();
let new_base = self.place_key(old_base, new_index, new_buf, cursor, i);
let mut new_slots: TinyArray<(LEN, u8), FNODE_SLOTS> = TinyArray::new();
for (plen, offset) in f.slots.as_slice() {
if *offset == FNODE_OFFSET_NULL {
new_slots.push((*plen, FNODE_OFFSET_NULL));
} else if *offset == 0 {
new_slots.push((*plen, 0));
} else {
let old_ki = old_base + *offset as usize;
let new_slot = self.place_key(old_ki, new_index, new_buf, cursor, i);
new_slots.push((*plen, (new_slot - new_base) as u8));
}
}
self.arena[child_phys] = ArenaNode::Fnode(FlatNode {
nibbles: f.nibbles,
base: PTR::from_usize(new_base),
terminal: f.terminal,
slots: new_slots,
});
new_base
}
};
if first.is_none() {
first = Some(child_first);
}
}
}
let leftmost = first.expect("walk_optimize: node must have at least one key in subtree");
self.inode_mut(phys_idx).leaf = OptNz::from_index(PTR::from_usize(leftmost));
leftmost
}
fn place_key(
&mut self,
old_ki: usize,
new_index: &mut Vec<Option<Slot<LEN, T>>>,
new_buf: &mut Vec<u8>,
cursor: &mut usize,
i: &mut usize,
) -> usize {
let slot = 2 * *i + 1;
*i += 1;
let (off, len, val) = self.index[old_ki].take().unwrap();
let old_off = off.get();
let start = *cursor;
new_buf.resize(start + len.as_usize(), 0);
new_buf[start..start + len.as_usize()]
.copy_from_slice(&self.buf[old_off..old_off + len.as_usize()]);
*cursor = start + len.as_usize();
new_index[slot] = Some((NonZero::new(start).unwrap(), len, val));
slot
}
pub fn flatten(&mut self) {
if self.arena.is_empty() {
return;
}
let mut counts: Vec<(usize, usize)> = vec![(0, 0); self.arena.len()];
self.count_subtree(0, &mut counts);
let mut new_arena: Vec<ArenaNode<PTR, LEN>> = Vec::with_capacity(self.arena.len());
self.rebuild_subtree(0, &mut new_arena, &counts);
self.arena = new_arena;
}
fn count_subtree(&self, phys: usize, counts: &mut [(usize, usize)]) {
if counts[phys] != (0, 0) {
return; }
let (keys, inodes) = match &self.arena[phys] {
ArenaNode::Fnode(f) => (f.key_count(), 1),
ArenaNode::Inode(node) => {
let mut k = if node.is_terminal() { 1 } else { 0 };
let mut i = 1;
for nib in 0..16 {
if !node.is_occupied(nib) {
continue;
}
if node.is_leaf(nib) {
k += 1;
} else {
let child = node.children[nib].get().as_usize();
self.count_subtree(child, counts);
let (ck, ci) = counts[child];
k += ck;
i += ci;
}
}
(k, i)
}
};
counts[phys] = (keys, inodes);
}
fn rebuild_subtree(
&self,
phys: usize,
new_arena: &mut Vec<ArenaNode<PTR, LEN>>,
counts: &[(usize, usize)],
) -> usize {
if let ArenaNode::Fnode(f) = &self.arena[phys] {
new_arena.push(ArenaNode::Fnode(*f));
return new_arena.len() - 1;
}
let node = *self.inode(phys);
let (n_keys, n_inodes) = counts[phys];
if phys != 0 && n_keys <= FNODE_CAP && n_inodes >= 2 {
if let Some(fnode) = self.build_fnode_subtree(phys) {
new_arena.push(ArenaNode::Fnode(fnode));
return new_arena.len() - 1;
}
}
let new_phys = new_arena.len();
new_arena.push(ArenaNode::Inode(node));
for nib in 0..16 {
if node.is_occupied(nib) && !node.is_leaf(nib) {
let child_old = node.children[nib].get().as_usize();
let child_new = self.rebuild_subtree(child_old, new_arena, counts);
match &mut new_arena[new_phys] {
ArenaNode::Inode(n) => {
n.children[nib] = OptNz::from_index(PTR::from_usize(child_new))
}
_ => unreachable!("rebuild_subtree: placeholder was Inode"),
}
}
}
new_phys
}
fn build_fnode_subtree(&self, phys: usize) -> Option<FlatNode<PTR, LEN>> {
if !matches!(self.arena[phys], ArenaNode::Inode(_)) {
return None;
}
let root = *self.inode(phys);
let base = root.leaf.get();
let terminal = root.is_terminal();
let mut plens: Vec<LEN> = Vec::new();
let mut key_idxs: Vec<Option<PTR>> = Vec::new();
let mut nibbles: u64 = 0;
let mut ok = true;
self.collect_flat_slots(phys, &mut plens, &mut key_idxs, &mut nibbles, &mut ok);
if !ok || plens.is_empty() || plens.len() > FNODE_SLOTS {
return None;
}
let base_u = base.as_usize();
let mut slots: TinyArray<(LEN, u8), FNODE_SLOTS> = TinyArray::new();
for (plen, kidx) in plens.into_iter().zip(key_idxs.into_iter()) {
let offset = match kidx {
None => FNODE_OFFSET_NULL,
Some(ki) => {
let off = ki.as_usize() - base_u;
if off >= FNODE_OFFSET_NULL as usize {
return None;
}
off as u8
}
};
slots.push((plen, offset));
}
Some(FlatNode { nibbles, base, terminal, slots })
}
fn collect_flat_slots(
&self,
phys: usize,
plens: &mut Vec<LEN>,
key_idxs: &mut Vec<Option<PTR>>,
nibbles: &mut u64,
ok: &mut bool,
) {
let node = *self.inode(phys);
let p = node.prefix_len;
for nib in 0..16 {
if !node.is_occupied(nib) {
continue;
}
let i = plens.len();
if i >= FNODE_SLOTS {
*ok = false;
return;
}
*nibbles |= (nib as u64) << (4 * i);
if node.is_leaf(nib) {
plens.push(p);
key_idxs.push(Some(node.children[nib].get()));
} else {
let child = node.children[nib].get().as_usize();
if matches!(self.arena[child], ArenaNode::Fnode(_)) {
*ok = false;
return;
}
let child_node = *self.inode(child);
let ptr = if child_node.is_terminal() {
Some(child_node.leaf.get())
} else {
None
};
plens.push(p);
key_idxs.push(ptr);
self.collect_flat_slots(child, plens, key_idxs, nibbles, ok);
if !*ok {
return;
}
}
}
}
}
impl<K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> Default for NibbleTrie<K, T, PTR, LEN> {
fn default() -> Self { Self::new() }
}
enum Case {
Terminal { phys: usize },
NewLeafChild { phys: usize, nib: usize },
SplitNode {
phys: usize,
diverge: usize,
new_is_terminal: bool,
new_nib: usize,
ref_nib: usize,
new_is_leftmost: bool,
},
SplitLeaf {
phys: usize,
nib: usize,
d: usize,
new_is_terminal: bool,
existing_is_terminal: bool,
new_nib: usize,
exist_nib: usize,
new_is_leftmost: bool,
},
}
impl<K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> NibbleTrie<K, T, PTR, LEN> {
pub fn insert(&mut self, key: K, value: T) -> Result<usize, ()> {
let key_bytes = key.bytes();
if self.arena.len() >= PTR::max_value() || self.index.len() >= PTR::max_value() {
return Err(());
}
if key_bytes.len() * 2 > LEN::max_value() {
return Err(());
}
if self.needs_optimize() && 2 * self.n_keys < PTR::max_value() {
self.optimize();
}
if self.arena.len() >= PTR::max_value() || self.index.len() >= PTR::max_value() {
return Err(());
}
let key_len = LEN::from_usize(key_bytes.len());
let off = self.buf.len();
self.buf.extend_from_slice(key_bytes);
self.n_keys += 1;
let max_nib = key_bytes.len() * 2;
if self.arena.is_empty() {
return Ok(self.insert_into_empty_trie(off, key_len, value, key_bytes, max_nib));
}
let (case, path) = match self.find_insert_case(key_bytes, max_nib) {
Ok(c) => c,
Err(()) => {
self.buf.truncate(off);
self.n_keys -= 1;
return Err(());
}
};
let p_opt = self.compute_p(&case, &path);
let (p, n) = match p_opt {
None => {
let p = self.index.len();
self.index
.push(Some((NonZero::new(off).unwrap(), key_len, value)));
self.execute_case(case, p, &path);
return Ok(p);
}
Some(p) => {
let mut n = 0;
while p + n < self.index.len() && self.index[p + n].is_some() {
n += 1;
}
(p, n)
}
};
if p + n >= self.index.len() {
self.index.push(None);
}
if n > 0 {
let succ_bytes = {
let (soff, slen, _) = self.index[p].as_ref().unwrap();
self.buf[soff.get()..soff.get() + slen.as_usize()].to_vec()
};
let stack: Vec<(usize, u16, usize)> = {
let mut it = self.walk_iter();
it.seek(&succ_bytes);
debug_assert_eq!(
it.current_index(),
Some(p),
"seek must land on the successor slot"
);
it.stack
.iter()
.map(|frame| match *frame {
Frame::Inode { encoded, mask, nib } => (encoded.as_usize(), mask, nib),
Frame::Fnode { .. } => panic!(
"bump_walk init: Fnode frame on stack — insert-into-Fnode is step 5"
),
})
.collect()
};
self.bump_walk(stack, p, n);
for i in (0..n).rev() {
self.index[p + i + 1] = self.index[p + i].take();
}
}
self.index[p] = Some((NonZero::new(off).unwrap(), key_len, value));
self.execute_case(case, p, &path);
Ok(p)
}
#[inline]
fn needs_optimize(&self) -> bool {
let idx_cap = self.index.capacity();
let buf_cap = self.buf.capacity();
(idx_cap > 0 && 10 * self.n_keys > 9 * idx_cap)
|| (buf_cap > 0 && 10 * self.buf.len() > 9 * buf_cap)
}
fn find_insert_case(
&self,
key: &[u8],
max_nib: usize,
) -> Result<(Case, Vec<(usize, usize)>), ()> {
let mut phys_idx: usize = 0;
let mut confirmed: usize = 0;
let mut path: Vec<(usize, usize)> = Vec::new();
loop {
let node = self.inode(phys_idx);
let ki = node.leaf.get();
let (off, ref_len, _) = self.index[ki.as_usize()].as_ref().unwrap();
let off = off.get();
let ref_key = &self.buf[off..off + ref_len.as_usize()];
let prefix_len = node.prefix_len.as_usize();
match simd_check_prefix::<8>(key, ref_key, confirmed, prefix_len) {
PrefixCheck::Diverges(diverge) => {
let new_nib = key_nibble_at(key, diverge) as usize;
let ref_nib = key_nibble_at(ref_key, diverge) as usize;
let new_is_terminal = diverge >= max_nib;
let new_is_leftmost = new_is_terminal || new_nib < ref_nib;
return Ok((
Case::SplitNode {
phys: phys_idx,
diverge,
new_is_terminal,
new_nib,
ref_nib,
new_is_leftmost,
},
path,
));
}
PrefixCheck::Matches => {
if max_nib == prefix_len {
if key.len() == ref_key.len() {
return Err(()); }
return Ok((Case::Terminal { phys: phys_idx }, path));
}
confirmed = prefix_len + 1;
let nib = key_nibble_at(key, prefix_len) as usize;
if !node.is_occupied(nib) {
return Ok((Case::NewLeafChild { phys: phys_idx, nib }, path));
}
if node.is_leaf(nib) {
path.push((phys_idx, nib));
let existing_key_index = node.children[nib].get();
let (eo, elen, _) =
self.index[existing_key_index.as_usize()].as_ref().unwrap();
let existing_key = &self.buf[eo.get()..eo.get() + elen.as_usize()];
match simd_find_divergence::<8>(key, existing_key, confirmed) {
DivergeResult::Duplicate => return Err(()),
DivergeResult::At(d) => {
let new_is_terminal = d >= max_nib;
let existing_is_terminal = d >= existing_key.len() * 2;
let new_nib = key_nibble_at(key, d) as usize;
let exist_nib = key_nibble_at(existing_key, d) as usize;
let new_is_leftmost = if new_is_terminal {
true
} else if existing_is_terminal {
false
} else {
new_nib < exist_nib
};
return Ok((
Case::SplitLeaf {
phys: phys_idx,
nib,
d,
new_is_terminal,
existing_is_terminal,
new_nib,
exist_nib,
new_is_leftmost,
},
path,
));
}
}
}
path.push((phys_idx, nib));
phys_idx = node.children[nib].get().as_usize();
}
}
}
}
fn compute_p(&self, case: &Case, path: &[(usize, usize)]) -> Option<usize> {
match case {
Case::Terminal { phys } => Some(self.inode(*phys).leaf.get().as_usize()),
Case::NewLeafChild { phys, nib } => self.right_anchor(*phys, *nib, path),
Case::SplitNode {
phys,
new_is_terminal,
new_nib,
ref_nib,
..
} => {
if *new_is_terminal || *new_nib < *ref_nib {
Some(self.inode(*phys).leaf.get().as_usize())
} else {
self.subtree_successor(path)
}
}
Case::SplitLeaf {
phys,
nib,
new_is_terminal,
existing_is_terminal,
new_nib,
exist_nib,
..
} => {
let existing_key_index = self.inode(*phys).children[*nib].get().as_usize();
if *new_is_terminal {
Some(existing_key_index)
} else if *existing_is_terminal {
self.right_anchor(*phys, *nib, path)
} else if *new_nib < *exist_nib {
Some(existing_key_index)
} else {
self.right_anchor(*phys, *nib, path)
}
}
}
}
fn right_anchor(&self, phys: usize, nib: usize, path: &[(usize, usize)]) -> Option<usize> {
let mask = self.inode(phys).children_mask();
let higher = if nib >= 15 { 0u16 } else { mask & !((1u16 << (nib + 1)) - 1) };
if higher != 0 {
let next_nib = higher.trailing_zeros() as usize;
let r = self.inode(phys).children[next_nib].get();
Some(if self.inode(phys).is_leaf(next_nib) {
r.as_usize()
} else {
self.inode(r.as_usize()).leaf.get().as_usize()
})
} else {
self.subtree_successor(path)
}
}
fn subtree_successor(&self, path: &[(usize, usize)]) -> Option<usize> {
for &(parent, nib) in path.iter().rev() {
let mask = self.inode(parent).children_mask();
let higher = if nib >= 15 { 0u16 } else { mask & !((1u16 << (nib + 1)) - 1) };
if higher != 0 {
let next_nib = higher.trailing_zeros() as usize;
let r = self.inode(parent).children[next_nib].get();
return Some(if self.inode(parent).is_leaf(next_nib) {
r.as_usize()
} else {
self.inode(r.as_usize()).leaf.get().as_usize()
});
}
}
None
}
fn bump_walk(&mut self, init_stack: Vec<(usize, u16, usize)>, lo: usize, n: usize) {
debug_assert!(n >= 1);
let hi = lo + n - 1; let mut stack = init_stack;
for &(phys, _mask, _nib) in &stack {
let l = self.inode(phys).leaf.get().as_usize();
if l >= lo && l <= hi {
self.inode_mut(phys).leaf = OptNz::from_index(PTR::from_usize(l + 1));
}
}
let mut seen = 0;
while seen < n {
let &(phys, _mask, nib) = stack.last().expect("bump_walk: stack emptied early");
if nib == TERMINAL_NIB {
seen += 1;
} else {
let k = self.inode(phys).children[nib].get().as_usize();
self.inode_mut(phys).children[nib] = OptNz::from_index(PTR::from_usize(k + 1));
seen += 1;
}
if seen == n {
break;
}
if !self.bump_advance(&mut stack, lo, hi) {
debug_assert!(seen >= n, "bump_walk: tree exhausted before n keys");
break;
}
}
}
fn bump_descend_first(
&mut self,
stack: &mut Vec<(usize, u16, usize)>,
mut phys: usize,
lo: usize,
hi: usize,
) {
loop {
let node = *self.inode(phys);
let l = node.leaf.get().as_usize();
if l >= lo && l <= hi {
self.inode_mut(phys).leaf = OptNz::from_index(PTR::from_usize(l + 1));
}
if node.is_terminal() {
let mask = node.children_mask();
stack.push((phys, mask, TERMINAL_NIB));
return;
}
let mask = node.children_mask();
debug_assert!(mask != 0, "bump_descend_first: non-terminal node with no children");
let nib = mask.trailing_zeros() as usize;
stack.push((phys, mask, nib));
if node.is_leaf(nib) {
return;
} else {
phys = node.children[nib].get().as_usize();
}
}
}
#[inline]
fn bump_push_next(
&mut self,
stack: &mut Vec<(usize, u16, usize)>,
encoded: usize,
mask: u16,
start_nib: usize,
lo: usize,
hi: usize,
) -> bool {
let shifted = if start_nib >= 16 { 0u16 } else { mask >> start_nib };
if shifted == 0 {
return false;
}
let nib = start_nib + shifted.trailing_zeros() as usize;
debug_assert!(nib < 16);
stack.push((encoded, mask, nib));
if !self.inode(encoded).is_leaf(nib) {
let addr = self.inode(encoded).children[nib].get().as_usize();
self.bump_descend_first(stack, addr, lo, hi);
}
true
}
#[inline]
fn bump_advance(
&mut self,
stack: &mut Vec<(usize, u16, usize)>,
lo: usize,
hi: usize,
) -> bool {
loop {
let (encoded, mask, nib) = match stack.pop() {
Some(v) => v,
None => return false,
};
if nib == TERMINAL_NIB {
if self.bump_push_next(stack, encoded, mask, 0, lo, hi) {
return true;
}
continue;
}
let search_start = if nib == usize::MAX { 0 } else { nib + 1 };
if self.bump_push_next(stack, encoded, mask, search_start, lo, hi) {
return true;
}
}
}
fn execute_case(&mut self, case: Case, p: usize, path: &[(usize, usize)]) {
let p_idx = PTR::from_usize(p);
match case {
Case::Terminal { phys } => {
self.inode_mut(phys).set_terminal(true);
self.inode_mut(phys).leaf = OptNz::from_index(p_idx);
self.up_walk_leftmost(phys, p_idx, path);
}
Case::NewLeafChild { phys, nib } => {
self.inode_mut(phys).set_leaf_child(nib, p_idx);
self.update_leftmost_on_leaf_insert(phys, nib, p_idx, path);
}
Case::SplitNode {
phys,
diverge,
new_is_terminal,
new_nib,
ref_nib,
new_is_leftmost,
} => {
let mut new_parent = Node::new();
new_parent.prefix_len = LEN::from_usize(diverge);
if new_is_terminal {
new_parent.set_terminal(true);
new_parent.leaf = OptNz::from_index(p_idx);
} else {
new_parent.set_leaf_child(new_nib, p_idx);
if new_is_leftmost {
new_parent.leaf = OptNz::from_index(p_idx);
}
}
let old_node = std::mem::replace(&mut self.arena[phys], ArenaNode::Inode(new_parent));
let old_addr = PTR::from_usize(self.arena.len()); self.arena.push(old_node);
self.inode_mut(phys).set_internal_child(ref_nib, old_addr);
if new_is_leftmost {
self.up_walk_leftmost(phys, p_idx, path);
} else {
let child_leaf = self.inode(old_addr.as_usize()).leaf;
self.inode_mut(phys).leaf = child_leaf;
}
}
Case::SplitLeaf {
phys,
nib,
d,
new_is_terminal,
existing_is_terminal,
new_nib,
exist_nib,
new_is_leftmost,
} => {
let existing_key_index = self.inode(phys).children[nib].get();
let mut split_node = Node::new();
split_node.prefix_len = LEN::from_usize(d);
if new_is_terminal {
split_node.set_terminal(true);
split_node.leaf = OptNz::from_index(p_idx);
split_node.set_leaf_child(exist_nib, existing_key_index);
} else if existing_is_terminal {
split_node.set_terminal(true);
split_node.leaf = OptNz::from_index(existing_key_index);
split_node.set_leaf_child(new_nib, p_idx);
} else {
split_node.set_leaf_child(new_nib, p_idx);
split_node.set_leaf_child(exist_nib, existing_key_index);
split_node.leaf = OptNz::from_index(if new_is_leftmost {
p_idx
} else {
existing_key_index
});
}
let split_addr = PTR::from_usize(self.arena.len());
self.arena.push(ArenaNode::Inode(split_node));
self.inode_mut(phys).set_internal_child(nib, split_addr);
if new_is_leftmost {
self.up_walk_leftmost(split_addr.as_usize(), p_idx, path);
}
}
}
}
#[inline]
fn update_leftmost_on_leaf_insert(
&mut self,
phys_idx: usize,
nib: usize,
new_index: PTR,
path: &[(usize, usize)],
) {
if self.inode(phys_idx).is_terminal() {
return;
}
let mask = self.inode(phys_idx).children_mask();
let lowest = mask.trailing_zeros() as usize;
if nib == lowest {
self.inode_mut(phys_idx).leaf = OptNz::from_index(new_index);
self.up_walk_leftmost(phys_idx, new_index, path);
}
}
#[inline]
fn up_walk_leftmost(&mut self, attach_phys: usize, new_leftmost: PTR, path: &[(usize, usize)]) {
let _ = attach_phys; let mut idx = path.len();
while idx > 0 {
idx -= 1;
let (parent_phys, nib) = path[idx];
if self.inode(parent_phys).is_terminal() {
break;
}
let parent_mask = self.inode(parent_phys).children_mask();
let lowest = parent_mask.trailing_zeros() as usize;
if nib == lowest {
self.inode_mut(parent_phys).leaf = OptNz::from_index(new_leftmost);
} else {
break;
}
}
}
#[inline]
fn insert_into_empty_trie(
&mut self,
off: usize,
key_len: LEN,
value: T,
key: &[u8],
max_nib: usize,
) -> usize {
let p = 1usize;
self.index
.push(Some((NonZero::new(off).unwrap(), key_len, value)));
let p_idx = PTR::from_usize(p);
if max_nib == 0 {
let mut root = Node::new();
root.set_terminal(true);
root.leaf = OptNz::from_index(p_idx);
self.arena.push(ArenaNode::Inode(root));
} else {
let first_nib = key_nibble_at(key, 0) as usize;
let mut root = Node::new();
root.set_leaf_child(first_nib, p_idx);
root.leaf = OptNz::from_index(p_idx);
self.arena.push(ArenaNode::Inode(root));
}
p
}
}
impl<K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> NibbleTrie<K, T, PTR, LEN> {
pub fn promote<NewPTR: TrieIndex>(self) -> NibbleTrie<K, T, NewPTR, LEN> {
let arena = self.arena.into_iter().map(|node| node.promote()).collect();
NibbleTrie {
arena,
buf: self.buf,
index: self.index,
n_keys: self.n_keys,
_key: PhantomData,
}
}
pub fn demote<NewPTR: TrieIndex>(self) -> Result<NibbleTrie<K, T, NewPTR, LEN>, Self> {
if self.arena.len() > NewPTR::max_value() || self.index.len() > NewPTR::max_value() {
return Err(self);
}
let arena = self.arena.into_iter().map(|node| {
node.demote().expect("demote capacity check should have caught this")
}).collect();
Ok(NibbleTrie {
arena,
buf: self.buf,
index: self.index,
n_keys: self.n_keys,
_key: PhantomData,
})
}
}
const TERMINAL_NIB: usize = 16;
#[derive(Clone, Copy)]
pub(crate) enum Frame<PTR: TrieIndex> {
Inode { encoded: PTR, mask: u16, nib: usize },
Fnode { arena_idx: PTR, pos: usize },
}
pub(crate) struct NibbleIter<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> {
trie: &'a NibbleTrie<K, T, PTR, LEN>,
pub(crate) stack: Vec<Frame<PTR>>,
}
impl<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> NibbleIter<'a, K, T, PTR, LEN> {
fn new(trie: &'a NibbleTrie<K, T, PTR, LEN>) -> Self {
if trie.arena.is_empty() {
return NibbleIter { trie, stack: Vec::new() };
}
let mask = trie.inode(0).children_mask();
let nib = if trie.inode(0).is_terminal() { TERMINAL_NIB } else { usize::MAX };
NibbleIter { trie, stack: vec![Frame::Inode { encoded: PTR::zero(), mask, nib }] }
}
fn descend_first(&mut self, mut phys_idx: usize) {
loop {
let fnode_pos = match &self.trie.arena[phys_idx] {
ArenaNode::Fnode(f) => Some(
f.first_terminal_pos()
.expect("descend_first: Fnode with no terminals"),
),
ArenaNode::Inode(_) => None,
};
if let Some(pos) = fnode_pos {
self.stack.push(Frame::Fnode { arena_idx: PTR::from_usize(phys_idx), pos });
return;
}
let node = *self.trie.inode(phys_idx);
if node.is_terminal() {
let mask = node.children_mask();
self.stack.push(Frame::Inode { encoded: PTR::from_usize(phys_idx), mask, nib: TERMINAL_NIB });
return;
}
let mask = node.children_mask();
debug_assert!(mask != 0, "descend_first: non-terminal node with no children");
let nib = mask.trailing_zeros() as usize;
self.stack.push(Frame::Inode { encoded: PTR::from_usize(phys_idx), mask, nib });
if node.is_leaf(nib) {
return;
}
phys_idx = node.children[nib].get().as_usize();
}
}
#[inline]
fn push_next_child(&mut self, encoded: PTR, mask: u16, start_nib: usize) -> bool {
let shifted = if start_nib >= 16 { 0u16 } else { mask >> start_nib };
if shifted == 0 {
return false;
}
let nib = start_nib + shifted.trailing_zeros() as usize;
debug_assert!(nib < 16);
debug_assert!(mask & (1 << nib) != 0);
let phys_idx = encoded.as_usize();
let node = *self.trie.inode(phys_idx);
self.stack.push(Frame::Inode { encoded, mask, nib });
if !node.is_leaf(nib) {
let addr = node.children[nib].get().as_usize();
self.descend_first(addr);
}
true
}
#[inline]
fn backtrack_to_next(&mut self) -> Option<(&[u8], &T)> {
loop {
let frame = self.stack.pop()?;
match frame {
Frame::Inode { encoded, mask, nib } => {
if self.push_next_child(encoded, mask, nib + 1) {
return self.current();
}
}
Frame::Fnode { .. } => {
continue;
}
}
}
}
pub fn current(&self) -> Option<(&[u8], &T)> {
let frame = self.stack.last()?;
match *frame {
Frame::Inode { encoded, nib, .. } => {
if nib == usize::MAX {
return None;
}
let phys_idx = encoded.as_usize();
let node = self.trie.inode(phys_idx);
if nib == TERMINAL_NIB {
let ki = node.leaf.get();
let (off, len, val) = self.trie.index[ki.as_usize()].as_ref().unwrap();
let off = off.get();
Some((&self.trie.buf[off..off + len.as_usize()], val))
} else if let Some(key_index) = node.leaf_key_index(nib) {
Some((self.trie.key_slice(key_index), &self.trie.index[key_index.as_usize()].as_ref().unwrap().2))
} else {
None
}
}
Frame::Fnode { arena_idx, pos } => {
let f = match &self.trie.arena[arena_idx.as_usize()] {
ArenaNode::Fnode(f) => f,
ArenaNode::Inode(_) => unreachable!("Fnode frame points at an Inode"),
};
let ki = f
.pos_key_index(pos)
.expect("current: Fnode frame positioned on a non-terminal");
Some((self.trie.key_slice(ki), &self.trie.index[ki.as_usize()].as_ref().unwrap().2))
}
}
}
pub fn current_index(&self) -> Option<usize> {
let frame = self.stack.last()?;
match *frame {
Frame::Inode { encoded, nib, .. } => {
if nib == usize::MAX {
return None;
}
let phys_idx = encoded.as_usize();
let node = self.trie.inode(phys_idx);
if nib == TERMINAL_NIB {
Some(node.leaf.get().as_usize())
} else {
node.leaf_key_index(nib).map(|ki| ki.as_usize())
}
}
Frame::Fnode { arena_idx, pos } => {
let f = match &self.trie.arena[arena_idx.as_usize()] {
ArenaNode::Fnode(f) => f,
ArenaNode::Inode(_) => unreachable!("Fnode frame points at an Inode"),
};
Some(
f.pos_key_index(pos)
.expect("current_index: Fnode frame positioned on a non-terminal")
.as_usize(),
)
}
}
}
#[inline]
fn advance_next(&mut self) -> bool {
loop {
let frame = match self.stack.pop() {
Some(v) => v,
None => return false,
};
match frame {
Frame::Inode { encoded, mask, nib } => {
if nib == TERMINAL_NIB {
if self.push_next_child(encoded, mask, 0) {
return true;
}
continue;
}
let search_start = if nib == usize::MAX { 0 } else { nib + 1 };
if self.push_next_child(encoded, mask, search_start) {
return true;
}
continue;
}
Frame::Fnode { arena_idx, pos } => {
let next_pos = match &self.trie.arena[arena_idx.as_usize()] {
ArenaNode::Fnode(f) => f.next_terminal_pos(pos),
ArenaNode::Inode(_) => unreachable!("Fnode frame points at an Inode"),
};
if let Some(np) = next_pos {
self.stack.push(Frame::Fnode { arena_idx, pos: np });
return true;
}
continue;
}
}
}
}
#[inline]
pub fn next(&mut self) -> Option<(&[u8], &T)> {
if self.advance_next() { self.current() } else { None }
}
fn fnode_seek(&mut self, arena_idx: usize, key: &[u8], _max_nib: usize) -> Option<(&[u8], &T)> {
let found_pos: Option<usize> = {
let f = match &self.trie.arena[arena_idx] {
ArenaNode::Fnode(f) => f,
ArenaNode::Inode(_) => unreachable!("fnode_seek on an Inode"),
};
if f.terminal && self.trie.key_slice(f.base) >= key {
Some(0)
} else {
let slots = f.slots.as_slice();
let base = f.base.as_usize();
let mut found = None;
for (i, (_plen, offset)) in slots.iter().enumerate() {
if *offset == FNODE_OFFSET_NULL {
continue;
}
let ki = PTR::from_usize(base + *offset as usize);
if self.trie.key_slice(ki) >= key {
found = Some(i + 1);
break;
}
}
found
}
};
if let Some(pos) = found_pos {
self.stack.push(Frame::Fnode { arena_idx: PTR::from_usize(arena_idx), pos });
return self.current();
}
match self.stack.pop() {
Some(Frame::Inode { encoded, mask, nib }) => {
if self.push_next_child(encoded, mask, nib + 1) {
return self.current();
}
self.backtrack_to_next()
}
other => {
if let Some(frm) = other {
self.stack.push(frm);
}
None
}
}
}
pub fn seek(&mut self, key: &[u8]) -> Option<(&[u8], &T)> {
if self.trie.arena.is_empty() {
self.stack.clear();
return None;
}
self.stack.clear();
let mut phys_idx: usize = 0;
let max_nib = key.len() * 2;
loop {
let node = *self.trie.inode(phys_idx);
let mask = node.children_mask();
if node.is_terminal() && node.prefix_len.as_usize() >= max_nib {
let ki = node.leaf.get();
let (off, len, _) = self.trie.index[ki.as_usize()].as_ref().unwrap();
let off = off.get();
let node_key = &self.trie.buf[off..off + len.as_usize()];
if node_key >= key {
self.stack.push(Frame::Inode { encoded: PTR::from_usize(phys_idx), mask, nib: TERMINAL_NIB });
return self.current();
}
}
if node.prefix_len.as_usize() >= max_nib {
if self.push_next_child(PTR::from_usize(phys_idx), mask, 0) {
return self.current();
}
return self.backtrack_to_next();
}
let nib = key_nibble_at(key, node.prefix_len.as_usize()) as usize;
if !node.is_occupied(nib) {
if self.push_next_child(PTR::from_usize(phys_idx), mask, nib + 1) {
return self.current();
}
return self.backtrack_to_next();
}
self.stack.push(Frame::Inode { encoded: PTR::from_usize(phys_idx), mask, nib });
if node.is_leaf(nib) {
let leaf_key = self.trie.key_slice(node.children[nib].get());
if leaf_key >= key {
return self.current();
}
return self.next();
} else {
let child_addr = node.children[nib].get().as_usize();
if matches!(self.trie.arena[child_addr], ArenaNode::Fnode(_)) {
return self.fnode_seek(child_addr, key, max_nib);
}
phys_idx = child_addr;
}
}
}
}
pub struct Cursor<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> {
trie: &'a NibbleTrie<K, T, PTR, LEN>,
pos: usize,
cur: Option<(&'a [u8], &'a T)>,
}
impl<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> Cursor<'a, K, T, PTR, LEN> {
#[inline]
fn park_slot(&mut self, pos: usize, slot: &'a Slot<LEN, T>) {
self.pos = pos;
let off = slot.0.get();
let klen = slot.1.as_usize();
self.cur = Some((&self.trie.buf[off..off + klen], &slot.2));
}
#[inline]
fn park_sentinel(&mut self, pos: usize) {
self.pos = pos;
self.cur = None;
}
pub fn new(trie: &'a NibbleTrie<K, T, PTR, LEN>) -> Self {
Cursor { trie, pos: 0, cur: None }
}
pub fn new_last(trie: &'a NibbleTrie<K, T, PTR, LEN>) -> Self {
let mut c = Cursor { trie, pos: 0, cur: None };
c.last();
c
}
pub fn first(&mut self) -> Option<(K::Borrowed<'a>, &'a T)> {
let len = self.trie.index.len();
let mut i = 1;
while i < len {
if let Some(slot) = self.trie.index[i].as_ref() {
self.park_slot(i, slot);
return self.cur.map(|(k, v)| (K::as_borrowed(k), v));
}
i += 1;
}
self.park_sentinel(0);
None
}
pub fn last(&mut self) -> Option<(K::Borrowed<'a>, &'a T)> {
let mut i = self.trie.index.len();
while i > 1 {
i -= 1;
if let Some(slot) = self.trie.index[i].as_ref() {
self.park_slot(i, slot);
return self.cur.map(|(k, v)| (K::as_borrowed(k), v));
}
}
self.park_sentinel(0);
None
}
#[inline]
pub fn current(&self) -> Option<(K::Borrowed<'a>, &'a T)> {
self.cur.map(|(k, v)| (K::as_borrowed(k), v))
}
#[inline]
pub fn current_index(&self) -> Option<usize> {
if self.cur.is_some() { Some(self.pos) } else { None }
}
#[inline]
pub fn next(&mut self) -> Option<(K::Borrowed<'a>, &'a T)> {
if self.advance_next() {
self.cur.map(|(k, v)| (K::as_borrowed(k), v))
} else {
None
}
}
#[inline]
pub fn prev(&mut self) -> Option<(K::Borrowed<'a>, &'a T)> {
if self.advance_prev() {
self.cur.map(|(k, v)| (K::as_borrowed(k), v))
} else {
None
}
}
#[inline]
pub fn next_index(&mut self) -> Option<usize> {
if self.advance_next() { Some(self.pos) } else { None }
}
#[inline]
pub fn prev_index(&mut self) -> Option<usize> {
if self.advance_prev() { Some(self.pos) } else { None }
}
pub fn seek(&mut self, key: &[u8]) -> Option<(K::Borrowed<'a>, &'a T)> {
let pos = {
let mut w = self.trie.walk_iter();
w.seek(key);
w.current_index()
};
match pos {
Some(p) => {
if let Some(slot) = self.trie.index[p].as_ref() {
self.park_slot(p, slot);
self.cur.map(|(k, v)| (K::as_borrowed(k), v))
} else {
self.park_sentinel(self.trie.index.len());
None
}
}
None => { self.park_sentinel(self.trie.index.len()); None }
}
}
#[inline]
fn advance_next(&mut self) -> bool {
let len = self.trie.index.len();
let mut i = self.pos + 1;
while i < len {
if let Some(slot) = self.trie.index[i].as_ref() {
self.park_slot(i, slot);
return true;
}
i += 1;
}
self.park_sentinel(len);
false
}
#[inline]
fn advance_prev(&mut self) -> bool {
let mut i = self.pos;
while i > 1 {
i -= 1;
if let Some(slot) = self.trie.index[i].as_ref() {
self.park_slot(i, slot);
return true;
}
}
self.park_sentinel(0);
false
}
}
pub struct CursorMut<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> {
trie: &'a mut NibbleTrie<K, T, PTR, LEN>,
pos: usize,
}
impl<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> CursorMut<'a, K, T, PTR, LEN> {
pub fn new(trie: &'a mut NibbleTrie<K, T, PTR, LEN>) -> Self {
CursorMut { trie, pos: 0 }
}
pub fn new_last(trie: &'a mut NibbleTrie<K, T, PTR, LEN>) -> Self {
let mut c = CursorMut { trie, pos: 0 };
c.last();
c
}
#[inline]
fn materialize<'b>(&'b mut self) -> Option<(K::Borrowed<'b>, &'b mut T)> {
let pos = self.pos;
let (off, len) = {
let slot = self.trie.index[pos].as_ref()?;
(slot.0.get(), slot.1.as_usize())
};
let k = K::as_borrowed(&self.trie.buf[off..off + len]);
let slot = self.trie.index[pos].as_mut()?;
Some((k, &mut slot.2))
}
pub fn first(&mut self) -> Option<(K::Borrowed<'_>, &mut T)> {
self.pos = 0;
if self.advance_next() { self.materialize() } else { None }
}
pub fn last(&mut self) -> Option<(K::Borrowed<'_>, &mut T)> {
self.pos = self.trie.index.len();
if self.advance_prev() { self.materialize() } else { None }
}
#[inline]
pub fn current(&mut self) -> Option<(K::Borrowed<'_>, &mut T)> {
let len = self.trie.index.len();
if self.pos == 0 || self.pos >= len {
return None;
}
self.materialize()
}
#[inline]
pub fn current_index(&self) -> Option<usize> {
let len = self.trie.index.len();
if self.pos != 0 && self.pos < len { Some(self.pos) } else { None }
}
#[inline]
pub fn next(&mut self) -> Option<(K::Borrowed<'_>, &mut T)> {
if self.advance_next() { self.materialize() } else { None }
}
#[inline]
pub fn prev(&mut self) -> Option<(K::Borrowed<'_>, &mut T)> {
if self.advance_prev() { self.materialize() } else { None }
}
#[inline]
pub fn next_index(&mut self) -> Option<usize> {
if self.advance_next() { Some(self.pos) } else { None }
}
#[inline]
pub fn prev_index(&mut self) -> Option<usize> {
if self.advance_prev() { Some(self.pos) } else { None }
}
pub fn seek(&mut self, key: &[u8]) -> Option<(K::Borrowed<'_>, &mut T)> {
let pos = {
let trie = &*self.trie;
let mut w = trie.walk_iter();
w.seek(key);
w.current_index()
};
let len = self.trie.index.len();
match pos {
Some(p) if self.trie.index[p].is_some() => {
self.pos = p;
self.materialize()
}
_ => { self.pos = len; None }
}
}
#[inline]
fn advance_next(&mut self) -> bool {
let len = self.trie.index.len();
let mut i = self.pos + 1;
while i < len {
if self.trie.index[i].is_some() {
self.pos = i;
return true;
}
i += 1;
}
self.pos = len;
false
}
#[inline]
fn advance_prev(&mut self) -> bool {
let mut i = self.pos;
while i > 1 {
i -= 1;
if self.trie.index[i].is_some() {
self.pos = i;
return true;
}
}
self.pos = 0;
false
}
}
pub struct Range<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> {
trie: &'a NibbleTrie<K, T, PTR, LEN>,
pos: usize,
end_pos: usize,
}
impl<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> Range<'a, K, T, PTR, LEN> {
pub(crate) fn new(
trie: &'a NibbleTrie<K, T, PTR, LEN>,
start: Bound<&[u8]>,
end: Bound<&[u8]>,
) -> Self {
let len = trie.index.len();
let pos = match start {
Bound::Included(k) => ceiling_index(trie, k).unwrap_or(len),
Bound::Excluded(k) => ceiling_strict_index(trie, k).unwrap_or(len),
Bound::Unbounded => 0, };
let end_pos = match end {
Bound::Included(k) => ceiling_strict_index(trie, k).unwrap_or(len),
Bound::Excluded(k) => ceiling_index(trie, k).unwrap_or(len),
Bound::Unbounded => len,
};
Range { trie, pos, end_pos }
}
}
impl<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> Iterator for Range<'a, K, T, PTR, LEN> {
type Item = (K::Borrowed<'a>, &'a T);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let end = self.end_pos;
let mut i = self.pos;
while i < end {
if let Some(slot) = self.trie.index[i].as_ref() {
let off = slot.0.get();
let klen = slot.1.as_usize();
let k = K::as_borrowed(&self.trie.buf[off..off + klen]);
self.pos = i + 1;
return Some((k, &slot.2));
}
i += 1;
}
self.pos = end;
None
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = self.end_pos.saturating_sub(self.pos);
(0, Some(remaining))
}
}
impl<'a, K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex> DoubleEndedIterator
for Range<'a, K, T, PTR, LEN>
{
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
let start = self.pos;
let mut i = self.end_pos;
while i > start {
i -= 1;
if let Some(slot) = self.trie.index[i].as_ref() {
let off = slot.0.get();
let klen = slot.1.as_usize();
let k = K::as_borrowed(&self.trie.buf[off..off + klen]);
self.end_pos = i;
return Some((k, &slot.2));
}
}
self.end_pos = start;
None
}
}
fn ceiling_index<K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex>(
trie: &NibbleTrie<K, T, PTR, LEN>,
key: &[u8],
) -> Option<usize> {
let mut w = trie.walk_iter();
w.seek(key);
w.current_index()
}
fn ceiling_strict_index<K: ByteKey, T, PTR: TrieIndex, LEN: TrieIndex>(
trie: &NibbleTrie<K, T, PTR, LEN>,
key: &[u8],
) -> Option<usize> {
let p = ceiling_index(trie, key)?;
let slot = trie.index[p].as_ref()?;
let off = slot.0.get();
let klen = slot.1.as_usize();
if &trie.buf[off..off + klen] == key {
let len = trie.index.len();
let mut i = p + 1;
while i < len {
if trie.index[i].is_some() {
return Some(i);
}
i += 1;
}
None
} else {
Some(p)
}
}
#[cfg(test)]
#[path = "tests/nibble_trie.rs"]
mod tests;