ip/concrete/prefix/set/
mod.rs

1use std::boxed::Box;
2use std::mem;
3
4use super::{Prefix, Range};
5use crate::traits::{self, Afi};
6
7mod iter;
8use self::iter::{Prefixes, Ranges};
9
10mod node;
11use self::node::Node;
12
13mod ops;
14
15/// A collection of IP prefixes, providing fast insertion and iteration,
16/// and set-theorectic arithmetic.
17///
18/// This is a Rust implementation derived in large part from the internal
19/// data-structure used in the widely used [`bgpq3`] tool by Alexandre Snarskii,
20/// packaged as a library, and with set-theoretic operations added.
21///
22/// # Examples
23///
24/// ``` rust
25/// use ip::{traits::PrefixSet as _, Error, Ipv6, Prefix, PrefixLength, PrefixRange, PrefixSet};
26///
27/// fn main() -> Result<(), Error> {
28///     // create a set by parsing a Vec<&str>
29///     let set: PrefixSet<Ipv6> = ["2001:db8::/37", "2001:db8:f00::/37"]
30///         .into_iter()
31///         .map(|s| s.parse::<Prefix<Ipv6>>())
32///         .collect::<Result<_, _>>()?;
33///
34///     // create a range by parsing a &str and providing the lower
35///     // and upper prefix lenth bounds
36///     let length = PrefixLength::<Ipv6>::from_primitive(37)?;
37///     let range = PrefixRange::<Ipv6>::new("2001:db8::/36".parse()?, length..=length)?;
38///
39///     assert_eq!(set.ranges().collect::<Vec<_>>(), vec![range]);
40///     Ok(())
41/// }
42/// ```
43///
44/// Most mutating methods return `&mut Self` for easy chaining, e.g.:
45///
46/// ``` rust
47/// # use ip::{traits::PrefixSet as _, Error, Ipv4, Prefix, PrefixSet};
48/// let set = PrefixSet::<Ipv4>::new()
49///     .insert("192.0.2.0/24".parse::<Prefix<Ipv4>>()?)
50///     .to_owned();
51/// assert_eq!(set.len(), 1);
52/// # Ok::<_, Error>(())
53/// ```
54///
55/// [`bgpq3`]: https://github.com/snar/bgpq3
56#[derive(Clone, Debug)]
57pub struct Set<A: Afi> {
58    root: Option<Box<Node<A>>>,
59}
60
61impl<A: Afi> Set<A> {
62    /// Construct a new, empty [`PrefixSet<A>`][Self].
63    #[must_use]
64    pub const fn new() -> Self {
65        Self::new_with_root(None)
66    }
67
68    const fn new_with_root(root: Option<Box<Node<A>>>) -> Self {
69        Self { root }
70    }
71
72    fn insert_node(&mut self, new: Box<Node<A>>) -> &mut Self {
73        match mem::take(&mut self.root) {
74            Some(root) => {
75                self.root = Some(root.add(new));
76            }
77            None => {
78                self.root = Some(new);
79            }
80        };
81        self
82    }
83
84    pub(crate) fn insert_only<T>(&mut self, item: T) -> &mut Self
85    where
86        T: Into<Node<A>>,
87    {
88        self.insert_node(item.into().boxed())
89    }
90
91    /// Insert a new `item` into `self`.
92    ///
93    /// `T` can be either a [`Prefix<A>`](crate::concrete::Prefix) or a
94    /// [`PrefixRange<A>`](crate::concrete::PrefixRange).
95    ///
96    /// ``` rust
97    /// # use ip::{traits::PrefixSet as _, Error, Ipv6, PrefixRange, PrefixSet};
98    /// let range: PrefixRange<Ipv6> = "2001:db8:f00::/48,64,64".parse()?;
99    /// let set = PrefixSet::<Ipv6>::new().insert(range).to_owned();
100    /// assert_eq!(set.len(), 1 << 16);
101    /// # Ok::<_, Error>(())
102    /// ```
103    pub fn insert<T>(&mut self, item: T) -> &mut Self
104    where
105        T: Into<Node<A>>,
106    {
107        self.insert_only(item).aggregate()
108    }
109
110    /// Insert items into `self` from an iterator yielding either
111    /// [`Prefix<A>`](crate::concrete::Prefix) or
112    /// [`PrefixRange<A>`](crate::concrete::PrefixRange).
113    ///
114    /// Aggregation occurs after all items are inserted, making this far more
115    /// efficient than calling [`PrefixSet::insert()`][Self::insert] repeatedly.
116    ///
117    /// ``` rust
118    /// # use ip::{traits::PrefixSet as _, Error, Ipv4, Prefix, PrefixSet};
119    /// let prefixes: Vec<_> = ["192.0.2.0/26", "192.0.2.64/26"]
120    ///     .into_iter()
121    ///     .map(|s| s.parse::<Prefix<Ipv4>>())
122    ///     .collect::<Result<_, _>>()?;
123    /// let set = PrefixSet::<Ipv4>::new().insert_from(prefixes).to_owned();
124    /// assert_eq!(set.len(), 2);
125    /// # Ok::<_, Error>(())
126    /// ```
127    pub fn insert_from<I, T>(&mut self, iter: I) -> &mut Self
128    where
129        I: IntoIterator<Item = T>,
130        T: Into<Node<A>>,
131    {
132        iter.into_iter()
133            .fold(self, |set, item| set.insert_only(item))
134            .aggregate()
135    }
136
137    fn remove_node(&mut self, mut old: Box<Node<A>>) -> &mut Self {
138        if let Some(root) = mem::take(&mut self.root) {
139            self.root = Some(root.remove(&mut old));
140        };
141        self
142    }
143
144    /// Remove an `item` from `self`.
145    ///
146    /// `T` can be either a [`Prefix<A>`](crate::concrete::Prefix) or a
147    /// [`PrefixRange<A>`](crate::concrete::PrefixRange).
148    ///
149    /// ``` rust
150    /// # use ip::{traits::PrefixSet as _, Error, Ipv6, Prefix, PrefixSet};
151    /// let set = ["2001:db8:f00::/48", "2001:db8:baa::/48"]
152    ///     .into_iter()
153    ///     .map(|s| s.parse::<Prefix<Ipv6>>())
154    ///     .collect::<Result<PrefixSet<Ipv6>, _>>()?
155    ///     .remove("2001:db8:f00::/48".parse::<Prefix<Ipv6>>()?)
156    ///     .to_owned();
157    /// assert_eq!(set.len(), 1);
158    /// # Ok::<_, Error>(())
159    /// ```
160    pub fn remove<T>(&mut self, item: T) -> &mut Self
161    where
162        T: Into<Node<A>>,
163    {
164        self.remove_node(item.into().boxed()).aggregate()
165    }
166
167    /// Remove items from `self` from an iterator yielding either
168    /// [`Prefix<A>`](crate::concrete::Prefix) or
169    /// [`PrefixRange<A>`](crate::concrete::PrefixRange).
170    ///
171    /// Aggregation occurs after all items are removed, making this far more
172    /// efficient than calling [`PrefixSet::remove()`][Self::remove] repeatedly.
173    ///
174    /// ``` rust
175    /// # use ip::{traits::PrefixSet as _, Error, Ipv4, Prefix, PrefixRange, PrefixSet};
176    /// let prefixes: Vec<_> = vec!["192.0.2.0/26", "192.0.2.64/26"]
177    ///     .into_iter()
178    ///     .map(|s| s.parse::<Prefix<Ipv4>>())
179    ///     .collect::<Result<_, _>>()?;
180    /// let mut set = PrefixSet::<Ipv4>::new()
181    ///     .insert("192.0.2.0/24,26,26".parse::<PrefixRange<Ipv4>>()?)
182    ///     .to_owned();
183    /// assert_eq!(set.remove_from(prefixes).len(), 2);
184    /// # Ok::<_, Error>(())
185    /// ```
186    pub fn remove_from<I, T>(&mut self, iter: I) -> &mut Self
187    where
188        I: IntoIterator<Item = T>,
189        T: Into<Node<A>>,
190    {
191        iter.into_iter()
192            .fold(self, |set, item| set.remove_node(item.into().boxed()))
193            .aggregate()
194    }
195
196    pub(crate) fn aggregate(&mut self) -> &mut Self {
197        if let Some(root) = mem::take(&mut self.root) {
198            self.root = root.aggregate(None);
199        }
200        self
201    }
202
203    /// Clear the contents of `self`
204    ///
205    /// ``` rust
206    /// # use ip::{traits::PrefixSet as _, Error, Ipv6, Prefix, PrefixSet};
207    /// let mut set = PrefixSet::<Ipv6>::new()
208    ///     .insert("2001:db8::/32".parse::<Prefix<Ipv6>>()?)
209    ///     .to_owned();
210    /// assert!(!set.is_empty());
211    /// set.clear();
212    /// assert!(set.is_empty());
213    /// # Ok::<_, Error>(())
214    /// ```
215    pub fn clear(&mut self) {
216        self.root = None;
217    }
218}
219
220impl<'a, A: Afi> traits::PrefixSet<'a> for Set<A> {
221    type Prefix = Prefix<A>;
222    type Range = Range<A>;
223    type Prefixes = Prefixes<'a, A>;
224    type Ranges = Ranges<'a, A>;
225
226    fn prefixes(&'a self) -> Self::Prefixes {
227        self.into()
228    }
229
230    fn ranges(&'a self) -> Self::Ranges {
231        self.into()
232    }
233
234    fn contains(&self, prefix: Self::Prefix) -> bool {
235        self.root
236            .as_ref()
237            .map_or(false, |root| root.search(&prefix.into()).is_some())
238    }
239}
240
241impl<A: Afi> Default for Set<A> {
242    fn default() -> Self {
243        Self::new()
244    }
245}
246
247impl<A: Afi, U> Extend<U> for Set<A>
248where
249    U: Into<Node<A>>,
250{
251    #[allow(unused_results)]
252    fn extend<T>(&mut self, iter: T)
253    where
254        T: IntoIterator<Item = U>,
255    {
256        self.insert_from(iter);
257    }
258}
259
260impl<A: Afi, T> FromIterator<T> for Set<A>
261where
262    T: Into<Node<A>>,
263{
264    fn from_iter<I>(iter: I) -> Self
265    where
266        I: IntoIterator<Item = T>,
267    {
268        Self::new().insert_from(iter).clone()
269    }
270}
271
272#[cfg(test)]
273mod tests;