Skip to main content

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}