dtrees_rs/caching/
trie.rs

1use crate::caching::entry::{CacheEntry, CacheEntryUpdater};
2use crate::caching::helpers::{CacheKey, Index};
3use crate::caching::Caching;
4use rustc_hash::FxHashMap;
5
6#[derive(Debug)]
7struct TrieNode {
8    index: usize,
9    entry: CacheEntry,
10    children: FxHashMap<usize, usize>,
11}
12
13impl Default for TrieNode {
14    fn default() -> Self {
15        Self {
16            index: usize::MAX,
17            children: FxHashMap::default(),
18            entry: CacheEntry::default(),
19        }
20    }
21}
22
23impl TrieNode {
24    pub fn new(item: usize) -> Self {
25        Self {
26            index: usize::MAX,
27            children: FxHashMap::default(),
28            entry: CacheEntry::new(item),
29        }
30    }
31}
32
33/// A [`Caching`] implementation as a trie of items stored in an arena.
34#[derive(Default)]
35pub struct Trie {
36    arena: Vec<TrieNode>,
37    root_index: usize,
38}
39
40impl Caching for Trie {
41    fn init(&mut self) -> Index {
42        if !self.arena.is_empty() {
43            self.arena.clear();
44        }
45        let index = self.add_default_root();
46        Index::new_at(index)
47    }
48
49    fn root_index(&mut self) -> Index {
50        Index::new_at(self.get_root_index())
51    }
52
53    fn root(&self) -> Option<&CacheEntry> {
54        self.get_node(self.get_root_index()).map(|node| &node.entry)
55    }
56
57    fn insert(&mut self, path: &[usize]) -> Index {
58        let mut current_index = self.root_index;
59        let mut is_new = false;
60        for &item in path {
61            if let Some(&child_index) = self
62                .get_node(current_index)
63                .and_then(|node| node.children.get(&item))
64            {
65                current_index = child_index;
66            } else {
67                is_new = true;
68                current_index = self.create_child(current_index, item);
69            }
70        }
71
72        let leaf_is_infinite = self
73            .get_node(current_index)
74            .map_or(true, |node| !node.entry.has_finite_leaf_error());
75
76        if is_new || leaf_is_infinite {
77            return Index::New(current_index);
78        }
79        Index::Existing(current_index)
80    }
81
82    fn node(&self, key: &CacheKey) -> Option<&CacheEntry> {
83        match key {
84            CacheKey::Index(index) => self.get_node(*index).map(|node| &node.entry),
85            CacheKey::Path(path) => self.find(path),
86        }
87    }
88
89    fn contains(&self, key: &CacheKey) -> bool {
90        self.node(key).is_some()
91    }
92
93    fn update_root(&mut self) -> Option<CacheEntryUpdater<'_>> {
94        self.get_node_mut(self.root_index)
95            .map(|node| CacheEntryUpdater::new(&mut node.entry))
96    }
97
98    fn update_node(&mut self, key: &CacheKey) -> Option<CacheEntryUpdater<'_>> {
99        match key {
100            CacheKey::Index(index) => self
101                .get_node_mut(*index)
102                .map(|node| CacheEntryUpdater::new(&mut node.entry)),
103            CacheKey::Path(path) => self.find_mut(path).map(CacheEntryUpdater::new),
104        }
105    }
106
107    fn size(&self) -> usize {
108        self.arena.len()
109    }
110
111    fn is_empty(&self) -> bool {
112        self.arena.is_empty()
113    }
114}
115
116impl Trie {
117    /// An empty trie.
118    pub fn new() -> Self {
119        Self {
120            arena: Vec::new(),
121            root_index: 0,
122        }
123    }
124
125    fn add_node(&mut self, parent_index: usize, mut node: TrieNode) -> usize {
126        let new_index = self.arena.len();
127        node.index = new_index;
128        let item = node.entry.item();
129        self.arena.push(node);
130        if new_index != 0 {
131            if let Some(parent) = self.arena.get_mut(parent_index) {
132                parent.children.insert(item, new_index);
133            }
134        }
135
136        new_index
137    }
138
139    fn find(&self, itemset: &[usize]) -> Option<&CacheEntry> {
140        let mut current_index = self.root_index;
141        for &item in itemset {
142            let current_node = self.get_node(current_index)?;
143            current_index = *current_node.children.get(&item)?;
144        }
145        self.get_node(current_index).map(|node| &node.entry)
146    }
147
148    fn find_mut(&mut self, itemset: &[usize]) -> Option<&mut CacheEntry> {
149        let mut current_index = self.root_index;
150        for &item in itemset {
151            let current_node = self.get_node(current_index)?;
152            current_index = *current_node.children.get(&item)?;
153        }
154        self.get_node_mut(current_index).map(|node| &mut node.entry)
155    }
156
157    fn add_root(&mut self, root: TrieNode) -> usize {
158        self.root_index = self.add_node(0, root);
159        self.root_index
160    }
161
162    fn add_default_root(&mut self) -> usize {
163        self.root_index = self.add_root(TrieNode::default());
164        self.root_index
165    }
166
167    #[inline]
168    fn get_root_index(&self) -> usize {
169        self.root_index
170    }
171
172    #[inline]
173    fn get_node(&self, index: usize) -> Option<&TrieNode> {
174        self.arena.get(index)
175    }
176
177    #[inline]
178    fn get_node_mut(&mut self, index: usize) -> Option<&mut TrieNode> {
179        self.arena.get_mut(index)
180    }
181
182    fn create_child(&mut self, parent: usize, item: usize) -> usize {
183        let node = TrieNode::new(item);
184        self.add_node(parent, node)
185    }
186}
187
188#[cfg(test)]
189mod trie_test {
190    use crate::caching::helpers::Index;
191    use crate::caching::trie::{Trie, TrieNode};
192    use crate::caching::Caching;
193
194    #[test]
195    fn test_cache_init() {
196        let mut cache = Trie::new();
197        assert!(cache.is_empty());
198
199        let root_data = TrieNode::default();
200        cache.add_root(root_data);
201        cache.add_node(0, TrieNode::default());
202        println!("Cache 0 {:?}", cache.arena[0].children);
203
204        assert!(!cache.is_empty());
205
206        let mut itemset = Vec::new();
207        itemset.push(0);
208        itemset.push(1);
209        itemset.push(3);
210
211        let idx = cache.insert(&itemset);
212
213        match idx {
214            Index::NewUnknown => {}
215            Index::New(_) | Index::Existing(_) => {}
216        }
217
218        itemset.remove(1);
219    }
220}
dtrees_rs/caching/trie.rs

dtrees_rs/caching/
trie.rs
