prefix_trie/map/mod.rs
1//! This module contains the implementation for the Dense Prefix Map.
2
3use std::marker::PhantomData;
4
5use crate::{aggregate::Aggregation, allocator::Loc, Prefix};
6
7mod entry;
8mod iter;
9pub use entry::{Entry, OccupiedEntry, VacantEntry};
10pub use iter::*;
11
12use super::table::{Location, Table, K};
13
14/// Prefix map implemented as a TreeBitMap.
15#[derive(Clone)]
16pub struct PrefixMap<P, T> {
17 table: Table<T>,
18 pub(crate) count: usize,
19 marker: PhantomData<P>,
20}
21
22impl<P: Prefix + PartialEq, T: PartialEq> PartialEq for PrefixMap<P, T> {
23 fn eq(&self, other: &Self) -> bool {
24 self.count == other.count && self.iter().eq(other.iter())
25 }
26}
27
28impl<P: Prefix + Eq, T: Eq> Eq for PrefixMap<P, T> {}
29
30impl<P, T> Default for PrefixMap<P, T> {
31 fn default() -> Self {
32 Self {
33 table: Table::default(),
34 count: 0,
35 marker: PhantomData,
36 }
37 }
38}
39
40impl<P, T> PrefixMap<P, T>
41where
42 P: Prefix,
43{
44 /// Create an empty prefix map.
45 pub fn new() -> Self {
46 Self::default()
47 }
48
49 /// Create an empty prefix map.
50 #[cfg(feature = "rkyv")]
51 pub(crate) fn from_table_count(table: Table<T>, count: usize) -> Self {
52 Self {
53 table,
54 count,
55 marker: PhantomData,
56 }
57 }
58
59 /// Returns the number of entries stored in the map.
60 ///
61 /// This is the number of stored prefixes, not the number of addresses they cover (see
62 /// [`address_count`](Self::address_count)).
63 #[inline(always)]
64 pub fn len(&self) -> usize {
65 self.count
66 }
67
68 /// Returns `true` if the map contains no entries.
69 #[inline(always)]
70 pub fn is_empty(&self) -> bool {
71 self.count == 0
72 }
73
74 /// Returns the amount of memory used by this datastructure in bytes.
75 ///
76 /// **Warning**: This number does not include any heap allocations of T!
77 pub fn mem_size(&self) -> usize {
78 self.table.mem_size() + std::mem::size_of::<Self>()
79 }
80
81 /// Count the number of unique addresses covered by all prefixes in the map. If the entire trie
82 /// is covered, the function returns `None` (as it contains `P::R::MAX + 1` addresses).
83 /// Overlapping prefixes are not double-counted.
84 ///
85 /// To avoid double-counting, the function traverses the (partial) tree once, skipping nodes
86 /// that are already covered.
87 ///
88 /// ```
89 /// use prefix_trie::PrefixMap;
90 ///
91 /// # #[cfg(feature = "ipnet")]
92 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
93 /// let mut pm: PrefixMap<ipnet::Ipv4Net, u32> = PrefixMap::new();
94 /// pm.insert("192.0.2.0/24".parse()?, 1);
95 /// pm.insert("192.0.2.128/25".parse()?, 2); // overlaps, counted once
96 /// pm.insert("198.51.100.0/24".parse()?, 3);
97 /// assert_eq!(pm.address_count(), Some(512));
98 ///
99 /// pm.insert("0.0.0.0/0".parse()?, 1);
100 /// assert_eq!(pm.address_count(), None);
101 /// # Ok(())
102 /// # }
103 /// # #[cfg(not(feature = "ipnet"))]
104 /// # fn main() {}
105 /// ```
106 pub fn address_count(&self) -> Option<P::R> {
107 self.table.address_count::<P>()
108 }
109
110 /// Return a reference to the underlying table (crate-internal use only).
111 #[inline(always)]
112 pub(crate) fn table(&self) -> &Table<T> {
113 &self.table
114 }
115
116 /// Return a reference to the underlying table (crate-internal use only).
117 #[inline(always)]
118 pub(crate) fn table_mut(&mut self) -> &mut Table<T> {
119 &mut self.table
120 }
121
122 /// Get the value stored at exactly `prefix`.
123 ///
124 /// ```
125 /// # use prefix_trie::*; use prefix_trie::*;
126 /// # #[cfg(feature = "ipnet")]
127 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
128 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
129 /// pm.insert("192.168.1.0/24".parse()?, 1);
130 /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&1));
131 /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), None);
132 /// assert_eq!(pm.get(&"192.168.0.0/23".parse()?), None);
133 /// assert_eq!(pm.get(&"192.168.1.128/25".parse()?), None);
134 /// # Ok(())
135 /// # }
136 /// # #[cfg(not(feature = "ipnet"))]
137 /// # fn main() {}
138 /// ```
139 pub fn get<'a>(&'a self, prefix: &P) -> Option<&'a T> {
140 let key = prefix.repr();
141 let prefix_len = prefix.prefix_len() as u32;
142 Some(self.table.find(key, prefix_len)?.get())
143 }
144
145 /// Get a mutable reference to the value stored at exactly `prefix`.
146 ///
147 /// ```
148 /// # use prefix_trie::*; use prefix_trie::*;
149 /// # #[cfg(feature = "ipnet")]
150 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
151 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
152 /// let prefix = "192.168.1.0/24".parse()?;
153 /// pm.insert(prefix, 1);
154 /// assert_eq!(pm.get_mut(&prefix), Some(&mut 1));
155 /// *pm.get_mut(&prefix).unwrap() += 1;
156 /// assert_eq!(pm.get_mut(&prefix), Some(&mut 2));
157 /// # Ok(())
158 /// # }
159 /// # #[cfg(not(feature = "ipnet"))]
160 /// # fn main() {}
161 /// ```
162 pub fn get_mut<'a>(&'a mut self, prefix: &P) -> Option<&'a mut T> {
163 let key = prefix.repr();
164 let prefix_len = prefix.prefix_len() as u32;
165 Some(self.table.find_mut(key, prefix_len).present()?.get_mut())
166 }
167
168 /// Get the value stored at exactly `prefix`, together with the canonical matched prefix.
169 ///
170 /// Prefixes are not stored verbatim. They are reconstructed from the trie position, so host
171 /// bits masked out by the prefix length are not preserved.
172 ///
173 /// ```
174 /// # use prefix_trie::*; use prefix_trie::*;
175 /// # #[cfg(feature = "ipnet")]
176 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
177 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
178 /// let prefix = "192.168.1.0/24".parse()?;
179 /// pm.insert(prefix, 1);
180 /// assert_eq!(pm.get_key_value(&prefix), Some((prefix, &1)));
181 /// # Ok(())
182 /// # }
183 /// # #[cfg(not(feature = "ipnet"))]
184 /// # fn main() {}
185 /// ```
186 ///
187 /// **Warning** The table does not store the prefix, but it is reconstructed. This means, that
188 /// any bits in the host part will be truncated:
189 ///
190 /// ```
191 /// # use prefix_trie::*; use prefix_trie::*;
192 /// # #[cfg(feature = "ipnet")]
193 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
194 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
195 /// let prefix = "192.168.1.0/24".parse()?;
196 /// pm.insert(prefix, 1);
197 /// assert_eq!(pm.get_key_value(&prefix), Some((prefix.trunc(), &1)));
198 /// # Ok(())
199 /// # }
200 /// # #[cfg(not(feature = "ipnet"))]
201 /// # fn main() {}
202 /// ```
203 pub fn get_key_value<'a>(&'a self, prefix: &P) -> Option<(P, &'a T)> {
204 let key = prefix.repr();
205 let prefix_len = prefix.prefix_len() as u32;
206 let r = self.table.find(key, prefix_len)?;
207 let p = r.prefix(key);
208 Some((p, r.get()))
209 }
210
211 /// Get the longest prefix in the map that contains `prefix`, together with its value.
212 ///
213 /// ```
214 /// # use prefix_trie::*; use prefix_trie::*;
215 /// # #[cfg(feature = "ipnet")]
216 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
217 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
218 /// pm.insert("192.168.1.0/24".parse()?, 1);
219 /// pm.insert("192.168.0.0/23".parse()?, 2);
220 /// assert_eq!(pm.get_lpm(&"192.168.1.1/32".parse()?), Some(("192.168.1.0/24".parse()?, &1)));
221 /// assert_eq!(pm.get_lpm(&"192.168.1.0/24".parse()?), Some(("192.168.1.0/24".parse()?, &1)));
222 /// assert_eq!(pm.get_lpm(&"192.168.0.0/24".parse()?), Some(("192.168.0.0/23".parse()?, &2)));
223 /// assert_eq!(pm.get_lpm(&"192.168.2.0/24".parse()?), None);
224 /// # Ok(())
225 /// # }
226 /// # #[cfg(not(feature = "ipnet"))]
227 /// # fn main() {}
228 /// ```
229 ///
230 /// **Warning** The table does not store the prefix, but it is reconstructed. This means, that
231 /// any bits in the host part will be truncated:
232 ///
233 /// ```
234 /// # use prefix_trie::*; use prefix_trie::*;
235 /// # #[cfg(feature = "ipnet")]
236 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
237 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
238 /// pm.insert("192.168.1.1/24".parse()?, 1);
239 /// assert_eq!(pm.get_lpm(&"192.168.1.1/32".parse()?), Some(("192.168.1.0/24".parse()?, &1)));
240 /// # Ok(())
241 /// # }
242 /// # #[cfg(not(feature = "ipnet"))]
243 /// # fn main() {}
244 /// ```
245 pub fn get_lpm<'a>(&'a self, prefix: &P) -> Option<(P, &'a T)> {
246 let key = prefix.repr();
247 let prefix_len = prefix.prefix_len() as u32;
248 let r = self.table.find_lpm(key, prefix_len)?;
249 let p = r.prefix(key);
250 Some((p, r.get()))
251 }
252
253 /// Get a mutable reference to the value of the longest prefix in the map that contains `prefix`.
254 ///
255 /// ```
256 /// # use prefix_trie::*; use prefix_trie::*;
257 /// # #[cfg(feature = "ipnet")]
258 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
259 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
260 /// pm.insert("192.168.1.0/24".parse()?, 1);
261 /// pm.insert("192.168.0.0/23".parse()?, 2);
262 /// assert_eq!(pm.get_lpm_mut(&"192.168.1.1/32".parse()?), Some(("192.168.1.0/24".parse()?, &mut 1)));
263 /// *pm.get_lpm_mut(&"192.168.1.64/26".parse()?).unwrap().1 += 1;
264 /// assert_eq!(pm.get_lpm_mut(&"192.168.1.1/32".parse()?), Some(("192.168.1.0/24".parse()?, &mut 2)));
265 /// # Ok(())
266 /// # }
267 /// # #[cfg(not(feature = "ipnet"))]
268 /// # fn main() {}
269 /// ```
270 ///
271 /// **Warning** The table does not store the prefix, but it is reconstructed. This means, that
272 /// any bits in the host part will be truncated.
273 pub fn get_lpm_mut<'a>(&'a mut self, prefix: &P) -> Option<(P, &'a mut T)> {
274 let key = prefix.repr();
275 let prefix_len = prefix.prefix_len() as u32;
276 let r = self.table.find_lpm_mut(key, prefix_len)?;
277 let p = r.prefix::<P>(key);
278 Some((p, r.get_mut()))
279 }
280
281 /// Get the longest prefix in the map that contains `prefix`.
282 ///
283 /// ```
284 /// # use prefix_trie::*; use prefix_trie::*;
285 /// # #[cfg(feature = "ipnet")]
286 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
287 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
288 /// pm.insert("192.168.1.0/24".parse()?, 1);
289 /// pm.insert("192.168.0.0/23".parse()?, 2);
290 /// assert_eq!(pm.get_lpm_prefix(&"192.168.1.1/32".parse()?), Some("192.168.1.0/24".parse()?));
291 /// assert_eq!(pm.get_lpm_prefix(&"192.168.1.0/24".parse()?), Some("192.168.1.0/24".parse()?));
292 /// assert_eq!(pm.get_lpm_prefix(&"192.168.0.0/24".parse()?), Some("192.168.0.0/23".parse()?));
293 /// assert_eq!(pm.get_lpm_prefix(&"192.168.2.0/24".parse()?), None);
294 /// # Ok(())
295 /// # }
296 /// # #[cfg(not(feature = "ipnet"))]
297 /// # fn main() {}
298 /// ```
299 ///
300 /// **Warning** The table does not store the prefix, but it is reconstructed. This means, that
301 /// any bits in the host part will be truncated:
302 ///
303 /// ```
304 /// # use prefix_trie::*; use prefix_trie::*;
305 /// # #[cfg(feature = "ipnet")]
306 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
307 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
308 /// pm.insert("192.168.1.1/24".parse()?, 1);
309 /// assert_eq!(pm.get_lpm_prefix(&"192.168.1.1/32".parse()?), Some("192.168.1.0/24".parse()?));
310 /// # Ok(())
311 /// # }
312 /// # #[cfg(not(feature = "ipnet"))]
313 /// # fn main() {}
314 /// ```
315 pub fn get_lpm_prefix(&self, prefix: &P) -> Option<P> {
316 self.get_lpm(prefix).map(|(p, _)| p)
317 }
318
319 /// Check whether `prefix` is present in the map.
320 ///
321 /// ```
322 /// # use prefix_trie::*; use prefix_trie::*;
323 /// # #[cfg(feature = "ipnet")]
324 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
325 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
326 /// pm.insert("192.168.1.0/24".parse()?, 1);
327 /// assert!(pm.contains_key(&"192.168.1.0/24".parse()?));
328 /// assert!(!pm.contains_key(&"192.168.2.0/24".parse()?));
329 /// assert!(!pm.contains_key(&"192.168.0.0/23".parse()?));
330 /// assert!(!pm.contains_key(&"192.168.1.128/25".parse()?));
331 /// # Ok(())
332 /// # }
333 /// # #[cfg(not(feature = "ipnet"))]
334 /// # fn main() {}
335 /// ```
336 pub fn contains_key(&self, prefix: &P) -> bool {
337 let key = prefix.repr();
338 let prefix_len = prefix.prefix_len() as u32;
339 self.table.find(key, prefix_len).is_some()
340 }
341
342 /// Get the shortest prefix in the map that contains `prefix`, together with its value.
343 ///
344 /// ```
345 /// # use prefix_trie::*; use prefix_trie::*;
346 /// # #[cfg(feature = "ipnet")]
347 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
348 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
349 /// pm.insert("192.168.1.0/24".parse()?, 1);
350 /// pm.insert("192.168.0.0/23".parse()?, 2);
351 /// assert_eq!(pm.get_spm(&"192.168.1.1/32".parse()?), Some(("192.168.0.0/23".parse()?, &2)));
352 /// assert_eq!(pm.get_spm(&"192.168.1.0/24".parse()?), Some(("192.168.0.0/23".parse()?, &2)));
353 /// assert_eq!(pm.get_spm(&"192.168.0.0/23".parse()?), Some(("192.168.0.0/23".parse()?, &2)));
354 /// assert_eq!(pm.get_spm(&"192.168.2.0/24".parse()?), None);
355 /// # Ok(())
356 /// # }
357 /// # #[cfg(not(feature = "ipnet"))]
358 /// # fn main() {}
359 /// ```
360 ///
361 /// **Warning** The table does not store the prefix, but it is reconstructed. This means, that
362 /// any bits in the host part will be truncated.
363 pub fn get_spm<'a>(&'a self, prefix: &P) -> Option<(P, &'a T)> {
364 let key = prefix.repr();
365 let prefix_len = prefix.prefix_len() as u32;
366 let r = self.table.find_spm(key, prefix_len)?;
367 let p = r.prefix(key);
368 Some((p, r.get()))
369 }
370
371 /// Get the shortest prefix in the map that contains `prefix`.
372 ///
373 /// ```
374 /// # use prefix_trie::*; use prefix_trie::*;
375 /// # #[cfg(feature = "ipnet")]
376 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
377 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
378 /// pm.insert("192.168.1.1/24".parse()?, 1);
379 /// pm.insert("192.168.0.0/23".parse()?, 2);
380 /// assert_eq!(pm.get_spm_prefix(&"192.168.1.1/32".parse()?), Some("192.168.0.0/23".parse()?));
381 /// assert_eq!(pm.get_spm_prefix(&"192.168.1.0/24".parse()?), Some("192.168.0.0/23".parse()?));
382 /// assert_eq!(pm.get_spm_prefix(&"192.168.0.0/23".parse()?), Some("192.168.0.0/23".parse()?));
383 /// assert_eq!(pm.get_spm_prefix(&"192.168.2.0/24".parse()?), None);
384 /// # Ok(())
385 /// # }
386 /// # #[cfg(not(feature = "ipnet"))]
387 /// # fn main() {}
388 /// ```
389 ///
390 /// **Warning** The table does not store the prefix, but it is reconstructed. This means, that
391 /// any bits in the host part will be truncated.
392 pub fn get_spm_prefix(&self, prefix: &P) -> Option<P> {
393 self.get_spm(prefix).map(|(p, _)| p)
394 }
395
396 /// Check whether `prefix` is covered by the map, i.e., whether the map contains an entry at
397 /// `prefix` itself or any less-specific prefix that contains it.
398 ///
399 /// This is equivalent to `self.cover(prefix).next().is_some()`, but stops at the first
400 /// (shortest) covering prefix. See [`cover`](Self::cover) to iterate over the covering
401 /// entries themselves.
402 ///
403 /// This function does not perform aggregation. That means that, even if both the left and
404 /// right children of `p` are present in the map, `is_covered(p)` may still return `false`. See
405 /// [`is_covered_in_aggregate`](Self::is_covered_in_aggregate) for that case.
406 ///
407 /// ```
408 /// # use prefix_trie::*; use prefix_trie::*;
409 /// # #[cfg(feature = "ipnet")]
410 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
411 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
412 /// pm.insert("10.0.0.0/8".parse()?, 1);
413 /// assert!(pm.is_covered(&"10.0.0.0/8".parse()?)); // exact member
414 /// assert!(pm.is_covered(&"10.1.2.0/24".parse()?)); // covered by 10.0.0.0/8
415 /// assert!(!pm.is_covered(&"11.0.0.0/8".parse()?)); // not covered
416 /// # Ok(())
417 /// # }
418 /// # #[cfg(not(feature = "ipnet"))]
419 /// # fn main() {}
420 /// ```
421 #[inline(always)]
422 pub fn is_covered(&self, prefix: &P) -> bool {
423 self.get_spm_prefix(prefix).is_some()
424 }
425
426 /// Check whether every address in `prefix` is covered by the map, i.e., whether `prefix`'s
427 /// entire range is tiled by entries in the map, even if no single entry covers `prefix` on its
428 /// own.
429 ///
430 /// This is equivalent to `{ let mut m = self.clone(); m.aggregate(); m.is_covered(prefix) }`,
431 /// but read-only and without cloning. See [`is_covered`](Self::is_covered) for the (cheaper,
432 /// stricter) single-entry check.
433 ///
434 /// ```
435 /// # use prefix_trie::*; use prefix_trie::*;
436 /// # #[cfg(feature = "ipnet")]
437 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
438 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
439 /// pm.insert("10.0.0.0/9".parse()?, 1);
440 /// pm.insert("10.128.0.0/9".parse()?, 2);
441 /// assert!(!pm.is_covered(&"10.0.0.0/8".parse()?)); // no single covering entry
442 /// assert!(pm.is_covered_in_aggregate(&"10.0.0.0/8".parse()?)); // the two /9s tile the /8
443 /// # Ok(())
444 /// # }
445 /// # #[cfg(not(feature = "ipnet"))]
446 /// # fn main() {}
447 /// ```
448 pub fn is_covered_in_aggregate(&self, prefix: &P) -> bool {
449 let key = prefix.repr();
450 let prefix_len = prefix.prefix_len() as u32;
451 self.table.covers_in_aggregate(key, prefix_len)
452 }
453
454 /// Insert a new item into the prefix-map. This function may return any value that existed
455 /// before.
456 ///
457 /// ```
458 /// # use prefix_trie::*; use prefix_trie::*;
459 /// # #[cfg(feature = "ipnet")]
460 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
461 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
462 /// assert_eq!(pm.insert("192.168.0.0/23".parse()?, 1), None);
463 /// assert_eq!(pm.insert("192.168.1.0/24".parse()?, 2), None);
464 /// assert_eq!(pm.insert("192.168.1.0/24".parse()?, 3), Some(2));
465 /// # Ok(())
466 /// # }
467 /// # #[cfg(not(feature = "ipnet"))]
468 /// # fn main() {}
469 /// ```
470 ///
471 /// **Warning**: You *cannot* store additional information in the host-part of the prefix.
472 /// Prefixes are reconstructed from the trie position, so host bits are not preserved.
473 ///
474 /// ```
475 /// # use prefix_trie::*; use prefix_trie::*;
476 /// # #[cfg(feature = "ipnet")]
477 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
478 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
479 ///
480 /// pm.insert("192.168.0.1/24".parse()?, 1);
481 /// assert_eq!(
482 /// pm.get_key_value(&"192.168.0.0/24".parse()?),
483 /// Some(("192.168.0.0/24".parse()?, &1)) // notice that the host part is zero.
484 /// );
485 /// # Ok(())
486 /// # }
487 /// # #[cfg(not(feature = "ipnet"))]
488 /// # fn main() {}
489 /// ```
490 pub fn insert(&mut self, prefix: P, value: T) -> Option<T> {
491 let key = prefix.repr();
492 let prefix_len = prefix.prefix_len() as u32;
493 match self.table.find_or_insert_mut(key, prefix_len) {
494 Ok(present) => Some(present.replace(value)),
495 Err(empty) => {
496 empty.insert(value);
497 self.count += 1;
498 None
499 }
500 }
501 }
502
503 /// Gets the given key's corresponding entry in the map for in-place manipulation.
504 ///
505 /// Prefixes are not stored verbatim. They are reconstructed from the trie position, so host
506 /// bits masked out by the prefix length are not preserved. See the documentation of
507 /// [`Entry`], [`OccupiedEntry`], and [`VacantEntry`].
508 ///
509 /// ```
510 /// # use prefix_trie::*; use prefix_trie::*;
511 /// # #[cfg(feature = "ipnet")]
512 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
513 /// let mut pm: PrefixMap<ipnet::Ipv4Net, Vec<i32>> = PrefixMap::new();
514 /// pm.insert("192.168.0.0/23".parse()?, vec![1]);
515 /// pm.entry("192.168.0.1/23".parse()?).or_default().push(2);
516 /// pm.entry("192.168.0.0/24".parse()?).or_default().push(3);
517 /// assert_eq!(pm.get(&"192.168.0.0/23".parse()?), Some(&vec![1, 2]));
518 /// assert_eq!(pm.get(&"192.168.0.0/24".parse()?), Some(&vec![3]));
519 /// # Ok(())
520 /// # }
521 /// # #[cfg(not(feature = "ipnet"))]
522 /// # fn main() {}
523 /// ```
524 pub fn entry(&mut self, prefix: P) -> Entry<'_, P, T> {
525 let key = prefix.repr();
526 let prefix_len = prefix.prefix_len() as u32;
527 // Split borrows so that `loc` (borrowing `table`) and `count` (borrowing `count`)
528 // can coexist inside the returned Entry without a full `&mut PrefixMap` borrow.
529 let table = &mut self.table;
530 let count = &mut self.count;
531 match table.find_mut(key, prefix_len) {
532 Location::Present(r) => Entry::Occupied(OccupiedEntry::new(r, count, prefix)),
533 Location::Empty(e) => Entry::Vacant(VacantEntry::empty(e, count, prefix)),
534 Location::NoNode(n) => Entry::Vacant(VacantEntry::no_node(n, count, prefix)),
535 }
536 }
537
538 /// Removes a key from the map, returning the value at the key if the key was previously in the
539 /// map. In contrast to [`Self::remove_keep_tree`], this operation may prune empty trie nodes,
540 /// reducing the memory footprint.
541 ///
542 /// ```
543 /// # use prefix_trie::*; use prefix_trie::*;
544 /// # #[cfg(feature = "ipnet")]
545 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
546 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
547 /// let prefix = "192.168.1.0/24".parse()?;
548 /// pm.insert(prefix, 1);
549 /// assert_eq!(pm.get(&prefix), Some(&1));
550 /// assert_eq!(pm.remove(&prefix), Some(1));
551 /// assert_eq!(pm.get(&prefix), None);
552 /// # Ok(())
553 /// # }
554 /// # #[cfg(not(feature = "ipnet"))]
555 /// # fn main() {}
556 /// ```
557 pub fn remove(&mut self, prefix: &P) -> Option<T> {
558 let key = prefix.repr();
559 let prefix_len = prefix.prefix_len() as u32;
560 let (loc_mut, mut path) = self.table.find_mut_with_path(key, prefix_len)?;
561
562 let node_loc = loc_mut.node_loc();
563 let old_value = if let Some(present) = loc_mut.present() {
564 let val = present.take();
565 self.count -= 1;
566 Some(val)
567 } else {
568 None
569 };
570
571 // cleanup_tree handles root internally (noop); call unconditionally.
572 // SAFETY: `node_loc` came from `find_mut_with_path`; `present.take()` only removes
573 // a data cell and does not alter node structure, so `node_loc` and `path` remain valid.
574 unsafe { self.table.cleanup_tree(node_loc, &mut path) };
575
576 old_value
577 }
578
579 /// Removes a key from the map, returning the value at the key if the key was previously in the
580 /// map. In contrast to [`Self::remove`], this operation only removes the stored value and may
581 /// leave empty trie nodes in place.
582 ///
583 /// ```
584 /// # use prefix_trie::*; use prefix_trie::*;
585 /// # #[cfg(feature = "ipnet")]
586 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
587 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
588 /// let prefix = "192.168.1.0/24".parse()?;
589 /// pm.insert(prefix, 1);
590 /// assert_eq!(pm.get(&prefix), Some(&1));
591 /// assert_eq!(pm.remove_keep_tree(&prefix), Some(1));
592 /// assert_eq!(pm.get(&prefix), None);
593 /// # Ok(())
594 /// # }
595 /// # #[cfg(not(feature = "ipnet"))]
596 /// # fn main() {}
597 /// ```
598 pub fn remove_keep_tree(&mut self, prefix: &P) -> Option<T> {
599 let key = prefix.repr();
600 let prefix_len = prefix.prefix_len() as u32;
601 let present = self.table.find_mut(key, prefix_len).present()?;
602 self.count -= 1;
603 Some(present.take())
604 }
605
606 /// Remove all entries that are contained within `prefix`. This will change the tree
607 /// structure. This operation is `O(n)`, as the entries must be freed up one-by-one. Like
608 /// [`Self::remove`], this prunes trie nodes that become empty.
609 ///
610 /// ```
611 /// # use prefix_trie::*; use prefix_trie::*;
612 /// # #[cfg(feature = "ipnet")]
613 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
614 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
615 /// pm.insert("192.168.0.0/21".parse()?, 1);
616 /// pm.insert("192.168.0.0/22".parse()?, 2);
617 /// pm.insert("192.168.0.0/23".parse()?, 3);
618 /// pm.insert("192.168.0.0/24".parse()?, 4);
619 /// pm.insert("192.168.4.0/22".parse()?, 5);
620 /// pm.insert("192.168.4.0/23".parse()?, 6);
621 ///
622 /// assert_eq!(pm.len(), 6);
623 /// pm.remove_children(&"192.168.0.0/22".parse()?);
624 /// assert_eq!(pm.len(), 3);
625 ///
626 /// assert_eq!(pm.get(&"192.168.0.0/22".parse()?), None);
627 /// assert_eq!(pm.get(&"192.168.0.0/23".parse()?), None);
628 /// assert_eq!(pm.get(&"192.168.0.0/24".parse()?), None);
629 /// assert_eq!(pm.get(&"192.168.4.0/22".parse()?), Some(&5));
630 /// assert_eq!(pm.get(&"192.168.4.0/23".parse()?), Some(&6));
631 /// # Ok(())
632 /// # }
633 /// # #[cfg(not(feature = "ipnet"))]
634 /// # fn main() {}
635 /// ```
636 pub fn remove_children(&mut self, prefix: &P) {
637 let key = prefix.repr();
638 let prefix_len = prefix.prefix_len() as u32;
639
640 if prefix_len == 0 {
641 return self.clear();
642 }
643
644 let Some((loc_mut, mut path)) = self.table.find_mut_with_path(key, prefix_len) else {
645 return;
646 };
647 let node = loc_mut.node_loc();
648 let depth = loc_mut.depth();
649
650 // fast-track delete this index if it covers the entire node
651 if prefix_len % K == 0 {
652 // SAFETY: `node` came from `find_mut_with_path` with no subsequent structural
653 // mutations.
654 self.count -= unsafe { self.table.clear_node_and_children(node) };
655 } else {
656 // Collect bitmap bits of covered data elements (from current node state).
657 // SAFETY: `node` came from `find_mut_with_path`; no structural mutations have
658 // occurred yet.
659 let covered_bits: Vec<u32> =
660 unsafe { self.table.data_descendants(node, depth, key, prefix_len) }
661 .map(|mp| mp.bit)
662 .collect();
663 for bit in covered_bits {
664 let idx = super::table::DataIdx { node, bit, depth };
665 // SAFETY: We only remove data cells in this loop; the node allocator structure
666 // (MultiBitNode slots, child pointers) is not modified, so `node` remains valid.
667 // resolve_mut re-reads the current AllocIdx + bitmap bit on each call, so it
668 // correctly handles any tier downgrades that occurred on prior iterations.
669 if let Some(r) = unsafe { idx.resolve_mut(&mut self.table) } {
670 r.take();
671 self.count -= 1;
672 }
673 }
674
675 // Collect bitmap bits of covered children (from original bitmap).
676 let covered_child_bits: Vec<u32> = self
677 .table
678 .node(node)
679 .child_cover_locs(depth, key, prefix_len)
680 .map(|loc| loc.bit)
681 .collect();
682
683 // First: clear each covered child's subtree using the original Loc (parent bitmap
684 // unchanged).
685 for &child_bit in &covered_child_bits {
686 // SAFETY: `node` is still valid (data-only removals above did not affect node
687 // structure). `child_bit` is set in the child_bitmap (from `child_cover_locs`).
688 let child_loc = unsafe { self.table.child(node, child_bit) }
689 .expect("child_bit should exist in bitmap");
690 // SAFETY: `child_loc` was just obtained from a valid `node` via `child()`.
691 self.count -= unsafe { self.table.clear_node_and_children(child_loc) };
692 }
693
694 // Then: remove covered children from parent. `remove_child_at` re-reads the current
695 // bitmap each time, so order does not matter.
696 for &child_bit in &covered_child_bits {
697 // SAFETY: `node` is still valid; each `clear_node_and_children` above only freed
698 // the *child's* allocation, not the parent's. The child_bitmap bit is still set.
699 unsafe { self.table.remove_child_at(node, child_bit) };
700 }
701 }
702
703 // Detach `node` (and any emptied ancestors) if the removal left it empty.
704 // SAFETY: everything above only touches `node`'s data and children allocations. `node`
705 // itself and every Loc in `path` live in their parents' children blocks, which are
706 // unaffected, so all locations are still valid.
707 unsafe { self.table.cleanup_tree(node, &mut path) };
708 }
709
710 /// Clear the map but keep the allocated memory.
711 ///
712 /// ```
713 /// # use prefix_trie::*; use prefix_trie::*;
714 /// # #[cfg(feature = "ipnet")]
715 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
716 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
717 /// pm.insert("192.168.0.0/24".parse()?, 1);
718 /// pm.insert("192.168.1.0/24".parse()?, 2);
719 /// pm.clear();
720 /// assert_eq!(pm.get(&"192.168.0.0/24".parse()?), None);
721 /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), None);
722 /// # Ok(())
723 /// # }
724 /// # #[cfg(not(feature = "ipnet"))]
725 /// # fn main() {}
726 /// ```
727 pub fn clear(&mut self) {
728 // SAFETY: `Loc::root()` is always a valid, live node location.
729 let deleted = unsafe { self.table.clear_node_and_children(Loc::root()) };
730 debug_assert_eq!(deleted, self.count);
731 self.count = 0;
732 }
733
734 /// Keep only the elements in the map that satisfy the given condition `f`.
735 ///
736 /// ```
737 /// # use prefix_trie::*; use prefix_trie::*;
738 /// # #[cfg(feature = "ipnet")]
739 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
740 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
741 /// pm.insert("192.168.0.0/24".parse()?, 1);
742 /// pm.insert("192.168.1.0/24".parse()?, 2);
743 /// pm.insert("192.168.2.0/24".parse()?, 3);
744 /// pm.insert("192.168.2.0/25".parse()?, 4);
745 /// pm.retain(|_, t| *t % 2 == 0);
746 /// assert_eq!(pm.get(&"192.168.0.0/24".parse()?), None);
747 /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), Some(&2));
748 /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), None);
749 /// assert_eq!(pm.get(&"192.168.2.0/25".parse()?), Some(&4));
750 /// # Ok(())
751 /// # }
752 /// # #[cfg(not(feature = "ipnet"))]
753 /// # fn main() {}
754 /// ```
755 ///
756 /// You can also use the prefix for filtering
757 ///
758 /// ```
759 /// # use prefix_trie::*; use prefix_trie::*;
760 /// # #[cfg(feature = "ipnet")]
761 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
762 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
763 /// pm.insert("192.168.0.0/24".parse()?, 1);
764 /// pm.insert("192.168.1.0/24".parse()?, 2);
765 /// pm.insert("192.168.2.0/24".parse()?, 3);
766 /// pm.insert("192.168.2.0/25".parse()?, 4);
767 /// pm.retain(|p, _| p.prefix_len() > 24);
768 /// assert_eq!(pm.get(&"192.168.0.0/24".parse()?), None);
769 /// assert_eq!(pm.get(&"192.168.1.0/24".parse()?), None);
770 /// assert_eq!(pm.get(&"192.168.2.0/24".parse()?), None);
771 /// assert_eq!(pm.get(&"192.168.2.0/25".parse()?), Some(&4));
772 /// # Ok(())
773 /// # }
774 /// # #[cfg(not(feature = "ipnet"))]
775 /// # fn main() {}
776 /// ```
777 pub fn retain<F>(&mut self, mut f: F)
778 where
779 F: FnMut(&P, &T) -> bool,
780 {
781 let removed = self.table.retain_all::<P, _>(&mut f);
782 self.count -= removed;
783 }
784
785 /// Removes every entry whose nearest covering ancestor (a less specific prefix) maps to the
786 /// **same value**, without merging adjacent prefixes.
787 ///
788 /// **Invariant**: for *any* prefix `p`, `before.get_lpm(p)` and `after.get_lpm(p)` return the
789 /// same value (the matched prefix may become less specific).
790 ///
791 /// ```
792 /// use prefix_trie::PrefixMap;
793 ///
794 /// # #[cfg(feature = "ipnet")]
795 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
796 /// let mut pm: PrefixMap<ipnet::Ipv4Net, u32> = PrefixMap::new();
797 /// pm.insert("10.0.0.0/16".parse()?, 1);
798 /// pm.insert("10.0.0.0/24".parse()?, 2); // exception under 10.0.0.0/16
799 /// pm.insert("10.0.1.0/24".parse()?, 2); // sibling of the above, same value
800 /// pm.insert("10.0.2.0/24".parse()?, 1); // same value as 10.0.0.0/16 -> redundant
801 /// pm.insert("192.168.0.0/16".parse()?, 1); // a separate branch, same value
802 /// pm.aggregate_consistent();
803 /// // Only the redundant 10.0.2.0/24 is dropped; nothing is merged.
804 /// assert_eq!(pm.iter().collect::<Vec<_>>(), vec![
805 /// ("10.0.0.0/16".parse()?, &1),
806 /// ("10.0.0.0/24".parse()?, &2),
807 /// ("10.0.1.0/24".parse()?, &2),
808 /// ("192.168.0.0/16".parse()?, &1),
809 /// ]);
810 /// # Ok(())
811 /// # }
812 /// # #[cfg(not(feature = "ipnet"))]
813 /// # fn main() {}
814 /// ```
815 pub fn aggregate_consistent(&mut self)
816 where
817 T: Clone + Eq,
818 {
819 // SAFETY: `Loc::root()` is always a valid, live node location.
820 let (_, count_delta) = unsafe { self.table.aggregate_consistent_map(Loc::root(), 0, None) };
821 self.count = (self.count as i64 + count_delta) as usize;
822 }
823
824 /// Reduce the map to the fewest entries that keep every lookup unchanged.
825 ///
826 /// For any address `a` (a host prefix, i.e. one of maximum length), `self.get_lpm(&a)` resolves
827 /// to the same value as before (only the matched prefix may differ). Holes remain uncovered.
828 /// Among all maps with that property this keeps the fewest entries. For prefixes, this may not
829 /// be the case; When two siblings with the same value get merged, the the parent prefix holds a
830 /// value after aggregation.
831 ///
832 /// The guarantee is per address, not per prefix: entries may be merged or moved. Use
833 /// [`aggregate_consistent`](Self::aggregate_consistent) instead to keep every prefix matching the
834 /// same entry.
835 ///
836 /// ```
837 /// use prefix_trie::PrefixMap;
838 ///
839 /// # #[cfg(feature = "ipnet")]
840 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
841 /// let mut pm: PrefixMap<ipnet::Ipv4Net, u32> = PrefixMap::new();
842 /// pm.insert("10.0.0.0/16".parse()?, 1);
843 /// pm.insert("10.0.0.0/24".parse()?, 2);
844 /// pm.insert("10.0.1.0/24".parse()?, 2);
845 /// pm.insert("10.0.2.0/24".parse()?, 1);
846 /// pm.insert("192.168.0.0/16".parse()?, 1);
847 /// pm.aggregate();
848 /// // The siblings merge into a /23 and the redundant /24 is dropped, but the two /16 branches
849 /// // cannot merge across the uncovered space between them.
850 /// assert_eq!(pm.iter().collect::<Vec<_>>(), vec![
851 /// ("10.0.0.0/16".parse()?, &1),
852 /// ("10.0.0.0/23".parse()?, &2),
853 /// ("192.168.0.0/16".parse()?, &1),
854 /// ]);
855 /// # Ok(())
856 /// # }
857 /// # #[cfg(not(feature = "ipnet"))]
858 /// # fn main() {}
859 /// ```
860 pub fn aggregate(&mut self)
861 where
862 T: Clone + Ord,
863 {
864 let delta = self
865 .table
866 .aggregate_map::<P::R, fn() -> T>(Aggregation::Drop);
867 self.count = (self.count as i64 + delta) as usize;
868 }
869
870 /// Reduce the map to the fewest entries, inserting otherwise-uncovered addresses to `default`.
871 ///
872 /// For any address `a` (a host prefix, i.e. one of maximum length), `self.get_lpm(&a)` under
873 /// `.unwrap_or_else(default)` remains unchanged: covered addresses keep their value, and
874 /// uncovered addresses now resolve to `default()`. Among all maps with that property this
875 /// keeps the fewest entries.
876 ///
877 /// Because no address is left uncovered, the result is always **total**: the root of the tree
878 /// (e.g., 0.0.0.0/0) will contain a value, so [`get_lpm`](Self::get_lpm) always returns `Some`.
879 ///
880 /// The guarantee is per address, not per prefix: entries may be merged or moved. Use
881 /// [`aggregate_consistent`](Self::aggregate_consistent) instead to keep every prefix matching the
882 /// same entry.
883 ///
884 /// ```
885 /// use prefix_trie::PrefixMap;
886 ///
887 /// # #[cfg(feature = "ipnet")]
888 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
889 /// let mut pm: PrefixMap<ipnet::Ipv4Net, u32> = PrefixMap::new();
890 /// pm.insert("10.0.0.0/16".parse()?, 1);
891 /// pm.insert("10.0.0.0/24".parse()?, 2);
892 /// pm.insert("10.0.1.0/24".parse()?, 2);
893 /// pm.insert("10.0.2.0/24".parse()?, 1);
894 /// pm.insert("192.168.0.0/16".parse()?, 1);
895 /// pm.aggregate_fill(|| 1);
896 /// // Filling the gaps with 1 lets both /16 branches and all uncovered space collapse into one
897 /// // default route; only the 10.0.0.0/23 = 2 exception survives.
898 /// assert_eq!(pm.iter().collect::<Vec<_>>(), vec![
899 /// ("0.0.0.0/0".parse()?, &1),
900 /// ("10.0.0.0/23".parse()?, &2),
901 /// ]);
902 /// # Ok(())
903 /// # }
904 /// # #[cfg(not(feature = "ipnet"))]
905 /// # fn main() {}
906 /// ```
907 pub fn aggregate_fill<F>(&mut self, default: F)
908 where
909 T: Clone + Ord,
910 F: Fn() -> T + Copy,
911 {
912 let delta = self
913 .table
914 .aggregate_map::<P::R, F>(Aggregation::Fill(default));
915 self.count = (self.count as i64 + delta) as usize;
916 }
917
918 /// [`aggregate_fill`](Self::aggregate_fill) with `T::default` as the fill value.
919 pub fn aggregate_fill_default(&mut self)
920 where
921 T: Clone + Ord + Default,
922 {
923 self.aggregate_fill(T::default)
924 }
925
926 /// Iterate over all entries in the map that cover `prefix`, including `prefix` itself if it is
927 /// present. The returned iterator yields `(P, &'a T)`, with reconstructed prefixes `P`.
928 ///
929 /// The iterator will always yield elements ordered by their prefix length, i.e., their depth in
930 /// the tree.
931 ///
932 /// ```
933 /// # use prefix_trie::*; use prefix_trie::*;
934 /// # #[cfg(feature = "ipnet")]
935 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
936 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
937 /// let p0 = "10.0.0.0/8".parse()?;
938 /// let p1 = "10.1.0.0/16".parse()?;
939 /// let p2 = "10.1.1.0/24".parse()?;
940 /// pm.insert(p0, 0);
941 /// pm.insert(p1, 1);
942 /// pm.insert(p2, 2);
943 /// pm.insert("10.1.2.0/24".parse()?, 3); // disjoint prefixes are not covered
944 /// pm.insert("10.1.1.0/25".parse()?, 4); // more specific prefixes are not covered
945 /// pm.insert("11.0.0.0/8".parse()?, 5); // Branch points that don't contain values are skipped
946 /// assert_eq!(
947 /// pm.cover(&p2).collect::<Vec<_>>(),
948 /// vec![(p0, &0), (p1, &1), (p2, &2)]
949 /// );
950 /// # Ok(())
951 /// # }
952 /// # #[cfg(not(feature = "ipnet"))]
953 /// # fn main() {}
954 /// ```
955 ///
956 /// This function also yields the root node *if* it is part of the map:
957 ///
958 /// ```
959 /// # use prefix_trie::*; use prefix_trie::*;
960 /// # #[cfg(feature = "ipnet")]
961 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
962 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
963 /// let root = "0.0.0.0/0".parse()?;
964 /// pm.insert(root, 0);
965 /// assert_eq!(pm.cover(&"10.0.0.0/8".parse()?).collect::<Vec<_>>(), vec![(root, &0)]);
966 /// # Ok(())
967 /// # }
968 /// # #[cfg(not(feature = "ipnet"))]
969 /// # fn main() {}
970 /// ```
971 pub fn cover<'a>(&'a self, prefix: &P) -> Cover<'a, P, T> {
972 Cover::new(self, prefix)
973 }
974
975 /// Iterate over all prefixes in the map that cover `prefix`, including `prefix` itself if it is
976 /// present. The returned iterator yields reconstructed prefixes `P`.
977 ///
978 /// The iterator will always yield elements ordered by their prefix length, i.e., their depth in
979 /// the tree.
980 ///
981 /// ```
982 /// # use prefix_trie::*; use prefix_trie::*;
983 /// # #[cfg(feature = "ipnet")]
984 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
985 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
986 /// let p0 = "10.0.0.0/8".parse()?;
987 /// let p1 = "10.1.0.0/16".parse()?;
988 /// let p2 = "10.1.1.0/24".parse()?;
989 /// pm.insert(p0, 0);
990 /// pm.insert(p1, 1);
991 /// pm.insert(p2, 2);
992 /// pm.insert("10.1.2.0/24".parse()?, 3); // disjoint prefixes are not covered
993 /// pm.insert("10.1.1.0/25".parse()?, 4); // more specific prefixes are not covered
994 /// pm.insert("11.0.0.0/8".parse()?, 5); // Branch points that don't contain values are skipped
995 /// assert_eq!(pm.cover_keys(&p2).collect::<Vec<_>>(), vec![p0, p1, p2]);
996 /// # Ok(())
997 /// # }
998 /// # #[cfg(not(feature = "ipnet"))]
999 /// # fn main() {}
1000 /// ```
1001 pub fn cover_keys<'a>(&'a self, prefix: &P) -> CoverKeys<'a, P, T> {
1002 CoverKeys(Cover::new(self, prefix))
1003 }
1004
1005 /// Iterate over the values of all prefixes in the map that cover `prefix`, including `prefix`
1006 /// itself if it is present. The returned iterator yields `&'a T`.
1007 ///
1008 /// The iterator will always yield elements ordered by their prefix length, i.e., their depth in
1009 /// the tree.
1010 ///
1011 /// ```
1012 /// # use prefix_trie::*; use prefix_trie::*;
1013 /// # #[cfg(feature = "ipnet")]
1014 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1015 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1016 /// let p0 = "10.0.0.0/8".parse()?;
1017 /// let p1 = "10.1.0.0/16".parse()?;
1018 /// let p2 = "10.1.1.0/24".parse()?;
1019 /// pm.insert(p0, 0);
1020 /// pm.insert(p1, 1);
1021 /// pm.insert(p2, 2);
1022 /// pm.insert("10.1.2.0/24".parse()?, 3); // disjoint prefixes are not covered
1023 /// pm.insert("10.1.1.0/25".parse()?, 4); // more specific prefixes are not covered
1024 /// pm.insert("11.0.0.0/8".parse()?, 5); // Branch points that don't contain values are skipped
1025 /// assert_eq!(pm.cover_values(&p2).collect::<Vec<_>>(), vec![&0, &1, &2]);
1026 /// # Ok(())
1027 /// # }
1028 /// # #[cfg(not(feature = "ipnet"))]
1029 /// # fn main() {}
1030 /// ```
1031 pub fn cover_values<'a>(&'a self, prefix: &P) -> CoverValues<'a, P, T> {
1032 CoverValues(Cover::new(self, prefix))
1033 }
1034
1035 /// An iterator visiting all key-value pairs in lexicographic order. The iterator element type
1036 /// is `(P, &T)`, with reconstructed prefixes `P`.
1037 ///
1038 /// ```
1039 /// # use prefix_trie::*; use prefix_trie::*;
1040 /// # #[cfg(feature = "ipnet")]
1041 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1042 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1043 /// pm.insert("192.168.0.0/22".parse()?, 1);
1044 /// pm.insert("192.168.0.0/23".parse()?, 2);
1045 /// pm.insert("192.168.2.0/23".parse()?, 3);
1046 /// pm.insert("192.168.0.0/24".parse()?, 4);
1047 /// pm.insert("192.168.2.0/24".parse()?, 5);
1048 /// assert_eq!(
1049 /// pm.iter().collect::<Vec<_>>(),
1050 /// vec![
1051 /// ("192.168.0.0/22".parse()?, &1),
1052 /// ("192.168.0.0/23".parse()?, &2),
1053 /// ("192.168.0.0/24".parse()?, &4),
1054 /// ("192.168.2.0/23".parse()?, &3),
1055 /// ("192.168.2.0/24".parse()?, &5),
1056 /// ]
1057 /// );
1058 /// # Ok(())
1059 /// # }
1060 /// # #[cfg(not(feature = "ipnet"))]
1061 /// # fn main() {}
1062 /// ```
1063 #[inline(always)]
1064 pub fn iter(&self) -> Iter<'_, P, T> {
1065 self.into_iter()
1066 }
1067
1068 /// Get a mutable iterator over all key-value pairs. The order of this iterator is lexicographic.
1069 pub fn iter_mut(&mut self) -> IterMut<'_, P, T> {
1070 IterMut::new(&mut self.table)
1071 }
1072
1073 /// Iterate over all entries starting at `prefix`, in lexicographic order.
1074 ///
1075 /// This enables stateless, cursor-based pagination: pass the last-seen prefix to resume.
1076 ///
1077 /// - If `inclusive` is `true`, the iterator includes the entry at `prefix` (if present).
1078 /// - If `inclusive` is `false`, the iterator starts after `prefix`. Entries more specific than
1079 /// `prefix` (its children) are still yielded.
1080 ///
1081 /// If `prefix` is not present in the map, the iterator starts at the first entry that would come
1082 /// after `prefix` in lexicographic order, regardless of `inclusive`.
1083 ///
1084 /// ```
1085 /// # use prefix_trie::*;
1086 /// # #[cfg(feature = "ipnet")]
1087 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1088 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1089 /// pm.insert("10.0.0.0/8".parse()?, 1);
1090 /// pm.insert("10.1.0.0/16".parse()?, 2);
1091 /// pm.insert("10.2.0.0/16".parse()?, 3);
1092 /// pm.insert("10.2.0.0/24".parse()?, 4);
1093 /// pm.insert("10.3.0.0/16".parse()?, 5);
1094 /// pm.insert("10.4.0.0/16".parse()?, 6);
1095 ///
1096 /// // Inclusive: start at 10.2.0.0/16 and take the next 2 entries
1097 /// let page: Vec<_> = pm.iter_from(&"10.2.0.0/16".parse()?, true).take(3).collect();
1098 /// assert_eq!(page, vec![
1099 /// ("10.2.0.0/16".parse()?, &3),
1100 /// ("10.2.0.0/24".parse()?, &4),
1101 /// ("10.3.0.0/16".parse()?, &5),
1102 /// ]);
1103 ///
1104 /// // Exclusive: cursor pagination — skip last seen, fetch next page
1105 /// let last_seen: ipnet::Ipv4Net = "10.2.0.0/16".parse()?;
1106 /// let next_page: Vec<_> = pm.iter_from(&last_seen, false).take(3).collect();
1107 /// assert_eq!(next_page, vec![
1108 /// ("10.2.0.0/24".parse()?, &4),
1109 /// ("10.3.0.0/16".parse()?, &5),
1110 /// ("10.4.0.0/16".parse()?, &6)
1111 /// ]);
1112 /// # Ok(())
1113 /// # }
1114 /// # #[cfg(not(feature = "ipnet"))]
1115 /// # fn main() {}
1116 /// ```
1117 pub fn iter_from<'a>(&'a self, prefix: &P, inclusive: bool) -> Iter<'a, P, T> {
1118 let key = prefix.mask();
1119 let prefix_len = prefix.prefix_len() as u32;
1120 let stack = self.table.build_iter_stack_at(key, prefix_len, inclusive);
1121 Iter::from_stack(&self.table, stack)
1122 }
1123
1124 /// Return a mutable iterator starting at the given prefix in lexicographic order.
1125 ///
1126 /// If `inclusive` is `true`, the iterator includes the entry at `prefix` (if present).
1127 /// If `inclusive` is `false`, the iterator starts after `prefix`.
1128 ///
1129 /// If `prefix` is not present in the map, the iterator starts at the first entry that
1130 /// would come after `prefix` in lexicographic order, regardless of `inclusive`.
1131 ///
1132 /// ```
1133 /// # use prefix_trie::*;
1134 /// # #[cfg(feature = "ipnet")]
1135 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1136 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1137 /// pm.insert("10.0.0.0/8".parse()?, 1);
1138 /// pm.insert("10.1.0.0/16".parse()?, 2);
1139 /// pm.insert("10.2.0.0/16".parse()?, 3);
1140 ///
1141 /// // Mutate all entries starting from 10.1.0.0/16 (inclusive)
1142 /// pm.iter_from_mut(&"10.1.0.0/16".parse()?, true).for_each(|(_, v)| *v *= 10);
1143 /// assert_eq!(pm.get(&"10.0.0.0/8".parse()?), Some(&1));
1144 /// assert_eq!(pm.get(&"10.1.0.0/16".parse()?), Some(&20));
1145 /// assert_eq!(pm.get(&"10.2.0.0/16".parse()?), Some(&30));
1146 /// # Ok(())
1147 /// # }
1148 /// # #[cfg(not(feature = "ipnet"))]
1149 /// # fn main() {}
1150 /// ```
1151 pub fn iter_from_mut<'a>(&'a mut self, prefix: &P, inclusive: bool) -> IterMut<'a, P, T> {
1152 let key = prefix.mask();
1153 let prefix_len = prefix.prefix_len() as u32;
1154 let stack = self.table.build_iter_stack_at(key, prefix_len, inclusive);
1155 IterMut::from_stack(&mut self.table, stack)
1156 }
1157
1158 /// An iterator visiting all keys in lexicographic order. The iterator element type is
1159 /// reconstructed prefixes `P`.
1160 ///
1161 /// ```
1162 /// # use prefix_trie::*; use prefix_trie::*;
1163 /// # #[cfg(feature = "ipnet")]
1164 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1165 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1166 /// pm.insert("192.168.0.0/22".parse()?, 1);
1167 /// pm.insert("192.168.0.0/23".parse()?, 2);
1168 /// pm.insert("192.168.2.0/23".parse()?, 3);
1169 /// pm.insert("192.168.0.0/24".parse()?, 4);
1170 /// pm.insert("192.168.2.0/24".parse()?, 5);
1171 /// assert_eq!(
1172 /// pm.keys().collect::<Vec<_>>(),
1173 /// vec![
1174 /// "192.168.0.0/22".parse()?,
1175 /// "192.168.0.0/23".parse()?,
1176 /// "192.168.0.0/24".parse()?,
1177 /// "192.168.2.0/23".parse()?,
1178 /// "192.168.2.0/24".parse()?,
1179 /// ]
1180 /// );
1181 /// # Ok(())
1182 /// # }
1183 /// # #[cfg(not(feature = "ipnet"))]
1184 /// # fn main() {}
1185 /// ```
1186 #[inline(always)]
1187 pub fn keys(&self) -> Keys<'_, P, T> {
1188 Keys(self.iter())
1189 }
1190
1191 /// Creates a consuming iterator visiting all keys in lexicographic order. The iterator element
1192 /// type is reconstructed prefixes `P`.
1193 #[inline(always)]
1194 pub fn into_keys(self) -> IntoKeys<P, T> {
1195 IntoKeys(self.into_iter())
1196 }
1197
1198 /// An iterator visiting all values in lexicographic order. The iterator element type is `&T`.
1199 ///
1200 /// ```
1201 /// # use prefix_trie::*; use prefix_trie::*;
1202 /// # #[cfg(feature = "ipnet")]
1203 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1204 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1205 /// pm.insert("192.168.0.0/22".parse()?, 1);
1206 /// pm.insert("192.168.0.0/23".parse()?, 2);
1207 /// pm.insert("192.168.2.0/23".parse()?, 3);
1208 /// pm.insert("192.168.0.0/24".parse()?, 4);
1209 /// pm.insert("192.168.2.0/24".parse()?, 5);
1210 /// assert_eq!(pm.values().collect::<Vec<_>>(), vec![&1, &2, &4, &3, &5]);
1211 /// # Ok(())
1212 /// # }
1213 /// # #[cfg(not(feature = "ipnet"))]
1214 /// # fn main() {}
1215 /// ```
1216 #[inline(always)]
1217 pub fn values(&self) -> Values<'_, P, T> {
1218 Values(self.iter())
1219 }
1220
1221 /// Creates a consuming iterator visiting all values in lexicographic order. The iterator
1222 /// element type is `T`.
1223 #[inline(always)]
1224 pub fn into_values(self) -> IntoValues<P, T> {
1225 IntoValues(self.into_iter())
1226 }
1227
1228 /// Get a mutable iterator over all values. The order of this iterator is lexicographic.
1229 pub fn values_mut(&mut self) -> ValuesMut<'_, P, T> {
1230 ValuesMut(self.iter_mut())
1231 }
1232}
1233
1234impl<P, T> PrefixMap<P, T>
1235where
1236 P: Prefix,
1237{
1238 /// Iterate over `prefix` and all more-specific entries contained within it, including `prefix`
1239 /// itself if it is present. The iterator yields `(P, &'a T)`, with reconstructed prefixes `P`,
1240 /// in lexicographic order.
1241 ///
1242 /// **Note**: Consider using [`crate::AsView::view_at`] as an alternative.
1243 ///
1244 /// ```
1245 /// # use prefix_trie::*; use prefix_trie::*;
1246 /// # #[cfg(feature = "ipnet")]
1247 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1248 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1249 /// pm.insert("192.168.0.0/22".parse()?, 1);
1250 /// pm.insert("192.168.0.0/23".parse()?, 2);
1251 /// pm.insert("192.168.2.0/23".parse()?, 3);
1252 /// pm.insert("192.168.0.0/24".parse()?, 4);
1253 /// pm.insert("192.168.2.0/24".parse()?, 5);
1254 /// assert_eq!(
1255 /// pm.children(&"192.168.0.0/23".parse()?).collect::<Vec<_>>(),
1256 /// vec![
1257 /// ("192.168.0.0/23".parse()?, &2),
1258 /// ("192.168.0.0/24".parse()?, &4),
1259 /// ]
1260 /// );
1261 /// # Ok(())
1262 /// # }
1263 /// # #[cfg(not(feature = "ipnet"))]
1264 /// # fn main() {}
1265 /// ```
1266 pub fn children<'a>(&'a self, prefix: &P) -> Iter<'a, P, T> {
1267 let lex = iter::lpm_children_iter_start(&self.table, prefix);
1268 Iter::at_node(&self.table, lex)
1269 }
1270
1271 /// Iterate with mutable references over `prefix` and all more-specific entries contained within
1272 /// it, including `prefix` itself if it is present. The iterator yields `(P, &'a mut T)`, with
1273 /// reconstructed prefixes `P`, in lexicographic order.
1274 ///
1275 /// **Note**: Consider using [`crate::AsView::view_at`] on a mutable map reference as an
1276 /// alternative.
1277 ///
1278 /// ```
1279 /// # use prefix_trie::*; use prefix_trie::*;
1280 /// # #[cfg(feature = "ipnet")]
1281 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1282 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1283 /// pm.insert("192.168.0.0/22".parse()?, 1);
1284 /// pm.insert("192.168.0.0/23".parse()?, 2);
1285 /// pm.insert("192.168.0.0/24".parse()?, 3);
1286 /// pm.insert("192.168.2.0/23".parse()?, 4);
1287 /// pm.insert("192.168.2.0/24".parse()?, 5);
1288 /// pm.children_mut(&"192.168.0.0/23".parse()?).for_each(|(_, x)| *x *= 10);
1289 /// assert_eq!(
1290 /// pm.into_iter().collect::<Vec<_>>(),
1291 /// vec![
1292 /// ("192.168.0.0/22".parse()?, 1),
1293 /// ("192.168.0.0/23".parse()?, 20),
1294 /// ("192.168.0.0/24".parse()?, 30),
1295 /// ("192.168.2.0/23".parse()?, 4),
1296 /// ("192.168.2.0/24".parse()?, 5),
1297 /// ]
1298 /// );
1299 /// # Ok(())
1300 /// # }
1301 /// # #[cfg(not(feature = "ipnet"))]
1302 /// # fn main() {}
1303 /// ```
1304 pub fn children_mut<'a>(&'a mut self, prefix: &P) -> IterMut<'a, P, T> {
1305 let lex = iter::lpm_children_iter_start(&self.table, prefix);
1306 IterMut::at_node(&mut self.table, lex)
1307 }
1308
1309 /// Consume the map and iterate over `prefix` and all more-specific entries contained within it,
1310 /// including `prefix` itself if it is present. This returns an iterator over the owned entries.
1311 ///
1312 /// ```
1313 /// # use prefix_trie::*; use prefix_trie::*;
1314 /// # #[cfg(feature = "ipnet")]
1315 /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
1316 /// let mut pm: PrefixMap<ipnet::Ipv4Net, _> = PrefixMap::new();
1317 /// pm.insert("192.168.0.0/22".parse()?, 1);
1318 /// pm.insert("192.168.0.0/23".parse()?, 2);
1319 /// pm.insert("192.168.2.0/23".parse()?, 3);
1320 /// pm.insert("192.168.0.0/24".parse()?, 4);
1321 /// pm.insert("192.168.2.0/24".parse()?, 5);
1322 /// assert_eq!(
1323 /// pm.into_children(&"192.168.0.0/23".parse()?).collect::<Vec<_>>(),
1324 /// vec![
1325 /// ("192.168.0.0/23".parse()?, 2),
1326 /// ("192.168.0.0/24".parse()?, 4),
1327 /// ]
1328 /// );
1329 /// # Ok(())
1330 /// # }
1331 /// # #[cfg(not(feature = "ipnet"))]
1332 /// # fn main() {}
1333 /// ```
1334 pub fn into_children(self, prefix: &P) -> IntoIter<P, T> {
1335 let lex = iter::lpm_children_iter_start(&self.table, prefix);
1336 IntoIter::at_node(self.table, lex)
1337 }
1338
1339 /// Check the allocator: No memory should be unreferenced, and no memory should be aliased
1340 /// (double referenced). This function returns `true` if the allocator is in a correct state,
1341 /// and `false` if the memory is corrupt.
1342 #[cfg(test)]
1343 pub fn check_memory_alloc(&self) -> bool {
1344 self.table.check_memory_alloc()
1345 }
1346
1347 /// Count the live nodes reachable from the root, including the root itself.
1348 #[cfg(test)]
1349 pub(crate) fn num_nodes(&self) -> usize {
1350 self.table.num_nodes()
1351 }
1352}
1353
1354#[cfg(test)]
1355mod tests {
1356 use super::*;
1357 use crate::prefix::Prefix;
1358
1359 // Minimal prefix type: (repr, len)
1360 type P = (u32, u8);
1361
1362 fn p(repr: u32, len: u8) -> P {
1363 P::from_repr_len(repr, len)
1364 }
1365
1366 fn map_from(entries: &[(u32, u8, i32)]) -> PrefixMap<P, i32> {
1367 let mut m = PrefixMap::new();
1368 for &(repr, len, val) in entries {
1369 m.insert(p(repr, len), val);
1370 }
1371 m
1372 }
1373
1374 fn iter_keys(m: &PrefixMap<P, i32>) -> Vec<P> {
1375 m.iter().map(|(p, _)| p).collect()
1376 }
1377
1378 struct DropCounter(std::rc::Rc<std::cell::Cell<usize>>);
1379
1380 impl Drop for DropCounter {
1381 fn drop(&mut self) {
1382 self.0.set(self.0.get() + 1);
1383 }
1384 }
1385
1386 // ---- basic storage ----
1387
1388 #[test]
1389 fn test_insert_and_get_root() {
1390 // /0 prefix (the single root prefix covering everything)
1391 let mut m = PrefixMap::new();
1392 m.insert(p(0, 0), 42);
1393 assert_eq!(m.get(&p(0, 0)), Some(&42));
1394 assert_eq!(m.len(), 1);
1395 }
1396
1397 #[test]
1398 fn test_insert_root_and_child_separate() {
1399 // /0 and 0/1 must be stored as distinct entries
1400 let mut m = PrefixMap::new();
1401 m.insert(p(0, 0), 1);
1402 m.insert(p(0, 1), 2);
1403 assert_eq!(m.len(), 2);
1404 assert_eq!(m.get(&p(0, 0)), Some(&1));
1405 assert_eq!(m.get(&p(0, 1)), Some(&2));
1406 }
1407
1408 #[test]
1409 fn test_insert_sibling_prefixes() {
1410 // 0/1 (left half) and 0x80000000/1 (right half)
1411 let mut m = PrefixMap::new();
1412 m.insert(p(0x00000000, 1), 1);
1413 m.insert(p(0x80000000, 1), 2);
1414 assert_eq!(m.len(), 2);
1415 assert_eq!(m.get(&p(0x00000000, 1)), Some(&1));
1416 assert_eq!(m.get(&p(0x80000000, 1)), Some(&2));
1417 }
1418
1419 #[test]
1420 fn test_drop_drops_values() {
1421 let drops = std::rc::Rc::new(std::cell::Cell::new(0));
1422 {
1423 let mut m = PrefixMap::new();
1424 m.insert(p(0, 0), DropCounter(drops.clone()));
1425 m.insert(p(0, 1), DropCounter(drops.clone()));
1426 m.insert(p(0x80000000, 1), DropCounter(drops.clone()));
1427 }
1428 assert_eq!(drops.get(), 3);
1429 }
1430
1431 #[test]
1432 fn test_partial_into_iter_drop_drops_remaining_values() {
1433 let drops = std::rc::Rc::new(std::cell::Cell::new(0));
1434 {
1435 let mut m = PrefixMap::new();
1436 m.insert(p(0, 0), DropCounter(drops.clone()));
1437 m.insert(p(0, 1), DropCounter(drops.clone()));
1438 m.insert(p(0x80000000, 1), DropCounter(drops.clone()));
1439
1440 let mut iter = m.into_iter();
1441 drop(iter.next().unwrap());
1442 assert_eq!(drops.get(), 1);
1443 }
1444 assert_eq!(drops.get(), 3);
1445 }
1446
1447 #[test]
1448 fn test_children() {
1449 let mut m = PrefixMap::new();
1450 m.insert(p(0x0a000000, 8), 1);
1451 m.insert(p(0x0a010000, 16), 2);
1452 m.insert(p(0x0a020000, 16), 3);
1453 m.insert(p(0x0a010000, 24), 4);
1454 // View at 10.1.0.0/16: should include /16 and /24, not /8 or 10.2.0.0/16
1455 let got: Vec<_> = m
1456 .children(&p(0x0a010000, 16))
1457 .map(|(p, x)| (p, *x))
1458 .collect();
1459 assert_eq!(got, vec![(p(0x0a010000, 16), 2), (p(0x0a010000, 24), 4)]);
1460 }
1461
1462 // ---- iterator ordering ----
1463
1464 #[test]
1465 fn test_iter_order_root_before_child() {
1466 // /0 must come before 0/1 in iteration
1467 let m = map_from(&[(0, 0, 1), (0, 1, 2)]);
1468 let keys = iter_keys(&m);
1469 assert_eq!(keys, vec![p(0, 0), p(0, 1)], "root must precede child");
1470 }
1471
1472 #[test]
1473 fn test_iter_order_left_before_right() {
1474 // 0/1 must come before 0x80000000/1
1475 let m = map_from(&[(0x00000000, 1, 1), (0x80000000, 1, 2)]);
1476 let keys = iter_keys(&m);
1477 assert_eq!(
1478 keys,
1479 vec![p(0x00000000, 1), p(0x80000000, 1)],
1480 "left sibling must precede right sibling"
1481 );
1482 }
1483
1484 #[test]
1485 fn test_iter_order_root_then_siblings() {
1486 // /0, 0/1, 0x80000000/1: root first, then left, then right
1487 let m = map_from(&[(0, 0, 0), (0x00000000, 1, 1), (0x80000000, 1, 2)]);
1488 let keys = iter_keys(&m);
1489 assert_eq!(keys, vec![p(0, 0), p(0, 1), p(0x80000000, 1)]);
1490 }
1491
1492 #[test]
1493 fn test_iter_order_matches_hashmap_sort() {
1494 // The key invariant: PrefixMap iter order == sorted-by-Ord order of keys.
1495 // (Both share the property that parent comes before child and left before right,
1496 // since Prefix::Ord orders by (repr, len) which puts containing prefixes earlier.)
1497 let entries: &[(u32, u8, i32)] = &[
1498 (0x00000000, 0, 10),
1499 (0x00000000, 1, 20),
1500 (0x80000000, 1, 30),
1501 (0x00000000, 2, 40),
1502 (0x40000000, 2, 50),
1503 ];
1504 let m = map_from(entries);
1505 let mut expected: Vec<P> = entries.iter().map(|&(r, l, _)| p(r, l)).collect();
1506 expected.sort();
1507 assert_eq!(iter_keys(&m), expected);
1508 }
1509
1510 #[test]
1511 fn test_iter_order_5_6() {
1512 let entries = &[(0xd0000000, 5, 1), (0xd0000000, 6, 2)];
1513 let m = map_from(entries);
1514 let mut expected: Vec<P> = entries.iter().map(|&(r, l, _)| p(r, l)).collect();
1515 expected.sort();
1516 assert_eq!(iter_keys(&m), expected);
1517 }
1518
1519 #[test]
1520 fn test_default_iterators_are_empty() {
1521 assert_eq!(Iter::<P, i32>::default().count(), 0);
1522 assert_eq!(Keys::<P, i32>::default().count(), 0);
1523 assert_eq!(Values::<P, i32>::default().count(), 0);
1524 assert_eq!(IterMut::<P, i32>::default().count(), 0);
1525 assert_eq!(ValuesMut::<P, i32>::default().count(), 0);
1526 }
1527
1528 #[test]
1529 fn test_remove_children_leak() {
1530 // Reproduce the quickcheck minimal failing case exactly
1531 use crate::fuzzing::TestPrefix;
1532 let tp = |repr: u32, len: u8| -> TestPrefix { crate::Prefix::from_repr_len(repr, len) };
1533 let mut pmap: PrefixMap<TestPrefix, i32> = PrefixMap::new();
1534 // Minimal case from quickcheck: /6 contains /7, remove_children(/6) should remove both
1535 pmap.insert(tp(0x00000000, 6), 0);
1536 pmap.insert(tp(0x00000000, 7), 0);
1537 assert!(pmap.check_memory_alloc(), "leak before remove_children");
1538 pmap.remove_children(&tp(0x00000000, 6));
1539 assert!(pmap.check_memory_alloc(), "leak after remove_children");
1540 }
1541
1542 #[test]
1543 fn test_remove_children_deep_tree() {
1544 // With K=5, inserting at /11 creates nodes at depths 0, 5, and 10.
1545 // remove_children(&/5) fast-tracks via clear_node_and_children on the
1546 // depth-5 node. That node has a child at depth 10 whose allocation
1547 // must be freed AND the depth-5 node's child_bitmap/children_idx must
1548 // be cleared. Otherwise check_memory_alloc detects the stale pointer
1549 // (slot referenced by live node AND on free list).
1550 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1551 m.insert(p(0x00000000, 11), 100);
1552 assert!(m.check_memory_alloc(), "before remove_children");
1553
1554 m.remove_children(&p(0x00000000, 5));
1555 assert_eq!(m.len(), 0);
1556 assert!(
1557 m.check_memory_alloc(),
1558 "after remove_children: stale child pointers"
1559 );
1560
1561 // Re-insert at the same depth to verify no corruption from stale pointers.
1562 m.insert(p(0x00000000, 11), 200);
1563 assert_eq!(m.get(&p(0x00000000, 11)), Some(&200));
1564 assert!(m.check_memory_alloc(), "after re-insert");
1565 }
1566
1567 #[test]
1568 fn test_remove_children_deep_tree_slot_reuse() {
1569 // Regression test for stale child_bitmap/children_idx after
1570 // clear_node_and_children on a non-root node.
1571 //
1572 // The scenario:
1573 // 1. Insert at /11 → creates nodes at depths 0, 5, 10
1574 // 2. remove_children(&/5) → frees depth-10 node (slot goes to free list)
1575 // 3. Insert into a DIFFERENT subtree at /11 → allocator reuses the freed
1576 // slot for a completely different node
1577 // 4. If the depth-5 node's child_bitmap was left stale, traversal through
1578 // it would follow the old children_idx into the reused slot, reading a
1579 // node that belongs to a different subtree → data corruption
1580 //
1581 // With the fix (child_bitmap cleared), step 4 correctly sees "no children"
1582 // and creates a fresh allocation instead.
1583 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1584
1585 // Step 1: build a 3-level subtree rooted at the left side (bit 0)
1586 m.insert(p(0x00000000, 11), 100);
1587 assert!(m.check_memory_alloc(), "after initial insert");
1588
1589 // Step 2: wipe that subtree
1590 m.remove_children(&p(0x00000000, 5));
1591 assert_eq!(m.len(), 0);
1592 assert!(m.check_memory_alloc(), "after remove_children");
1593
1594 // Step 3: insert into a DIFFERENT subtree (bit 31 set → right side of root)
1595 // This forces the allocator to allocate a new depth-10 node, which reuses
1596 // the freed slot from step 2.
1597 m.insert(p(0x80000000, 11), 200);
1598 assert!(
1599 m.check_memory_alloc(),
1600 "after insert into different subtree"
1601 );
1602
1603 // Step 4: insert back into the ORIGINAL subtree path
1604 // If child_bitmap on the old depth-5 node is stale, find_or_insert_mut
1605 // follows the stale children_idx to the slot now owned by the right
1606 // subtree → wrong node → corruption.
1607 m.insert(p(0x00000000, 11), 300);
1608 assert!(
1609 m.check_memory_alloc(),
1610 "after re-insert into original subtree"
1611 );
1612
1613 // Verify both entries exist independently with correct values
1614 assert_eq!(m.len(), 2);
1615 assert_eq!(m.get(&p(0x00000000, 11)), Some(&300));
1616 assert_eq!(m.get(&p(0x80000000, 11)), Some(&200));
1617
1618 // Verify iteration yields exactly the two entries
1619 let mut entries: Vec<_> = m.iter().map(|(k, v)| (k, *v)).collect();
1620 entries.sort_by_key(|(k, _)| *k);
1621 assert_eq!(
1622 entries,
1623 vec![(p(0x00000000, 11), 300), (p(0x80000000, 11), 200)],
1624 );
1625 }
1626
1627 #[test]
1628 fn test_remove_children_prunes_empty_node_fast_path() {
1629 // /11 creates nodes at depths 0, 5, and 10. remove_children(&/5) takes the fast path
1630 // (5 % K == 0) and clears the depth-5 node. The emptied node must also be detached
1631 // from the root instead of remaining as an empty shell.
1632 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1633 m.insert(p(0x00000000, 11), 1);
1634 assert_eq!(m.num_nodes(), 3);
1635
1636 m.remove_children(&p(0x00000000, 5));
1637 assert_eq!(m.len(), 0);
1638 assert_eq!(m.num_nodes(), 1, "empty node shells left behind");
1639 assert!(m.check_memory_alloc());
1640 }
1641
1642 #[test]
1643 fn test_remove_children_prunes_empty_node_slow_path() {
1644 // /6 and /7 both live in the depth-5 node. remove_children(&/6) takes the slow path
1645 // (6 % K != 0) and removes both entries, leaving the node empty. It must be detached.
1646 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1647 m.insert(p(0x00000000, 6), 1);
1648 m.insert(p(0x00000000, 7), 2);
1649 assert_eq!(m.num_nodes(), 2);
1650
1651 m.remove_children(&p(0x00000000, 6));
1652 assert_eq!(m.len(), 0);
1653 assert_eq!(m.num_nodes(), 1, "empty node shells left behind");
1654 assert!(m.check_memory_alloc());
1655 }
1656
1657 #[test]
1658 fn test_remove_children_prunes_empty_ancestors() {
1659 // A single /16 creates nodes at depths 0, 5, 10, and 15. remove_children(&/12) frees
1660 // the depth-15 child and empties the depth-10 node; pruning must cascade through the
1661 // (now empty) depth-5 node all the way up to the root.
1662 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1663 m.insert(p(0x00000000, 16), 1);
1664 assert_eq!(m.num_nodes(), 4);
1665
1666 m.remove_children(&p(0x00000000, 12));
1667 assert_eq!(m.len(), 0);
1668 assert_eq!(m.num_nodes(), 1, "empty node shells left behind");
1669 assert!(m.check_memory_alloc());
1670 }
1671
1672 #[test]
1673 fn test_remove_children_keeps_nonempty_node() {
1674 // Two /6 entries share the depth-5 node, but only one is covered by the removed
1675 // prefix. The node must survive with the other entry intact.
1676 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1677 m.insert(p(0x00000000, 6), 1);
1678 m.insert(p(0x04000000, 6), 2);
1679 assert_eq!(m.num_nodes(), 2);
1680
1681 m.remove_children(&p(0x00000000, 6));
1682 assert_eq!(m.len(), 1);
1683 assert_eq!(m.get(&p(0x04000000, 6)), Some(&2));
1684 assert_eq!(m.num_nodes(), 2);
1685 assert!(m.check_memory_alloc());
1686 }
1687
1688 #[test]
1689 fn test_retain_leak() {
1690 use crate::fuzzing::TestPrefix;
1691 let tp = |repr: u32, len: u8| -> TestPrefix { crate::Prefix::from_repr_len(repr, len) };
1692 let mut pmap: PrefixMap<TestPrefix, i32> = PrefixMap::new();
1693 pmap.insert(tp(0xf0000000, 5), 0);
1694 pmap.insert(tp(0xf8000000, 5), 0);
1695 assert!(pmap.check_memory_alloc(), "leak before retain");
1696 pmap.retain(|pp, _| pp.prefix_len() < 2);
1697 assert!(pmap.check_memory_alloc(), "leak after retain");
1698 }
1699
1700 #[test]
1701 fn test_remove_children_minimal() {
1702 use crate::Prefix;
1703
1704 let mut pmap: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1705
1706 let p1 = <(u32, u8) as Prefix>::from_repr_len(0u32, 1);
1707 let p2 = <(u32, u8) as Prefix>::from_repr_len(0x40000000u32, 2); // bit 30 set
1708 let p3 = <(u32, u8) as Prefix>::from_repr_len(0x80000000u32, 2); // bit 31 set
1709
1710 pmap.insert(p1, 0);
1711 pmap.insert(p2, 1);
1712 pmap.insert(p3, 0);
1713
1714 pmap.remove_children(&p1);
1715
1716 let want: Vec<_> = vec![(p3, 0)];
1717 let actual: Vec<_> = pmap.into_iter().collect();
1718
1719 assert_eq!(want, actual, "mismatch in remove_children result");
1720 }
1721
1722 #[test]
1723 fn test_retain_minimal() {
1724 use crate::Prefix;
1725
1726 let mut pmap: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1727
1728 let p1 = <(u32, u8) as Prefix>::from_repr_len(0x50000000u32, 5);
1729 let p2 = <(u32, u8) as Prefix>::from_repr_len(0x50000000u32, 6);
1730 let p3 = <(u32, u8) as Prefix>::from_repr_len(0x5c000000u32, 6);
1731
1732 pmap.insert(p1, 0);
1733 pmap.insert(p2, 1);
1734 pmap.insert(p3, 1);
1735
1736 // Retain: keep elements where !(root.contains(p) && p.1 >= root.1 + 2)
1737 let predicate = |_: &(u32, u8), v: &i32| *v == 0;
1738
1739 let want: Vec<_> = pmap
1740 .iter()
1741 .filter(|(p, v)| predicate(p, v))
1742 .map(|(p, v)| (p, *v))
1743 .collect();
1744
1745 pmap.retain(predicate);
1746
1747 let actual: Vec<_> = pmap.into_iter().collect();
1748
1749 assert_eq!(want, actual, "mismatch in retain result");
1750 }
1751
1752 // /32 host routes require depth=30 with K=5, which means depth+K=35 > 32 (num_bits).
1753 // The `data_offset` and `get_mask` functions compute a shift of `32-30-5 = -3`,
1754 // which underflows u32: panics in debug, wraps in release causing collisions.
1755 mod max_prefix_length {
1756 use super::*;
1757
1758 #[test]
1759 fn distinct_offsets() {
1760 // Four /32 addresses whose bottom 2 bits differ (bits 30-31 of the u32).
1761 // In a correct implementation each must map to a distinct internal offset.
1762 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1763 let addrs: &[(u32, i32)] = &[
1764 (0x01020300, 1), // bits 30,31 = 0b00
1765 (0x01020301, 2), // bits 30,31 = 0b01
1766 (0x01020302, 3), // bits 30,31 = 0b10
1767 (0x01020303, 4), // bits 30,31 = 0b11
1768 ];
1769 for &(repr, val) in addrs {
1770 m.insert(p(repr, 32), val);
1771 }
1772 assert_eq!(
1773 m.len(),
1774 4,
1775 "all four /32s must be stored as distinct entries"
1776 );
1777 for &(repr, val) in addrs {
1778 assert_eq!(
1779 m.get(&p(repr, 32)),
1780 Some(&val),
1781 "wrong value for /32 addr {:#010x}",
1782 repr,
1783 );
1784 }
1785 }
1786
1787 #[test]
1788 fn lpm() {
1789 // /24 parent + /32 child: LPM on the /32 address must return the /32 value.
1790 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1791 m.insert(p(0x01020300, 24), 10); // 1.2.3.0/24
1792 m.insert(p(0x01020304, 32), 42); // 1.2.3.4/32
1793 assert_eq!(
1794 m.get_lpm(&p(0x01020304, 32)),
1795 Some((p(0x01020304, 32), &42))
1796 );
1797 assert_eq!(
1798 m.get_lpm(&p(0x01020305, 32)),
1799 Some((p(0x01020300, 24), &10))
1800 );
1801 }
1802
1803 #[test]
1804 fn iter() {
1805 // All /32 entries must appear in the iterator with correct (prefix, value) pairs.
1806 let addrs: &[(u32, i32)] = &[
1807 (0xc0000000, 10),
1808 (0xc0000001, 20),
1809 (0xc0000002, 30),
1810 (0xc0000003, 40),
1811 ];
1812 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1813 for &(repr, val) in addrs {
1814 m.insert(p(repr, 32), val);
1815 }
1816 let mut got: Vec<_> = m.iter().map(|(k, v)| (k.0, *v)).collect();
1817 got.sort_by_key(|&(r, _)| r);
1818 let want: Vec<_> = addrs.to_vec();
1819 assert_eq!(got, want);
1820 }
1821
1822 #[test]
1823 fn remove() {
1824 // Insert four /32s, remove two, verify the remaining two are correct.
1825 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1826 m.insert(p(0x01020300, 32), 1);
1827 m.insert(p(0x01020301, 32), 2);
1828 m.insert(p(0x01020302, 32), 3);
1829 m.insert(p(0x01020303, 32), 4);
1830
1831 assert_eq!(m.remove(&p(0x01020301, 32)), Some(2));
1832 assert_eq!(m.remove(&p(0x01020302, 32)), Some(3));
1833
1834 assert_eq!(m.len(), 2);
1835 assert_eq!(m.get(&p(0x01020300, 32)), Some(&1));
1836 assert_eq!(m.get(&p(0x01020301, 32)), None);
1837 assert_eq!(m.get(&p(0x01020302, 32)), None);
1838 assert_eq!(m.get(&p(0x01020303, 32)), Some(&4));
1839 }
1840
1841 #[test]
1842 fn remove_children_of_slash31() {
1843 // A /31 (no value) covers exactly two /32 host routes (.2 and .3).
1844 // A third /32 (.0) sits outside the /31.
1845 // remove_children(&/31) must drop the two covered /32s but leave the outsider.
1846 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1847 let parent = p(0x01020302, 31); // 1.2.3.2/31 (covers .2 and .3, no value)
1848 m.insert(p(0x01020300, 32), 10); // outside /31
1849 m.insert(p(0x01020302, 32), 1); // inside /31
1850 m.insert(p(0x01020303, 32), 2); // inside /31
1851
1852 m.remove_children(&parent);
1853
1854 assert_eq!(m.len(), 1);
1855 assert_eq!(
1856 m.get(&p(0x01020300, 32)),
1857 Some(&10),
1858 ".0/32 outside /31 must survive"
1859 );
1860 assert_eq!(m.get(&p(0x01020302, 32)), None, ".2/32 must be gone");
1861 assert_eq!(m.get(&p(0x01020303, 32)), None, ".3/32 must be gone");
1862 }
1863
1864 #[test]
1865 fn retain_slash32() {
1866 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1867 m.insert(p(0x01020300, 32), 1);
1868 m.insert(p(0x01020301, 32), 2);
1869 m.insert(p(0x01020302, 32), 3);
1870 m.insert(p(0x01020303, 32), 4);
1871 m.insert(p(0x01020300, 24), 10);
1872
1873 m.retain(|k, _| k.1 == 32 && k.0 % 2 == 0);
1874
1875 assert_eq!(m.len(), 2);
1876 assert_eq!(m.get(&p(0x01020300, 32)), Some(&1));
1877 assert_eq!(m.get(&p(0x01020301, 32)), None);
1878 assert_eq!(m.get(&p(0x01020302, 32)), Some(&3));
1879 assert_eq!(m.get(&p(0x01020303, 32)), None);
1880 assert_eq!(m.get(&p(0x01020300, 24)), None);
1881 assert!(m.check_memory_alloc(), "leak after retain on /32s");
1882 }
1883
1884 #[test]
1885 fn remove_children_of_slash32() {
1886 // remove_children of a /32 removes only that exact entry.
1887 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1888 m.insert(p(0x01020300, 32), 1);
1889 m.insert(p(0x01020301, 32), 2);
1890
1891 m.remove_children(&p(0x01020300, 32));
1892
1893 assert_eq!(m.len(), 1);
1894 assert_eq!(m.get(&p(0x01020300, 32)), None);
1895 assert_eq!(m.get(&p(0x01020301, 32)), Some(&2));
1896 assert!(m.check_memory_alloc(), "leak after remove_children /32");
1897 }
1898
1899 #[test]
1900 fn cover_slash32() {
1901 // cover() on a /32 should yield the /32 itself plus all ancestor prefixes.
1902 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1903 m.insert(p(0x01020300, 24), 10);
1904 m.insert(p(0x01020304, 32), 42);
1905
1906 let cover: Vec<_> = m.cover(&p(0x01020304, 32)).map(|(k, v)| (k, *v)).collect();
1907 assert_eq!(
1908 cover,
1909 vec![(p(0x01020300, 24), 10), (p(0x01020304, 32), 42)]
1910 );
1911 }
1912
1913 #[test]
1914 fn lpm_all_depths_to_slash32() {
1915 // Build a chain: /0, /5, /10, /15, /20, /25, /30, /32
1916 // LPM for the /32 address should return the /32 entry.
1917 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1918 let key = 0xAABBCCDDu32;
1919 for &len in &[0, 5, 10, 15, 20, 25, 30, 32] {
1920 m.insert(p(key, len), len as i32);
1921 }
1922 assert_eq!(m.len(), 8);
1923 assert_eq!(m.get_lpm(&p(key, 32)), Some((p(key, 32), &32)));
1924 // Verify each intermediate entry is retrievable
1925 for &len in &[0, 5, 10, 15, 20, 25, 30, 32] {
1926 assert_eq!(
1927 m.get(&p(key, len)),
1928 Some(&(len as i32)),
1929 "missing entry at /{}",
1930 len
1931 );
1932 }
1933 assert!(m.check_memory_alloc(), "leak with all-depth chain");
1934 }
1935
1936 #[test]
1937 fn clear_with_slash32() {
1938 let mut m: PrefixMap<(u32, u8), i32> = PrefixMap::new();
1939 for i in 0..4u32 {
1940 m.insert(p(0x01020300 | i, 32), i as i32);
1941 }
1942 m.insert(p(0x01020300, 24), 99);
1943 assert_eq!(m.len(), 5);
1944
1945 m.clear();
1946 assert_eq!(m.len(), 0);
1947 assert!(m.check_memory_alloc(), "leak after clear with /32s");
1948
1949 // Re-insert should work
1950 m.insert(p(0x01020304, 32), 1);
1951 assert_eq!(m.get(&p(0x01020304, 32)), Some(&1));
1952 }
1953
1954 /// Verify that clear_node_and_children is panic-safe: if T::drop() panics,
1955 /// Table::drop() during unwinding must not read already-uninit slots (UB).
1956 /// Under Miri, the old code would fail; this test documents the fix.
1957 #[test]
1958 fn clear_panic_safety() {
1959 use std::panic::{self, AssertUnwindSafe};
1960 use std::sync::atomic::{AtomicU32, Ordering};
1961
1962 static DROP_COUNT: AtomicU32 = AtomicU32::new(0);
1963 static PANIC_AT: AtomicU32 = AtomicU32::new(u32::MAX);
1964
1965 #[derive(Debug)]
1966 struct PanicDrop(#[allow(dead_code)] u32);
1967 impl Drop for PanicDrop {
1968 fn drop(&mut self) {
1969 if DROP_COUNT.fetch_add(1, Ordering::Relaxed)
1970 == PANIC_AT.load(Ordering::Relaxed)
1971 {
1972 panic!("intentional panic in Drop");
1973 }
1974 }
1975 }
1976
1977 // Use prefix lengths 0-4 so entries land in the ROOT node (depth 0,
1978 // covers /0../4 with K=5). This is critical: if the panic happens in
1979 // a child node, the root's child_bitmap is already cleared from a prior
1980 // iteration, so drop_values() never reaches the child — masking the UB.
1981 // With root-level entries, drop_values() immediately reads the root's
1982 // still-set bitmap and hits the uninit slots.
1983 let mut m: PrefixMap<(u32, u8), PanicDrop> = PrefixMap::new();
1984 m.insert(p(0x00000000, 0), PanicDrop(1));
1985 m.insert(p(0x00000000, 1), PanicDrop(2));
1986 m.insert(p(0x80000000, 1), PanicDrop(3));
1987
1988 // Panic on the 2nd drop during clear_node_and_children.
1989 DROP_COUNT.store(0, Ordering::Relaxed);
1990 PANIC_AT.store(1, Ordering::Relaxed);
1991
1992 let result = panic::catch_unwind(AssertUnwindSafe(|| {
1993 m.clear();
1994 }));
1995 assert!(result.is_err());
1996
1997 // Disable panics and drop the partially-cleared map. With the fix,
1998 // bitmaps are cleared before T::drop() runs, so drop_values() won't
1999 // read uninit slots. Without the fix, this would be UB under Miri.
2000 PANIC_AT.store(u32::MAX, Ordering::Relaxed);
2001 drop(m);
2002 }
2003 }
2004}