reth_trie_sparse/arena/
branch_child_idx.rs

1use alloy_trie::{TrieMask, TrieMaskIter};
2use core::{
3    iter::Enumerate,
4    ops::{Index, IndexMut},
5};
6use smallvec::SmallVec;
7
8/// A dense index into a branch node's children array.
9///
10/// Branch nodes store children densely — only occupied nibble slots have entries. This type
11/// wraps the `u8` index into that dense array, providing safe construction from a
12/// `(TrieMask, nibble)` pair and ergonomic indexing into `SmallVec` or slices.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub(super) struct BranchChildIdx(u8);
15
16impl BranchChildIdx {
17    /// Returns the dense index for `nibble` within the children array of a branch whose
18    /// occupied slots are described by `state_mask`.
19    ///
20    /// Returns `None` if the nibble's bit is not set in `state_mask`.
21    pub(super) const fn new(state_mask: TrieMask, nibble: u8) -> Option<Self> {
22        if !state_mask.is_bit_set(nibble) {
23            return None;
24        }
25        Some(Self::new_unchecked(state_mask, nibble))
26    }
27
28    /// Returns the dense insertion point for `nibble` — the number of occupied child slots
29    /// below `nibble`. Unlike [`Self::new`], this does **not** require the nibble's bit to be
30    /// set, making it suitable for computing the position at which a new child should be
31    /// inserted.
32    pub(super) const fn insertion_point(state_mask: TrieMask, nibble: u8) -> Self {
33        Self(Self::count_below(state_mask, nibble))
34    }
35
36    /// Returns the dense index as a `usize`, suitable for indexing into a `SmallVec` or slice.
37    pub(super) const fn get(self) -> usize {
38        self.0 as usize
39    }
40
41    /// Counts the number of occupied child slots below `nibble` in the dense children array.
42    const fn count_below(state_mask: TrieMask, nibble: u8) -> u8 {
43        (state_mask.get() & ((1u16 << nibble) - 1)).count_ones() as u8
44    }
45
46    /// Computes the dense index for `nibble` without checking whether the bit is set.
47    const fn new_unchecked(state_mask: TrieMask, nibble: u8) -> Self {
48        Self(Self::count_below(state_mask, nibble))
49    }
50}
51
52impl<T> Index<BranchChildIdx> for SmallVec<[T; 4]> {
53    type Output = T;
54
55    fn index(&self, idx: BranchChildIdx) -> &Self::Output {
56        &self.as_slice()[idx.get()]
57    }
58}
59
60impl<T> IndexMut<BranchChildIdx> for SmallVec<[T; 4]> {
61    fn index_mut(&mut self, idx: BranchChildIdx) -> &mut Self::Output {
62        &mut self.as_mut_slice()[idx.get()]
63    }
64}
65
66/// An iterator over a branch's children that yields `(BranchChildIdx, nibble)` pairs.
67///
68/// Wraps `TrieMask::iter().enumerate()` to produce [`BranchChildIdx`] values instead of raw
69/// `usize` indices.
70pub(super) struct BranchChildIter {
71    inner: Enumerate<TrieMaskIter>,
72}
73
74impl BranchChildIter {
75    /// Creates a new iterator over the occupied children of the given `state_mask`.
76    pub(super) fn new(state_mask: TrieMask) -> Self {
77        Self { inner: state_mask.iter().enumerate() }
78    }
79}
80
81impl Iterator for BranchChildIter {
82    type Item = (BranchChildIdx, u8);
83
84    fn next(&mut self) -> Option<Self::Item> {
85        self.inner.next().map(|(dense, nibble)| (BranchChildIdx(dense as u8), nibble))
86    }
87
88    fn size_hint(&self) -> (usize, Option<usize>) {
89        self.inner.size_hint()
90    }
91}
reth_trie_sparse/arena/branch_child_idx.rs

reth_trie_sparse/arena/
branch_child_idx.rs
