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;