Skip to main content

ts_bart/iptrie/
insert.rs

1use crate::{
2    BaseIndex, StrideOpsExt,
3    iptrie::util,
4    node::{Child, StrideOps},
5};
6
7/// Insert a route into the trie rooted at `node` with the given prefix.
8#[inline]
9pub fn insert<N>(node: &mut N, prefix: ipnet::IpNet, val: N::T) -> Option<N::T>
10where
11    N: StrideOps,
12{
13    // Wrapper to avoid exposing the depth parameter
14    insert_inner(node, prefix.trunc(), val, 0)
15}
16
17pub fn insert_inner<N>(
18    mut node: &mut N,
19    prefix: ipnet::IpNet,
20    val: N::T,
21    depth: usize,
22) -> Option<N::T>
23where
24    N: StrideOps,
25{
26    let (pfx_full_strides, pfx_overflow_bits) = util::stride_count_and_overflow(&prefix);
27
28    for (depth, &octet) in util::ip_octets(&prefix.addr())
29        .iter()
30        .enumerate()
31        .skip(depth)
32    {
33        if depth == pfx_full_strides {
34            return node.insert_prefix(BaseIndex::from_prefix(octet, pfx_overflow_bits), val);
35        }
36
37        if !node.child_bitset().test(octet as _) {
38            let child = if util::is_fringe(depth, &prefix) {
39                Child::Fringe(val)
40            } else {
41                Child::Leaf { prefix, value: val }
42            };
43
44            return node.insert_child(octet, child)?.into_value();
45        }
46
47        // Child replacement/recursion is done in three phases:
48        //
49        // - If the child at `octet` is a match for target prefix, the child already
50        //   exists, just replace the value and return.
51        // - If it's not a match and it's not a path node, it needs to be upgraded to
52        //   one, because there's now more than one prefix under the same address at
53        //   this depth. Remove the child, replace it with a path node, then reinsert
54        //   the old route value into the new child.
55        // - Now the child must be a path node: descend into it.
56        //
57        // The phases need to be separate in order for the &mut borrow to the child to
58        // be released before manipulating the current node.
59
60        // invariant: child is present (bitset test succeeded)
61        let replace_child: bool = match node.get_child_mut(octet).unwrap() {
62            Child::Leaf {
63                prefix: leaf_pfx,
64                value: leaf_value,
65            } if prefix == leaf_pfx => {
66                return Some(core::mem::replace(leaf_value, val));
67            }
68            Child::Fringe(old) if util::is_fringe(depth, &prefix) => {
69                return Some(core::mem::replace(old, val));
70            }
71
72            Child::Leaf { .. } => true,
73            Child::Fringe(..) => true,
74
75            Child::Path(..) => false,
76        };
77
78        if replace_child {
79            // invariant: child is still present, so we must be able to replace it
80            let removed = node.insert_child(octet, N::default().into_child()).unwrap();
81
82            let Some(Child::Path(n)) = node.get_child_mut(octet) else {
83                unreachable!();
84            };
85
86            match removed {
87                Child::Fringe(old) => {
88                    n.insert_prefix(BaseIndex::new(1), old);
89                }
90                Child::Leaf {
91                    prefix: leaf_prefix,
92                    value: leaf_value,
93                } => {
94                    // Recursive call: known to be depth 1 because the node was just inserted and
95                    // is empty.
96                    insert_inner(n, leaf_prefix, leaf_value, depth + 1);
97                }
98
99                _ => unreachable!(),
100            }
101        }
102
103        // invariant: child is present _and_ a path node
104        let next = node.get_child_mut(octet).unwrap().into_node().unwrap();
105        node = next;
106    }
107
108    None
109}