Skip to main content

dtrees_rs/caching/
trie.rs

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