ts_bart/node/
mod.rs

1//! Single-level trie nodes and operations.
2//!
3//! Separate from [`iptrie`][crate::iptrie] to maintain a clear dependency
4//! ordering: the multi-level trie operations are strictly dependent on the
5//! single-level ops.
6
7use core::{
8    borrow::Borrow,
9    fmt::{Debug, Formatter},
10};
11
12use ts_array256::{Array256, ArrayStorage};
13use ts_bitset::Bitset256;
14
15use crate::BaseIndex;
16
17mod child;
18mod child_storage;
19mod descendants;
20mod stride_ops;
21
22pub use child::Child;
23pub use child_storage::{BoxStorage, InlineStorage, Storage};
24pub use stride_ops::{
25    NodePrefixIter, PrefixOps, PrefixOpsExt, PrefixReadOps, Stats, StrideBase, StrideOps,
26    StrideOpsExt,
27};
28
29/// Type alias defining the default child storage type.
30///
31/// This storage type is considered the best for most applications based on
32/// benchmarks and storage size profiling.
33pub type DefaultStorage = BoxStorage;
34
35/// Type alias for [`Node`] with [`DefaultStorage`]. Provided to improve
36/// inference by concretizing the storage type.
37pub type DefaultNode<T> = Node<T, DefaultStorage>;
38
39cfg_if::cfg_if! {
40    if #[cfg(feature = "smallvec")] {
41        /// Storage for elements in [`Node`] child and prefix arrays.
42        /// [`SmallVec`][smallvec::SmallVec] wins marginally in memory size according
43        /// to benchmarks, improves majorly on mutation performance (avoiding allocations,
44        /// presumably), and has mixed effects on lookup performance.
45        type ArrayStore<T> = smallvec::SmallVec<[T; 1]>;
46    } else {
47        /// Storage for elements in [`Node`] child and prefix arrays.
48        type ArrayStore<T> = alloc::vec::Vec<T>;
49    }
50}
51
52/// A trie node logically covering a single (octet) address stride.
53///
54/// Contains prefixes (directly owned by this node) and children (descendants
55/// in the trie).
56///
57/// This is the canonical stride-level [`Node`] type for the crate, mirroring
58/// go-bart's Full nodes. The operations its supports are abstracted for
59/// via the [`StrideOps`] trait -- other implementations are possible, e.g. to
60/// support similar functionality to go-bart's `Lite` and `Fast` node types.
61#[derive(PartialEq, Eq, Hash)]
62pub struct Node<T, C = DefaultStorage>
63where
64    C: Storage + ?Sized,
65{
66    /// The prefixes this node covers directly, indexed by [`BaseIndex`].
67    ///
68    /// If this node is at trie depth 3 via octet path `[192, 168]` for
69    /// instance, the prefix entry `12/5` => `34` would represent the overall
70    /// trie route mapping `192.168.12.0/21` => 34.
71    pub prefixes: Array256<ArrayStore<T>>,
72
73    /// The children contained in this node, indexed by complete prefix
74    /// octet.
75    pub children: Array256<ArrayStore<Child<C::Container<Self>, T>>>,
76}
77
78impl<T, C> Clone for Node<T, C>
79where
80    C: Storage + ?Sized,
81    T: Clone,
82{
83    #[inline]
84    fn clone(&self) -> Self {
85        Self {
86            prefixes: self.prefixes.clone(),
87            children: self.children.clone_with(&|c| match c {
88                Child::Path(p) => Child::Path(C::clone(p)),
89                Child::Fringe(p) => Child::Fringe(p.clone()),
90                Child::Leaf { prefix, value } => Child::Leaf {
91                    prefix: *prefix,
92                    value: value.clone(),
93                },
94            }),
95        }
96    }
97}
98
99impl<T, C> Debug for Node<T, C>
100where
101    T: Debug,
102    C: Storage + ?Sized,
103{
104    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
105        struct IdArrayFormatter<'a, S> {
106            ary: &'a Array256<S>,
107        }
108
109        impl<S> Debug for IdArrayFormatter<'_, S>
110        where
111            S: ArrayStorage + AsRef<[S::T]>,
112            S::T: Debug,
113        {
114            fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
115                f.debug_map()
116                    .entries(self.ary.bitset().bits().map(|i| {
117                        (
118                            BaseIndex::new(i as _).fmt_prefix(),
119                            self.ary.get(i as _).unwrap(),
120                        )
121                    }))
122                    .finish()
123            }
124        }
125
126        f.debug_struct("Node")
127            .field(
128                "prefixes",
129                &IdArrayFormatter {
130                    ary: &self.prefixes,
131                },
132            )
133            .field(
134                "children",
135                &self
136                    .children
137                    .custom_storage_fmt(&Child::as_node_ref::<C, _>),
138            )
139            .finish()
140    }
141}
142
143impl<T, C> Default for Node<T, C>
144where
145    C: Storage + ?Sized,
146{
147    #[inline]
148    fn default() -> Self {
149        Self::EMPTY
150    }
151}
152
153impl<T, C> Node<T, C>
154where
155    C: Storage + ?Sized,
156{
157    /// The empty node value (no children, no prefixes).
158    pub const EMPTY: Self = Self {
159        prefixes: Array256::EMPTY,
160        children: Array256::EMPTY,
161    };
162
163    /// Returns whether the node is empty.
164    #[inline]
165    pub const fn is_empty(&self) -> bool {
166        self.prefixes.is_empty() && self.children.is_empty()
167    }
168
169    /// Return usage statistics for the trie rooted at this node.
170    pub fn stats(&self) -> Stats
171    where
172        Self: StrideOps,
173    {
174        self.descendant_nodes(true)
175            .fold(Stats::default(), |mut stats, (_path, node)| {
176                let (leaves, fringes) = node.direct_children().fold(
177                    (0, 0),
178                    |(leaves, fringes), (_, child)| match child {
179                        Child::Fringe(..) => (leaves, fringes + 1),
180                        Child::Leaf { .. } => (leaves + 1, fringes),
181                        _ => (leaves, fringes),
182                    },
183                );
184
185                stats.node_count += 1;
186
187                stats.prefix_count += node.prefix_count();
188                stats.child_count += node.child_count();
189
190                stats.leaf_count += leaves;
191                stats.fringe_count += fringes;
192
193                stats
194            })
195    }
196}
197
198impl<T, C> StrideBase for Node<T, C>
199where
200    T: 'static,
201    C: Storage + ?Sized,
202{
203    type T = T;
204}
205
206impl<T, C> PrefixReadOps for Node<T, C>
207where
208    T: 'static,
209    C: Storage + ?Sized,
210{
211    fn prefix_bitset(&self) -> &Bitset256 {
212        self.prefixes.bitset()
213    }
214
215    fn prefix_count(&self) -> usize {
216        self.prefixes.len()
217    }
218
219    fn lookup_index(&self, idx: BaseIndex) -> Option<(BaseIndex, &Self::T)> {
220        let top = self.prefixes.intersection_top(crate::lpm(idx).borrow())?;
221        Some((BaseIndex::try_new(top)?, self.prefixes.get(top)?))
222    }
223
224    fn get_prefix_exact(&self, idx: BaseIndex) -> Option<&Self::T> {
225        self.prefixes.get(idx.get())
226    }
227}
228
229impl<T, C> PrefixOps for Node<T, C>
230where
231    T: 'static,
232    C: Storage + ?Sized,
233{
234    fn insert_prefix(&mut self, idx: BaseIndex, value: Self::T) -> Option<Self::T> {
235        self.prefixes.insert(idx.into(), value)
236    }
237
238    fn remove_prefix(&mut self, idx: BaseIndex) -> Option<Self::T> {
239        self.prefixes.remove(idx.into())
240    }
241
242    fn get_prefix_exact_mut(&mut self, idx: BaseIndex) -> Option<&mut Self::T> {
243        self.prefixes.get_mut(idx.get())
244    }
245}
246
247impl<T, C> StrideOps for Node<T, C>
248where
249    T: 'static,
250    C: Storage + ?Sized,
251{
252    fn child_bitset(&self) -> &Bitset256 {
253        self.children.bitset()
254    }
255
256    fn child_count(&self) -> usize {
257        self.children.len()
258    }
259
260    fn stats(&self) -> Stats {
261        self.stats()
262    }
263
264    fn insert_child(
265        &mut self,
266        addr: u8,
267        child: Child<Self, Self::T>,
268    ) -> Option<Child<Self, Self::T>> {
269        self.children
270            .insert(addr, child.map_node(C::new))
271            .map(|ret| ret.map_node(C::into_inner))
272    }
273
274    fn direct_children(&self) -> impl Iterator<Item = (u8, Child<&Self, &T>)> {
275        self.children
276            .iter()
277            .map(|(addr, child)| (addr, child.as_node_ref::<C, _>()))
278    }
279
280    fn remove_child(&mut self, addr: u8) -> Option<Child<Self, Self::T>> {
281        self.children
282            .remove(addr)
283            .map(|ret| ret.map_node(C::into_inner))
284    }
285
286    fn get_child(&self, addr: u8) -> Option<Child<&Self, &Self::T>> {
287        self.children.get(addr).map(Child::as_node_ref::<C, _>)
288    }
289
290    fn get_child_mut(&mut self, addr: u8) -> Option<Child<&mut Self, &mut Self::T>> {
291        self.children.get_mut(addr).map(Child::as_node_mut::<C, _>)
292    }
293
294    fn direct_prefixes(&self) -> impl Iterator<Item = (BaseIndex, &T)> {
295        self.prefixes
296            .iter()
297            .map(|(idx, value)| (BaseIndex::new(idx), value))
298    }
299}
300
301#[cfg(test)]
302mod test {
303    use super::*;
304
305    #[test]
306    fn bart_examples_prefix_crud() {
307        let mut node = Node::<usize>::EMPTY;
308        let index = BaseIndex::new(32);
309
310        assert_eq!(None, node.insert_prefix(index, 100));
311        assert_eq!(1, node.prefix_count());
312
313        assert_eq!(Some(100), node.insert_prefix(index, 111));
314        assert_eq!(1, node.prefix_count());
315        assert_eq!(Some(111), node.lookup(index).copied());
316
317        assert_eq!(Some(111), node.remove_prefix(index));
318        assert_eq!(0, node.prefix_count());
319        assert_eq!(None, node.lookup(index));
320        assert_eq!(None, node.remove_prefix(index));
321    }
322
323    #[test]
324    fn bart_examples_contains() {
325        let mut node = Node::<()>::EMPTY;
326        node.insert_prefix(BaseIndex::new(32), ());
327
328        for idx in [32, 64, 65, 128, 129, 130, 131] {
329            assert!(node.supersets_prefix(BaseIndex::new(idx)));
330        }
331
332        for idx in [1, 16, 33, 63, 127, 132, 255] {
333            assert!(!node.supersets_prefix(BaseIndex::new(idx)));
334        }
335    }
336
337    #[test]
338    fn bart_examples_lookup() {
339        let mut node = Node::<usize>::EMPTY;
340
341        node.insert_prefix(BaseIndex::new(32), 1);
342        node.insert_prefix(BaseIndex::new(64), 2);
343
344        assert_eq!(Some(&2), node.lookup(BaseIndex::new(128)));
345        assert_eq!(
346            Some((BaseIndex::new(64), &2)),
347            node.lookup_index(BaseIndex::new(128))
348        );
349        assert_eq!(None, node.lookup(BaseIndex::new(127)));
350    }
351
352    #[test]
353    fn bart_examples_children_crud() {
354        let mut child = Node::<usize>::EMPTY;
355        child.insert_prefix(BaseIndex::new(1), 10);
356
357        let mut node = Node::<usize>::EMPTY;
358        assert!(node.insert_child(10, Child::Path(child)).is_none());
359        assert_eq!(1, node.child_count());
360
361        assert!(node.get_child(10).is_some_and(|c| match c {
362            Child::Path(inner) => inner.supersets_prefix(BaseIndex::new(1)),
363            _ => false,
364        }));
365
366        assert!(node.remove_child(10).is_some());
367        assert_eq!(0, node.child_count());
368        assert!(node.remove_child(10).is_none());
369    }
370}
ts_bart/node/mod.rs

ts_bart/node/
mod.rs
