iptrie/
set.rs

1//! Generic prefix trie set structures
2use crate::prefix::*;
3use crate::trie::lctrie::LevelCompressedTrie;
4use crate::trie::patricia::RadixTrie;
5use std::num::NonZeroUsize;
6
7#[cfg(feature = "graphviz")]
8use crate::graphviz::DotWriter;
9use crate::trie::common::Leaf;
10#[cfg(feature = "graphviz")]
11use std::fmt::Display;
12
13/// A set of Ip prefixes based on a radix binary trie
14#[derive(Clone)]
15pub struct RTrieSet<P: IpPrefix>(pub(crate) RadixTrie<P, ()>);
16
17/// A set of Ip prefixes based on a level-compressed trie
18pub struct LCTrieSet<P: IpPrefix>(pub(crate) LevelCompressedTrie<P, ()>);
19
20impl<P: IpRootPrefix> RTrieSet<P> {
21    /// Creates a new set which contains the root prefix.
22    #[inline]
23    pub fn new() -> Self {
24        Self::with_capacity(1000)
25    }
26
27    /// Creates a new set with an initial capacity.
28    ///
29    /// The returned set already contains the root prefix.
30    #[inline]
31    pub fn with_capacity(capacity: usize) -> Self {
32        Self(RadixTrie::new((), capacity))
33    }
34}
35
36impl<P: IpPrefix> RTrieSet<P> {
37    /// Returns the size of the set.
38    ///
39    /// Notice that it never equals zero since the top prefix is
40    /// always present in the set.
41    #[inline]
42    #[allow(clippy::len_without_is_empty)]
43    pub fn len(&self) -> NonZeroUsize {
44        self.0.len()
45    }
46
47    /// Compress this Patricia trie in a LC-Trie.
48    ///
49    /// For lookup algorithms, a Patricia trie performs unit bit checking and LC-Trie
50    /// performs multi bits checking. So the last one is more performant but it
51    /// cannot be modified (no insertion or removal operations are provided).
52    #[inline]
53    pub fn compress(self) -> LCTrieSet<P> {
54        LCTrieSet(LevelCompressedTrie::new(self.0))
55    }
56
57    #[inline]
58    pub fn shrink_to_fit(&mut self) {
59        self.0.shrink_to_fit()
60    }
61
62    /// Inserts a new element in the set.
63    ///
64    /// If the specified element already exists in the set, `false` is returned.
65    ///
66    /// # Example
67    /// ```
68    /// #  use iptrie::*;
69    /// use std::net::Ipv4Addr;
70    ///
71    /// let addr = Ipv4Addr::new(1,1,1,1);
72    /// let mut trie = Ipv4RTrieSet::new();
73    ///
74    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
75    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
76    ///
77    /// assert_eq!( trie.insert(ip20), true);
78    /// assert_eq!( trie.insert(ip22), true);
79    /// assert_eq!( trie.insert(ip20), false);
80    /// ```
81    #[inline]
82    pub fn insert(&mut self, k: P) -> bool {
83        self.0.insert(k, ()).is_none()
84    }
85
86    /// Checks if an element is present (exact match).
87    ///
88    /// # Example
89    /// ```
90    /// #  use iptrie::*;
91    /// use std::net::Ipv4Addr;
92    ///
93    /// let addr = Ipv4Addr::new(1,1,1,1);
94    ///
95    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
96    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
97    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
98    ///
99    /// let trie = Ipv4RTrieSet::from_iter([ip20,ip24]);
100    ///
101    /// assert_eq!( trie.contains(&ip20), true);
102    /// assert_eq!( trie.contains(&ip22), false);
103    /// assert_eq!( trie.contains(&ip24), true);
104    /// ```
105    #[inline]
106    pub fn contains<Q>(&self, k: &Q) -> bool
107    where
108        Q: IpPrefix<Addr = P::Addr>,
109        P: IpPrefixCovering<Q>,
110    {
111        self.0.get(k).is_some()
112    }
113
114    /// Removes a previously inserted prefix (exact match).
115    ///
116    /// Returns `false` is the element was not present in the set
117    /// and `true` if the removal is effective.
118    ///
119    /// # Example
120    /// ```
121    /// #  use iptrie::*;
122    /// use std::net::Ipv4Addr;
123    ///
124    /// let addr = Ipv4Addr::new(1,1,1,1);
125    ///
126    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
127    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
128    ///
129    /// let mut trie = Ipv4RTrieSet::from_iter([ip20,ip22]);
130    ///
131    /// assert_eq!( trie.contains(&ip20), true);
132    /// assert_eq!(trie.remove(&ip20), true);
133    /// assert_eq!(trie.remove(&ip20), false);
134    ///
135    /// assert_eq!( trie.contains(&ip22), true);
136    /// assert_eq!( trie.contains(&ip20), false);
137    /// ```
138    #[inline]
139    pub fn remove<Q>(&mut self, k: &Q) -> bool
140    where
141        Q: IpPrefix<Addr = P::Addr>,
142        P: IpPrefixCovering<Q>,
143    {
144        self.0.remove(k).is_some()
145    }
146
147    /// Replace an existing prefix.
148    ///
149    /// Adds a prefix to the set, replacing the existing one, if any (exact match performed).
150    /// Returns the replaced value.
151    ///
152    /// # Example
153    /// ```
154    /// # use iptrie::*;
155    /// # use iptrie::set::*;
156    /// use std::net::Ipv4Addr;
157    /// use ipnet::Ipv4Net;
158    /// let mut trie = RTrieSet::new();
159    ///
160    /// let addr1 = Ipv4Addr::new(1,1,1,1);
161    /// let addr2 = Ipv4Addr::new(1,1,1,2);
162    /// let ip20 = Ipv4Net::new(addr1, 20).unwrap();
163    /// let ip20b = Ipv4Net::new(addr2, 20).unwrap();
164    ///
165    /// trie.insert(ip20);
166    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
167    ///
168    /// assert_eq!(trie.insert(ip20b), false);
169    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
170    ///
171    /// assert_eq!(trie.replace(ip20b), Some(ip20));
172    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.2/20".to_string());
173    /// ```
174    #[inline]
175    pub fn replace(&mut self, k: P) -> Option<P> {
176        self.0.replace(k, ()).map(|l| *l.prefix())
177    }
178
179    /// Gets the value associated with an exact match of the key.
180    ///
181    /// To access to the longest prefix match, use [`Self::lookup`].
182    ///
183    /// # Example
184    /// ```
185    /// # use iptrie::*;
186    /// # use iptrie::set::*;
187    /// use std::net::Ipv4Addr;
188    /// use ipnet::Ipv4Net;
189    /// let mut trie = RTrieSet::new();
190    ///
191    /// let addr1 = Ipv4Addr::new(1,1,1,1);
192    /// let addr2 = Ipv4Addr::new(1,1,1,2);
193    /// let ip20 = Ipv4Net::new(addr1, 20).unwrap();
194    /// let ip20b = Ipv4Net::new(addr2, 20).unwrap();
195    ///
196    /// trie.insert(ip20);
197    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
198    ///
199    /// trie.insert(ip20b);
200    /// assert_eq!(trie.get(&ip20).unwrap().to_string(), "1.1.1.1/20".to_string());
201    /// ```
202    #[inline]
203    pub fn get<Q>(&self, k: &Q) -> Option<&P>
204    where
205        Q: IpPrefix<Addr = P::Addr>,
206        P: IpPrefixCovering<Q>,
207    {
208        self.0.get(k).map(|(k, _)| k)
209    }
210
211    /// Gets the longest prefix which matches the given key.
212    ///
213    /// As the top prefix always matches, it never fails.
214    ///
215    /// To access to the exact prefix match, use [`Self::get`].
216    ///
217    /// # Example
218    /// ```
219    /// # use iptrie::*;
220    /// use std::net::Ipv4Addr;
221    /// let mut trie = Ipv4RTrieSet::new();
222    ///
223    /// let addr = Ipv4Addr::new(1,1,1,1);
224    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
225    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
226    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
227    ///
228    /// trie.insert(ip20);
229    /// trie.insert(ip24);
230    ///
231    /// assert_eq!( trie.lookup(&ip20), &ip20);
232    /// assert_eq!( trie.lookup(&ip22), &ip20);
233    /// assert_eq!( trie.lookup(&ip24), &ip24);
234    ///
235    /// assert_eq!( trie.lookup(&addr), &ip24);
236    /// ```
237    #[inline]
238    pub fn lookup<Q>(&self, k: &Q) -> &P
239    where
240        Q: IpPrefix<Addr = P::Addr>,
241        P: IpPrefixCovering<Q>,
242    {
243        self.0.lookup(k).0
244    }
245
246    /// Iterates over all the prefixes of this set.
247    #[inline]
248    pub fn iter(&self) -> impl Iterator<Item = &P> + '_ {
249        self.0.leaves.0.iter().map(Leaf::prefix)
250    }
251
252    #[inline]
253    pub fn info(&self) {
254        self.0.info()
255    }
256}
257
258impl<P: IpRootPrefix> Default for RTrieSet<P> {
259    #[inline]
260    fn default() -> Self {
261        Self::new()
262    }
263}
264
265impl<P: IpPrefix> Extend<P> for RTrieSet<P> {
266    fn extend<I: IntoIterator<Item = P>>(&mut self, iter: I) {
267        iter.into_iter().for_each(|item| {
268            self.insert(item);
269        })
270    }
271}
272
273impl<P: IpRootPrefix> FromIterator<P> for RTrieSet<P> {
274    fn from_iter<I: IntoIterator<Item = P>>(iter: I) -> Self {
275        let mut trieset = Self::default();
276        trieset.extend(iter);
277        trieset
278    }
279}
280
281impl<P: IpPrefix> LCTrieSet<P> {
282    /// Returns the size of the set.
283    ///
284    /// Notice that it never equals zero since the top prefix is
285    /// always present in the set.
286    #[inline]
287    #[allow(clippy::len_without_is_empty)]
288    pub fn len(&self) -> NonZeroUsize {
289        self.0.len()
290    }
291
292    #[inline]
293    pub fn info(&self) {
294        self.0.info()
295    }
296
297    /// Checks if an element is present (exact match).
298    ///
299    /// # Example
300    /// ```
301    /// #  use iptrie::*;
302    /// use std::net::Ipv4Addr;
303    ///
304    /// let addr = Ipv4Addr::new(1,1,1,1);
305    ///
306    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
307    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
308    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
309    ///
310    /// let trie = Ipv4RTrieSet::from_iter([ip20,ip24]);
311    /// let lctrie = trie.compress();
312    ///
313    /// assert_eq!( lctrie.contains(&ip20), true);
314    /// assert_eq!( lctrie.contains(&ip22), false);
315    /// assert_eq!( lctrie.contains(&ip24), true);
316    /// ```
317    #[inline]
318    pub fn contains<Q>(&self, k: &Q) -> bool
319    where
320        Q: IpPrefix<Addr = P::Addr>,
321        P: IpPrefixCovering<Q>,
322    {
323        self.0.get(k).is_some()
324    }
325
326    /// Gets the value associated with an exact match of the key.
327    ///
328    /// To access to the longest prefix match, use [`Self::lookup`].
329    ///
330    /// # Example
331    /// ```
332    /// # use iptrie::*;
333    /// # use iptrie::set::*;
334    /// use std::net::Ipv4Addr;
335    /// let mut trie = RTrieSet::new();
336    ///
337    /// let addr = Ipv4Addr::new(1,1,1,1);
338    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
339    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
340    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
341    ///
342    /// trie.insert(ip24);
343    /// trie.insert(ip20);
344    ///
345    /// let lctrie = trie.compress();
346    ///
347    /// assert!( lctrie.get(&ip24).is_some());
348    /// assert!( lctrie.get(&ip22).is_none());
349    /// assert!( lctrie.get(&ip20).is_some());
350    /// ```
351    #[inline]
352    pub fn get<Q>(&self, k: &Q) -> Option<&P>
353    where
354        Q: IpPrefix<Addr = P::Addr>,
355        P: IpPrefixCovering<Q>,
356    {
357        self.0.get(k).map(|(k, _)| k)
358    }
359
360    /// Gets the longest prefix which matches the given key.
361    ///
362    /// As the top prefix always matches, it never fails.
363    ///
364    /// To access to the exact prefix match, use [`Self::get`].
365    ///
366    /// # Example
367    /// ```
368    /// # use iptrie::*;
369    /// use std::net::Ipv4Addr;
370    /// let mut trie = Ipv4RTrieSet::new();
371    ///
372    /// let addr = Ipv4Addr::new(1,1,1,1);
373    /// let ip20 = Ipv4Prefix::new(addr, 20).unwrap();
374    /// let ip22 = Ipv4Prefix::new(addr, 22).unwrap();
375    /// let ip24 = Ipv4Prefix::new(addr, 24).unwrap();
376    ///
377    /// trie.insert(ip20);
378    /// trie.insert(ip24);
379    ///
380    /// let lctrie = trie.compress();
381    ///
382    /// assert_eq!( lctrie.lookup(&ip20), &ip20);
383    /// assert_eq!( lctrie.lookup(&ip22), &ip20);
384    /// assert_eq!( lctrie.lookup(&ip24), &ip24);
385    ///
386    /// assert_eq!( lctrie.lookup(&addr), &ip24);
387    /// ```
388    #[inline]
389    pub fn lookup<Q>(&self, k: &Q) -> &P
390    where
391        Q: IpPrefix<Addr = P::Addr>,
392        P: IpPrefixCovering<Q>,
393    {
394        self.0.lookup(k).0
395    }
396
397    /// Iterates over all the prefixes of this set.
398    #[inline]
399    pub fn iter(&self) -> impl Iterator<Item = &P> + '_ {
400        self.0.leaves.0.iter().map(Leaf::prefix)
401    }
402}
403
404impl<P: IpRootPrefix> FromIterator<P> for LCTrieSet<P> {
405    fn from_iter<I: IntoIterator<Item = P>>(iter: I) -> Self {
406        RTrieSet::from_iter(iter).compress()
407    }
408}
409
410#[cfg(feature = "graphviz")]
411impl<P: IpPrefix + Display> DotWriter for RTrieSet<P> {
412    fn write_dot(&self, dot: &mut dyn std::io::Write) -> std::io::Result<()> {
413        self.0.write_dot(dot)
414    }
415}
416
417#[cfg(feature = "graphviz")]
418impl<P: IpPrefix + Display> DotWriter for LCTrieSet<P> {
419    fn write_dot(&self, dot: &mut dyn std::io::Write) -> std::io::Result<()> {
420        self.0.write_dot(dot)
421    }
422}
iptrie/set.rs

iptrie/
set.rs
