ts_bart/node/stride_ops/
prefix.rs

1use ts_bitset::Bitset256;
2
3use crate::{BaseIndex, StrideBase};
4
5/// Trie node operations for accessing and mutating prefixes.
6pub trait PrefixOps: PrefixReadOps {
7    /// Insert a value at the given base index.
8    fn insert_prefix(&mut self, idx: BaseIndex, value: Self::T) -> Option<Self::T>;
9    /// Remove the value for the given base index.
10    fn remove_prefix(&mut self, idx: BaseIndex) -> Option<Self::T>;
11    /// Get a mutable reference to the prefix value that matches the given
12    /// index exactly.
13    fn get_prefix_exact_mut(&mut self, idx: BaseIndex) -> Option<&mut Self::T>;
14}
15
16static_assertions::assert_obj_safe!(PrefixOps<T = ()>);
17
18impl<T> PrefixOps for &mut T
19where
20    T: PrefixOps + ?Sized,
21{
22    fn insert_prefix(&mut self, idx: BaseIndex, value: Self::T) -> Option<Self::T> {
23        (*self).insert_prefix(idx, value)
24    }
25
26    fn remove_prefix(&mut self, idx: BaseIndex) -> Option<Self::T> {
27        (*self).remove_prefix(idx)
28    }
29
30    fn get_prefix_exact_mut(&mut self, idx: BaseIndex) -> Option<&mut Self::T> {
31        (*self).get_prefix_exact_mut(idx)
32    }
33}
34
35/// Read-only operations on prefixes stored in nodes.
36pub trait PrefixReadOps: StrideBase {
37    /// Get a reference to the prefix bitset.
38    fn prefix_bitset(&self) -> &Bitset256;
39    /// Lookup the prefix value that covers the given index.
40    fn lookup_index(&self, idx: BaseIndex) -> Option<(BaseIndex, &Self::T)>;
41    /// Get a reference to the prefix value that matches the given index
42    /// exactly.
43    fn get_prefix_exact(&self, idx: BaseIndex) -> Option<&Self::T>;
44
45    /// Get the number of prefixes in this node.
46    fn prefix_count(&self) -> usize {
47        self.prefix_bitset().count_ones()
48    }
49}
50
51static_assertions::assert_obj_safe!(PrefixReadOps<T = ()>);
52
53impl<T> PrefixReadOps for &T
54where
55    T: PrefixReadOps + ?Sized,
56{
57    fn prefix_bitset(&self) -> &Bitset256 {
58        (*self).prefix_bitset()
59    }
60
61    fn lookup_index(&self, idx: BaseIndex) -> Option<(BaseIndex, &Self::T)> {
62        (*self).lookup_index(idx)
63    }
64
65    fn get_prefix_exact(&self, idx: BaseIndex) -> Option<&Self::T> {
66        (*self).get_prefix_exact(idx)
67    }
68}
69
70impl<T> PrefixReadOps for &mut T
71where
72    T: PrefixReadOps + ?Sized,
73{
74    fn prefix_bitset(&self) -> &Bitset256 {
75        (**self).prefix_bitset()
76    }
77
78    fn lookup_index(&self, idx: BaseIndex) -> Option<(BaseIndex, &Self::T)> {
79        (**self).lookup_index(idx)
80    }
81
82    fn get_prefix_exact(&self, idx: BaseIndex) -> Option<&Self::T> {
83        (**self).get_prefix_exact(idx)
84    }
85}
86
87/// Extension methods relating to prefixes.
88pub trait PrefixOpsExt: PrefixOps {
89    /// Report whether this node's prefix set covers the prefix with the given
90    /// index. It does not need to match exactly, `idx` just needs to be
91    /// contained in this node's prefix set.
92    ///
93    /// # Examples
94    ///
95    /// ```rust
96    /// # use ts_bart::{DefaultNode, BaseIndex, PrefixOps, PrefixOpsExt};
97    /// let idx = BaseIndex::from_prefix(0, 2);
98    /// let node = DefaultNode::EMPTY.with_prefix(idx, 32);
99    /// assert!(node.supersets_prefix(idx)); // exact match
100    /// // parent is not contained
101    /// assert!(!node.supersets_prefix(idx.parent().unwrap()));
102    /// // child indexes are contained
103    /// let (child1, child2) = idx.children().unwrap();
104    /// assert!(node.supersets_prefix(child1));
105    /// assert!(node.supersets_prefix(child2));
106    /// ```
107    #[inline]
108    fn supersets_prefix(&self, idx: BaseIndex) -> bool {
109        use core::borrow::Borrow;
110        self.prefix_bitset().intersects(crate::lpm(idx).borrow())
111    }
112
113    /// Lookup a value by [`BaseIndex`].
114    ///
115    /// This is sugar over [`PrefixReadOps::lookup_index`] to only return the
116    /// matched value.
117    #[inline]
118    fn lookup(&self, idx: BaseIndex) -> Option<&Self::T> {
119        let (_idx, ret) = self.lookup_index(idx)?;
120        Some(ret)
121    }
122
123    /// Return this node with the given prefix added.
124    ///
125    /// Sugar for easily constructing nodes directly.
126    ///
127    /// # Examples
128    ///
129    /// ```rust
130    /// # use ts_bart::{DefaultNode, BaseIndex, PrefixOps, PrefixReadOps, PrefixOpsExt};
131    /// let idx = BaseIndex::from_prefix(1, 1);
132    /// let node = DefaultNode::EMPTY.with_prefix(idx, 12);
133    /// assert_eq!(node.get_prefix_exact(idx).copied(), Some(12));
134    /// ```
135    #[inline]
136    fn with_prefix(mut self, idx: BaseIndex, value: Self::T) -> Self
137    where
138        Self: Sized,
139    {
140        self.insert_prefix(idx, value);
141        self
142    }
143
144    /// Iterate prefixes in this node matching `octet`.
145    ///
146    /// The prefixes are returned in reverse order (most specific to least specific).
147    ///
148    /// # Examples
149    ///
150    /// ```rust
151    /// # use ts_bart::{BaseIndex, DefaultNode, PrefixOpsExt};
152    /// let zero_pfx = BaseIndex::from_prefix(0, 0);
153    /// let second_half_pfx = BaseIndex::from_prefix(128, 1);
154    ///
155    /// let node = ts_bart::DefaultNode::EMPTY
156    ///     .with_prefix(zero_pfx, 123)
157    ///     .with_prefix(second_half_pfx, 456);
158    ///
159    /// assert_eq!(vec![(zero_pfx, &123)], node.matching_prefixes(1).collect::<Vec<_>>());
160    /// assert_eq!(vec![(second_half_pfx, &456), (zero_pfx, &123)], node.matching_prefixes(200).collect::<Vec<_>>());
161    /// ```
162    fn matching_prefixes(&self, octet: u8) -> NodePrefixIter<'_, Self::T>
163    where
164        Self: Sized,
165    {
166        NodePrefixIter::for_octet(self, octet)
167    }
168}
169
170impl<T> PrefixOpsExt for T where T: PrefixOps + ?Sized {}
171
172/// Iterator for matching prefixes in a single node.
173pub struct NodePrefixIter<'n, T> {
174    node: &'n dyn PrefixReadOps<T = T>,
175    yield_state: Bitset256,
176}
177
178impl<'n, T> NodePrefixIter<'n, T>
179where
180    T: 'static,
181{
182    /// Construct a [`NodePrefixIter`] yielding the prefixes stored in `node` that
183    /// cover `octet`, yielded in reverse order (most to least specific).
184    pub fn for_octet(node: &'n dyn PrefixReadOps<T = T>, octet: u8) -> Self {
185        use core::borrow::Borrow;
186
187        let idx = BaseIndex::from_pfx_7(octet);
188
189        let mut yield_state = *node.prefix_bitset();
190        yield_state.intersect_inplace(crate::lpm(idx).borrow());
191
192        Self { node, yield_state }
193    }
194}
195
196impl<'n, T> Iterator for NodePrefixIter<'n, T>
197where
198    T: 'static,
199{
200    type Item = (BaseIndex, &'n T);
201
202    fn next(&mut self) -> Option<Self::Item> {
203        while let Some(bit) = self.yield_state.last_set() {
204            self.yield_state.clear(bit);
205
206            let index = BaseIndex::new(bit as _);
207            if let Some(val) = self.node.get_prefix_exact(index) {
208                return Some((index, val));
209            }
210        }
211
212        None
213    }
214}
215
216impl<T> core::iter::FusedIterator for NodePrefixIter<'_, T> where T: 'static {}
ts_bart/node/stride_ops/prefix.rs

ts_bart/node/stride_ops/
prefix.rs
