Skip to main content

hara_native/lang/data/
trie.rs

1use std::cell::Cell;
2use std::collections::BTreeMap;
3use std::rc::Rc;
4
5use crate::lang::hash::JavaHash;
6use crate::lang::protocol::{
7    HashType, IAssoc, IColl, IConj, ICount, IDisplay, IDissoc, IEmpty, IEquality, IFind, IHash,
8    ILookup, IMetadata, IMutable, IObjType, IPersistent, IToMutable, IToPersistent, MetaType,
9    ObjType,
10};
11
12#[derive(Debug, Clone)]
13struct Node<V> {
14    children: BTreeMap<char, Rc<Node<V>>>,
15    value: Option<V>,
16}
17impl<V> Default for Node<V> {
18    fn default() -> Self {
19        Self {
20            children: BTreeMap::new(),
21            value: None,
22        }
23    }
24}
25fn assoc_node<V: Clone>(node: &Rc<Node<V>>, chars: &[char], value: V) -> Rc<Node<V>> {
26    if chars.is_empty() {
27        return Rc::new(Node {
28            children: node.children.clone(),
29            value: Some(value),
30        });
31    }
32    let mut children = node.children.clone();
33    let child = children
34        .get(&chars[0])
35        .cloned()
36        .unwrap_or_else(|| Rc::new(Node::default()));
37    children.insert(chars[0], assoc_node(&child, &chars[1..], value));
38    Rc::new(Node {
39        children,
40        value: node.value.clone(),
41    })
42}
43fn dissoc_node<V: Clone>(node: &Rc<Node<V>>, chars: &[char]) -> Option<Rc<Node<V>>> {
44    let mut children = node.children.clone();
45    let value = if chars.is_empty() {
46        None
47    } else {
48        let child = children.get(&chars[0])?;
49        match dissoc_node(child, &chars[1..]) {
50            Some(next) => {
51                children.insert(chars[0], next);
52            }
53            None => {
54                children.remove(&chars[0]);
55            }
56        }
57        node.value.clone()
58    };
59    if value.is_none() && children.is_empty() {
60        None
61    } else {
62        Some(Rc::new(Node { children, value }))
63    }
64}
65#[derive(Debug, Clone)]
66pub struct Standard<V> {
67    metadata: Option<Rc<crate::lang::data::Metadata>>,
68    root: Rc<Node<V>>,
69    size: usize,
70}
71impl<V> Default for Standard<V> {
72    fn default() -> Self {
73        Self {
74            metadata: None,
75            root: Rc::new(Node::default()),
76            size: 0,
77        }
78    }
79}
80impl<V: Clone> Standard<V> {
81    pub fn new() -> Self {
82        Self::default()
83    }
84    pub fn len(&self) -> usize {
85        self.size
86    }
87    pub fn is_empty(&self) -> bool {
88        self.size == 0
89    }
90    pub fn get(&self, key: &str) -> Option<&V> {
91        let mut node = self.root.as_ref();
92        for ch in key.chars() {
93            node = node.children.get(&ch)?.as_ref();
94        }
95        node.value.as_ref()
96    }
97    pub fn assoc_value(&self, key: impl Into<String>, value: V) -> Self {
98        let key = key.into();
99        let added = self.get(&key).is_none();
100        let chars = key.chars().collect::<Vec<_>>();
101        Self {
102            metadata: self.metadata.clone(),
103            root: assoc_node(&self.root, &chars, value),
104            size: self.size + usize::from(added),
105        }
106    }
107    pub fn dissoc_value(&self, key: &str) -> Self {
108        if self.get(key).is_none() {
109            return self.clone();
110        }
111        let chars = key.chars().collect::<Vec<_>>();
112        Self {
113            metadata: self.metadata.clone(),
114            root: dissoc_node(&self.root, &chars).unwrap_or_else(|| Rc::new(Node::default())),
115            size: self.size - 1,
116        }
117    }
118    pub fn entries(&self) -> Vec<(String, &V)> {
119        fn collect<'a, V>(n: &'a Node<V>, prefix: &mut String, out: &mut Vec<(String, &'a V)>) {
120            if let Some(v) = &n.value {
121                out.push((prefix.clone(), v));
122            }
123            for (ch, child) in &n.children {
124                prefix.push(*ch);
125                collect(child, prefix, out);
126                prefix.pop();
127            }
128        }
129        let mut out = Vec::with_capacity(self.size);
130        collect(&self.root, &mut String::new(), &mut out);
131        out
132    }
133    pub fn iter(&self) -> impl Iterator<Item = String> + '_ {
134        self.entries().into_iter().map(|(key, _)| key)
135    }
136}
137impl<V: Clone> ICount for Standard<V> {
138    fn count(&self) -> usize {
139        self.len()
140    }
141}
142impl<V: Clone> IFind<String> for Standard<V> {
143    type Output = (String, V);
144    fn find(&self, key: &String) -> Option<Self::Output> {
145        self.get(key).map(|v| (key.clone(), v.clone()))
146    }
147}
148impl<V: Clone> ILookup<String, V> for Standard<V> {
149    type Keys = std::vec::IntoIter<String>;
150    type Values = std::vec::IntoIter<V>;
151    fn keys(&self) -> Self::Keys {
152        self.iter().collect::<Vec<_>>().into_iter()
153    }
154    fn vals(&self) -> Self::Values {
155        self.entries()
156            .into_iter()
157            .map(|(_, v)| v.clone())
158            .collect::<Vec<_>>()
159            .into_iter()
160    }
161}
162impl<V: Clone> IAssoc<String, V> for Standard<V> {
163    type Output = Self;
164    fn assoc(&self, k: String, v: V) -> Self {
165        self.assoc_value(k, v)
166    }
167}
168impl<V: Clone> IDissoc<String> for Standard<V> {
169    type Output = Self;
170    fn dissoc(&self, k: &String) -> Self {
171        self.dissoc_value(k)
172    }
173}
174impl<V: Clone + Default> IConj<String> for Standard<V> {
175    type Output = Self;
176    fn conj(&self, k: String) -> Self {
177        self.assoc_value(k, V::default())
178    }
179}
180impl<V: Clone> IEmpty for Standard<V> {
181    type Output = Self;
182    fn empty(&self) -> Self {
183        Self::new().with_meta(self.metadata.clone())
184    }
185}
186impl<V: Clone> IMetadata for Standard<V> {
187    type Metadata = Rc<crate::lang::data::Metadata>;
188    fn meta(&self) -> Option<&Self::Metadata> {
189        self.metadata.as_ref()
190    }
191    fn with_meta(&self, metadata: Option<Self::Metadata>) -> Self {
192        Self {
193            metadata,
194            ..self.clone()
195        }
196    }
197
198    fn metatype(&self) -> MetaType {
199        MetaType::Map
200    }
201}
202impl<V: Clone> IPersistent for Standard<V> {}
203impl<V: Clone> IntoIterator for Standard<V> {
204    type Item = String;
205    type IntoIter = std::vec::IntoIter<String>;
206    fn into_iter(self) -> Self::IntoIter {
207        self.iter().collect::<Vec<_>>().into_iter()
208    }
209}
210impl<V: Clone + PartialEq> IEquality for Standard<V> {
211    fn equality(&self, other: &Self) -> bool {
212        self.len() == other.len() && self.entries().iter().all(|(k, v)| other.get(k) == Some(*v))
213    }
214}
215impl<V: Clone + std::fmt::Debug> IDisplay for Standard<V> {
216    fn display(&self) -> String {
217        format!(
218            "#{{{}}}",
219            self.entries()
220                .iter()
221                .map(|(k, v)| format!("{k:?} {v:?}"))
222                .collect::<Vec<_>>()
223                .join(" ")
224        )
225    }
226}
227impl<V: Clone + std::hash::Hash + JavaHash> IHash for Standard<V> {
228    fn hash_calc(&self, hash_type: HashType) -> u64 {
229        // Java Trie.hashCalc override: acc = "::MAP".hashCode(), then
230        // acc += hash(key) + hash(value) per entry (NOT the MapEntry
231        // composition used by maps). Keys are plain Strings, so they hash
232        // via Java String.hashCode under every hash type (see lang::hash).
233        let mut acc = crate::lang::hash::hash_seed("MAP") as i64;
234        for (k, v) in self.entries() {
235            acc = acc
236                .wrapping_add(crate::lang::hash::java_string_hash(&k) as i64)
237                .wrapping_add(v.java_hash(hash_type));
238        }
239        acc as u64
240    }
241}
242impl<V: Clone + std::fmt::Debug> IObjType for Standard<V> {
243    fn obj_type(&self) -> ObjType {
244        ObjType::Map
245    }
246}
247impl<V> IColl<String> for Standard<V>
248where
249    V: Clone + Default + PartialEq + std::hash::Hash + JavaHash + std::fmt::Debug,
250{
251    fn start_string(&self) -> &'static str {
252        "#{"
253    }
254    fn end_string(&self) -> &'static str {
255        "}"
256    }
257}
258impl<V: Clone> IToMutable for Standard<V> {
259    type Mutable = Mutable<V>;
260    fn to_mutable(&self) -> Self::Mutable {
261        Mutable {
262            editable: Cell::new(true),
263            trie: self.clone(),
264        }
265    }
266}
267#[derive(Debug, Clone)]
268pub struct Mutable<V> {
269    editable: Cell<bool>,
270    trie: Standard<V>,
271}
272impl<V: Clone> Mutable<V> {
273    fn check(&self) {
274        assert!(self.editable.get(), "mutable trie used after to_persistent")
275    }
276    pub fn assoc(&mut self, k: impl Into<String>, v: V) -> &mut Self {
277        self.check();
278        self.trie = self.trie.assoc_value(k, v);
279        self
280    }
281    pub fn dissoc(&mut self, k: &str) -> &mut Self {
282        self.check();
283        self.trie = self.trie.dissoc_value(k);
284        self
285    }
286}
287impl<V: Clone> std::ops::Deref for Mutable<V> {
288    type Target = Standard<V>;
289    fn deref(&self) -> &Self::Target {
290        self.check();
291        &self.trie
292    }
293}
294impl<V> IMutable for Mutable<V> {}
295impl<V: Clone> IToPersistent for Mutable<V> {
296    type Persistent = Standard<V>;
297    fn to_persistent(&mut self) -> Self::Persistent {
298        self.check();
299        self.editable.set(false);
300        self.trie.clone()
301    }
302}
303#[cfg(test)]
304mod tests {
305    use super::Standard;
306    #[test]
307    fn updates_empty_and_mutable_preserve_metadata() {
308        use crate::lang::protocol::{IEmpty, IMetadata, IToMutable, IToPersistent};
309        let trie = Standard::new()
310            .assoc_value("cat", 1)
311            .with_meta(Some(crate::lang::data::Metadata::document("doc")));
312        assert_eq!(
313            trie.assoc_value("car", 2).meta().map(|m| m.doc().unwrap()),
314            Some("doc")
315        );
316        assert_eq!(
317            trie.dissoc_value("cat").meta().map(|m| m.doc().unwrap()),
318            Some("doc")
319        );
320        assert_eq!(trie.empty().meta().map(|m| m.doc().unwrap()), Some("doc"));
321        let mut mutable = trie.to_mutable();
322        mutable.assoc("car", 2);
323        assert_eq!(
324            mutable.to_persistent().meta().map(|m| m.doc().unwrap()),
325            Some("doc")
326        );
327    }
328
329    #[test]
330    fn shares_prefixes_and_iterates_lexically() {
331        let a = Standard::new()
332            .assoc_value("car", 1)
333            .assoc_value("cat", 2)
334            .assoc_value("dog", 3);
335        assert_eq!(a.iter().collect::<Vec<_>>(), vec!["car", "cat", "dog"]);
336        let b = a.dissoc_value("cat");
337        assert_eq!(b.get("cat"), None);
338        assert_eq!(a.get("cat"), Some(&2));
339        assert_eq!(b.get("car"), Some(&1));
340    }
341
342    #[test]
343    fn dissoc_prunes_childless_ancestors_bottom_up() {
344        // Java dissocHelper: a node is removed once it is non-terminal and
345        // childless, and the removal cascades up the spine.
346        let trie = Standard::new()
347            .assoc_value("cat", 1)
348            .assoc_value("cats", 2)
349            .assoc_value("car", 3);
350        // removing the leaf word prunes the 's' node; "cat" stays terminal
351        let trie = trie.dissoc_value("cats");
352        assert_eq!(trie.get("cats"), None);
353        assert_eq!(trie.get("cat"), Some(&1));
354        // removing "car" prunes the 'r' node but keeps the "cat" spine
355        let trie = trie.dissoc_value("car");
356        assert_eq!(trie.get("car"), None);
357        assert_eq!(trie.entries().len(), 1);
358        // removing the last word prunes the whole spine back to an empty root
359        let trie = trie.dissoc_value("cat");
360        assert_eq!(trie.len(), 0);
361        assert!(trie.root.children.is_empty());
362        assert!(trie.root.value.is_none());
363    }
364
365    #[test]
366    fn dissoc_prefix_clears_terminal_but_keeps_children() {
367        // Java: dissoc of a word that is a prefix of another word only
368        // clears the terminal flag/value; the subtree is untouched.
369        let trie = Standard::new()
370            .assoc_value("cat", 1)
371            .assoc_value("cats", 2)
372            .dissoc_value("cat");
373        assert_eq!(trie.get("cat"), None);
374        assert_eq!(trie.get("cats"), Some(&2));
375        assert_eq!(trie.len(), 1);
376        assert_eq!(trie.iter().collect::<Vec<_>>(), vec!["cats"]);
377    }
378
379    #[test]
380    fn empty_string_key_sets_the_root_terminal() {
381        // Java assoc("") walks zero chars and marks the root node terminal;
382        // iteration yields "" first (it is a prefix of every key).
383        let trie = Standard::new().assoc_value("", 7).assoc_value("a", 1);
384        assert_eq!(trie.get(""), Some(&7));
385        assert_eq!(trie.len(), 2);
386        assert_eq!(trie.iter().collect::<Vec<_>>(), vec!["", "a"]);
387        // dissoc("") clears the root terminal but keeps the children
388        let trie = trie.dissoc_value("");
389        assert_eq!(trie.get(""), None);
390        assert_eq!(trie.get("a"), Some(&1));
391        assert_eq!(trie.len(), 1);
392    }
393
394    #[test]
395    fn dissoc_absent_prefix_or_divergent_key_is_a_noop() {
396        let trie = Standard::new().assoc_value("cat", 1);
397        for key in ["cow", "ca", "cats", "dog"] {
398            let next = trie.dissoc_value(key);
399            assert_eq!(next.len(), 1);
400            assert_eq!(next.get("cat"), Some(&1));
401        }
402    }
403}