1use crate::{
2    key_part::KeyPart,
3    prefix_tree_map::{Node, PrefixTreeMap},
4    std_lib::{BinaryHeap, Ordering},
5};
6
7#[derive(Clone)]
9pub struct PrefixTreeMapBuilder<E, W, V> {
10    root: NodeBuilder<E, W, V>,
11    max_wildcard_depth: usize,
12}
13
14#[derive(Clone)]
15struct NodeBuilder<E, W, V> {
16    key_part: Option<KeyPart<E, W>>,
17    value: Option<V>,
18    children: Option<BinaryHeap<NodeBuilder<E, W, V>>>,
19}
20
21impl<E, W, V> PrefixTreeMapBuilder<E, W, V>
22where
23    E: Clone + Ord,
24    W: Clone + Ord,
25{
26    pub fn new() -> Self {
28        Self {
29            root: NodeBuilder {
30                key_part: None,
31                value: None,
32                children: None,
33            },
34            max_wildcard_depth: 0,
35        }
36    }
37
38    pub fn insert(&mut self, key: impl IntoIterator<Item = KeyPart<E, W>>, value: V) {
44        let mut node = &mut self.root as *mut NodeBuilder<E, W, V>;
45        let mut wildcard_depth = 0;
46
47        for key_part in key {
48            if key_part.is_wildcard() {
49                wildcard_depth += 1;
50            }
51
52            if unsafe { (*node).children.is_none() } {
53                let mut children = BinaryHeap::new();
54                children.push(NodeBuilder::new(key_part));
55
56                unsafe {
57                    (*node).children = Some(children);
58                }
59
60                let child = unsafe {
61                    (*node)
62                        .children
63                        .as_ref()
64                        .unwrap_unchecked()
65                        .peek()
66                        .unwrap_unchecked()
67                };
68
69                let child_const_ptr = child as *const NodeBuilder<E, W, V>;
70                node = child_const_ptr as *mut NodeBuilder<E, W, V>;
71            } else {
72                let children = unsafe { (*node).children.as_mut().unwrap_unchecked() };
73
74                if let Some(child) = children
75                    .iter()
76                    .find(|child| child.key_part.as_ref() == Some(&key_part))
77                {
78                    let child_const_ptr = child as *const NodeBuilder<E, W, V>;
79                    node = child_const_ptr as *mut NodeBuilder<E, W, V>;
80                } else {
81                    let key_part_cloned = key_part.clone();
82                    children.push(NodeBuilder::new(key_part_cloned));
83
84                    let child = unsafe {
85                        children
86                            .iter()
87                            .find(|child| child.key_part.as_ref() == Some(&key_part))
88                            .unwrap_unchecked()
89                    };
90
91                    let child_const_ptr = child as *const NodeBuilder<E, W, V>;
92                    node = child_const_ptr as *mut NodeBuilder<E, W, V>;
93                }
94            }
95        }
96
97        unsafe {
98            (*node).value = Some(value);
99        }
100
101        self.max_wildcard_depth = self.max_wildcard_depth.max(wildcard_depth);
102    }
103
104    pub fn insert_exact(&mut self, key: impl IntoIterator<Item = E>, value: V) {
106        self.insert(key.into_iter().map(KeyPart::Exact), value);
107    }
108
109    pub fn build(self) -> PrefixTreeMap<E, W, V> {
111        PrefixTreeMap {
112            root: Self::node_builder_to_node(self.root),
113            max_wildcard_depth: self.max_wildcard_depth,
114        }
115    }
116
117    fn node_builder_to_node(node_builder: NodeBuilder<E, W, V>) -> Node<E, W, V> {
118        let key_part = node_builder.key_part;
119        let value = node_builder.value;
120
121        let children = node_builder.children.map(|children| {
122            children
123                .into_sorted_vec()
124                .into_iter()
125                .map(Self::node_builder_to_node)
126                .collect()
127        });
128
129        Node {
130            key_part,
131            value,
132            children,
133        }
134    }
135}
136
137impl<E, W, V> Default for PrefixTreeMapBuilder<E, W, V>
138where
139    E: Clone + Ord,
140    W: Clone + Ord,
141{
142    fn default() -> Self {
143        Self::new()
144    }
145}
146
147impl<E, W, V> NodeBuilder<E, W, V>
148where
149    E: Clone + Ord,
150    W: Clone + Ord,
151{
152    fn new(key_part: KeyPart<E, W>) -> Self {
153        Self {
154            key_part: Some(key_part),
155            value: None,
156            children: None,
157        }
158    }
159}
160
161impl<E, W, V> PartialEq for NodeBuilder<E, W, V>
162where
163    E: Clone + Ord,
164    W: Clone + Ord,
165{
166    fn eq(&self, other: &Self) -> bool {
167        self.key_part == other.key_part
168    }
169}
170
171impl<E, W, V> Eq for NodeBuilder<E, W, V>
172where
173    E: Clone + Ord,
174    W: Clone + Ord,
175{
176}
177
178impl<E, W, V> PartialOrd for NodeBuilder<E, W, V>
179where
180    E: Clone + Ord,
181    W: Clone + Ord,
182{
183    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
184        self.key_part.partial_cmp(&other.key_part)
185    }
186}
187
188impl<E, W, V> Ord for NodeBuilder<E, W, V>
189where
190    E: Clone + Ord,
191    W: Clone + Ord,
192{
193    fn cmp(&self, other: &Self) -> Ordering {
194        self.key_part.cmp(&other.key_part)
195    }
196}