Skip to main content

ts_bart/node/stride_ops/
mod.rs

1use ts_bitset::Bitset256;
2
3use crate::{BaseIndex, node::Child};
4
5mod prefix;
6
7pub use prefix::{NodePrefixIter, PrefixOps, PrefixOpsExt, PrefixReadOps};
8
9/// Stats about a node and its descendants.
10#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
11pub struct Stats {
12    /// The total number of direct prefixes: does not count leaves and fringes.
13    pub prefix_count: usize,
14    /// The total number of parent-child node relations.
15    pub child_count: usize,
16    /// The total number of full nodes in the trie.
17    pub node_count: usize,
18    /// The total number of leaf nodes.
19    pub leaf_count: usize,
20    /// The total number of fringe nodes.
21    pub fringe_count: usize,
22}
23
24/// Base trait for stride operations.
25pub trait StrideBase {
26    /// The kind of value held inside nodes.
27    type T: 'static;
28}
29
30impl<T> StrideBase for &T
31where
32    T: StrideBase + ?Sized,
33{
34    type T = T::T;
35}
36
37impl<T> StrideBase for &mut T
38where
39    T: StrideBase + ?Sized,
40{
41    type T = T::T;
42}
43
44/// Single-stride operations supported by a trie node.
45pub trait StrideOps: Default + PrefixOps {
46    // The `Default` bound is needed to construct empty values of this type for trie
47    // operations.
48
49    /// Iterate the prefixes directly contained in this node.
50    fn direct_prefixes(&self) -> impl Iterator<Item = (BaseIndex, &Self::T)>;
51
52    /// Get a reference to the child bitset.
53    fn child_bitset(&self) -> &Bitset256;
54
55    /// Get a child node by address.
56    fn get_child(&self, addr: u8) -> Option<Child<&Self, &Self::T>>;
57    /// Get a mutable reference to a child node by address.
58    fn get_child_mut(&mut self, addr: u8) -> Option<Child<&mut Self, &mut Self::T>>;
59    /// Insert a child node at the given address.
60    fn insert_child(
61        &mut self,
62        addr: u8,
63        child: Child<Self, Self::T>,
64    ) -> Option<Child<Self, Self::T>>;
65    /// Remove a child node at the given address.
66    fn remove_child(&mut self, addr: u8) -> Option<Child<Self, Self::T>>;
67    /// Iterate the direct children of this node.
68    fn direct_children(&self) -> impl Iterator<Item = (u8, Child<&Self, &Self::T>)>;
69
70    /// Get the number of direct children of this node.
71    #[inline]
72    fn child_count(&self) -> usize {
73        self.child_bitset().count_ones()
74    }
75
76    /// Calculate trie occupancy stats for this node and its descendants.
77    fn stats(&self) -> Stats;
78}
79
80mod private {
81    pub trait Sealed {}
82}
83
84/// Extension methods for nodes implementing [`StrideOps`].
85pub trait StrideOpsExt: StrideOps + private::Sealed {
86    /// Report whether the node has no children and prefixes.
87    #[inline]
88    fn is_empty(&self) -> bool {
89        self.prefix_count() == 0 && self.child_count() == 0
90    }
91
92    /// Return this node with the given child added.
93    ///
94    /// Meant as sugar for easily constructing nodes directly.
95    ///
96    /// # Examples
97    ///
98    /// ```rust
99    /// # use ts_bart::{DefaultNode, Child, StrideOps, StrideOpsExt};
100    /// DefaultNode::EMPTY.with_child(
101    ///     0,
102    ///     DefaultNode::EMPTY
103    ///         .with_child(1, Child::Fringe(123))
104    ///         .into_child(),
105    /// );
106    /// ```
107    #[inline]
108    fn with_child(mut self, addr: u8, child: impl Into<Child<Self, Self::T>>) -> Self {
109        self.insert_child(addr, child.into());
110        self
111    }
112
113    /// Wrap this node in a [`Child::Path`].
114    #[inline]
115    fn into_child(self) -> Child<Self, Self::T> {
116        Child::Path(self)
117    }
118}
119
120impl<T> private::Sealed for T where T: StrideOps {}
121impl<T> StrideOpsExt for T where T: StrideOps {}