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