radixtarget_rust/node/
ip.rs

1use super::base::BaseNode;
2use ipnet::{IpNet, Ipv4Net, Ipv6Net};
3use std::collections::HashMap;
4
5#[derive(Debug, Clone)]
6pub struct IPNode {
7    pub children: HashMap<u64, Box<IPNode>>,
8    pub network: Option<IpNet>,
9}
10
11impl Default for IPNode {
12    fn default() -> Self {
13        Self::new()
14    }
15}
16
17impl IPNode {
18    pub fn new() -> Self {
19        IPNode {
20            children: HashMap::new(),
21            network: None,
22        }
23    }
24    // Recursively prune dead child nodes (no network and no children).
25    pub fn prune(&mut self) -> usize {
26        let mut pruned = 0;
27        loop {
28            let mut new_pruned = 0;
29            let keys: Vec<u64> = self.children.keys().copied().collect();
30            for key in keys {
31                if let Some(child) = self.children.get_mut(&key) {
32                    new_pruned += child.prune();
33                    if child.network.is_none() && child.children.is_empty() {
34                        self.children.remove(&key);
35                        new_pruned += 1;
36                    }
37                }
38            }
39            if new_pruned == 0 {
40                break;
41            }
42            pruned += new_pruned;
43        }
44        pruned
45    }
46    /// Public iterator over all mergeable pairs of child nodes and their supernet.
47    /// Yields (left_net, right_net, supernet) for each mergeable pair in the subtree.
48    pub fn mergeable_pairs(&self) -> Vec<(IpNet, IpNet, IpNet)> {
49        let mut pairs = Vec::new();
50        // Check this node
51        if self.children.len() == 2 {
52            let mut iter = self.children.values();
53            let left = iter.next().unwrap();
54            let right = iter.next().unwrap();
55            if let (Some(left_net), Some(right_net)) = (&left.network, &right.network)
56                && left_net.prefix_len() == right_net.prefix_len()
57                && (matches!(left_net, IpNet::V4(_)) && matches!(right_net, IpNet::V4(_))
58                    || matches!(left_net, IpNet::V6(_)) && matches!(right_net, IpNet::V6(_)))
59            {
60                let prefix = left_net.prefix_len();
61                let supernet = match (left_net, right_net) {
62                    (IpNet::V4(l), IpNet::V4(r)) => {
63                        if prefix > 0 && prefix <= 32 {
64                            let mask = if prefix == 32 {
65                                0xffffffffu32
66                            } else {
67                                !(0xffffffffu32 >> prefix)
68                            };
69                            let l_addr = u32::from(l.network());
70                            let r_addr = u32::from(r.network());
71                            if (l_addr ^ r_addr) == (1 << (32 - prefix)) {
72                                let min_addr = l_addr.min(r_addr) & mask;
73                                Some(IpNet::V4(
74                                    Ipv4Net::new(min_addr.into(), prefix - 1).unwrap(),
75                                ))
76                            } else {
77                                None
78                            }
79                        } else {
80                            None
81                        }
82                    }
83                    (IpNet::V6(l), IpNet::V6(r)) => {
84                        if prefix > 0 && prefix <= 128 {
85                            let mask = if prefix == 128 {
86                                0xffffffffffffffffffffffffffffffffu128
87                            } else {
88                                !(0xffffffffffffffffffffffffffffffffu128 >> prefix)
89                            };
90                            let l_addr = u128::from(l.network());
91                            let r_addr = u128::from(r.network());
92                            if (l_addr ^ r_addr) == (1 << (128 - prefix)) {
93                                let min_addr = l_addr.min(r_addr) & mask;
94                                Some(IpNet::V6(
95                                    Ipv6Net::new(min_addr.into(), prefix - 1).unwrap(),
96                                ))
97                            } else {
98                                None
99                            }
100                        } else {
101                            None
102                        }
103                    }
104                    _ => None,
105                };
106                if let Some(supernet) = supernet {
107                    pairs.push((*left_net, *right_net, supernet));
108                }
109            }
110        }
111        // Recurse into children
112        for child in self.children.values() {
113            pairs.extend(child.mergeable_pairs());
114        }
115        pairs
116    }
117}
118
119impl BaseNode for IPNode {
120    fn is_dead(&self) -> bool {
121        self.network.is_none() && self.children.is_empty()
122    }
123    fn children_mut(&mut self) -> &mut HashMap<u64, Box<IPNode>> {
124        &mut self.children
125    }
126    fn children(&self) -> &HashMap<u64, Box<IPNode>> {
127        &self.children
128    }
129    fn host_string(&self) -> Option<String> {
130        self.network.as_ref().map(|n| n.to_string())
131    }
132}