general_sam/
trie.rs

1//! Trie, supporting `TrieNodeAlike`.
2
3use std::{borrow::Borrow, ops::Deref};
4
5use crate::{ConstructiveTransitionTable, GeneralSamNodeID, TransitionTable, TrieNodeAlike};
6
7pub type TrieNodeID = GeneralSamNodeID;
8pub const TRIE_NIL_NODE_ID: TrieNodeID = 0;
9pub const TRIE_ROOT_NODE_ID: TrieNodeID = 1;
10
11#[derive(Clone, Debug)]
12pub struct TrieNode<TransTable: TransitionTable> {
13    trans: TransTable,
14    parent: TrieNodeID,
15    pub accept: bool,
16}
17
18#[derive(Clone, Debug)]
19pub struct Trie<TransTable: TransitionTable> {
20    node_pool: Vec<TrieNode<TransTable>>,
21}
22
23#[derive(Debug)]
24pub struct TrieState<TransTable: TransitionTable, TrieRef: Deref<Target = Trie<TransTable>>> {
25    pub trie: TrieRef,
26    pub node_id: TrieNodeID,
27}
28
29impl<TransTable: TransitionTable, TrieRef: Deref<Target = Trie<TransTable>> + Clone> Clone
30    for TrieState<TransTable, TrieRef>
31{
32    fn clone(&self) -> Self {
33        Self {
34            trie: self.trie.clone(),
35            node_id: self.node_id,
36        }
37    }
38}
39
40impl<TransTable: ConstructiveTransitionTable> TrieNode<TransTable> {
41    fn new(parent: TrieNodeID) -> Self {
42        Self {
43            trans: Default::default(),
44            parent,
45            accept: Default::default(),
46        }
47    }
48}
49
50impl<TransTable: TransitionTable> TrieNode<TransTable> {
51    pub fn get_trans(&self) -> &TransTable {
52        &self.trans
53    }
54
55    pub fn get_parent(&self) -> TrieNodeID {
56        self.parent
57    }
58
59    fn alter_trans_table<NewTableType: TransitionTable<KeyType = TransTable::KeyType>>(
60        &self,
61    ) -> TrieNode<NewTableType> {
62        TrieNode {
63            trans: NewTableType::from_kv_iter(self.trans.iter()),
64            parent: self.parent,
65            accept: self.accept,
66        }
67    }
68}
69
70impl<TransTable: ConstructiveTransitionTable> Default for Trie<TransTable> {
71    fn default() -> Self {
72        Self {
73            node_pool: vec![
74                TrieNode::new(TRIE_NIL_NODE_ID),
75                TrieNode::new(TRIE_NIL_NODE_ID),
76            ],
77        }
78    }
79}
80
81impl<TransTable: TransitionTable> Trie<TransTable> {
82    pub fn num_of_nodes(&self) -> usize {
83        self.node_pool.len()
84    }
85
86    pub fn get_state(&self, node_id: TrieNodeID) -> TrieState<TransTable, &Trie<TransTable>> {
87        if node_id >= self.node_pool.len() {
88            return TrieState {
89                trie: self,
90                node_id: TRIE_NIL_NODE_ID,
91            };
92        }
93        TrieState {
94            trie: self,
95            node_id,
96        }
97    }
98
99    pub fn get_node(&self, node_id: TrieNodeID) -> Option<&TrieNode<TransTable>> {
100        self.node_pool.get(node_id)
101    }
102
103    pub fn get_root_node(&self) -> &TrieNode<TransTable> {
104        self.get_node(TRIE_ROOT_NODE_ID).unwrap()
105    }
106
107    pub fn get_root_state(&self) -> TrieState<TransTable, &Trie<TransTable>> {
108        self.get_state(TRIE_ROOT_NODE_ID)
109    }
110
111    pub fn alter_trans_table<NewTableType: TransitionTable<KeyType = TransTable::KeyType>>(
112        &self,
113    ) -> Trie<NewTableType> {
114        Trie {
115            node_pool: self
116                .node_pool
117                .iter()
118                .map(|x| x.alter_trans_table())
119                .collect(),
120        }
121    }
122}
123
124impl<TransTable: ConstructiveTransitionTable> Trie<TransTable> {
125    fn alloc_node(&mut self, parent: TrieNodeID) -> TrieNodeID {
126        let node_id = self.node_pool.len();
127        self.node_pool.push(TrieNode::new(parent));
128        node_id
129    }
130
131    pub fn insert<Iter: IntoIterator<Item = TransTable::KeyType>>(
132        &mut self,
133        iter: Iter,
134    ) -> TrieNodeID {
135        let mut current = TRIE_ROOT_NODE_ID;
136        iter.into_iter().for_each(|t| {
137            current = match self.node_pool[current].trans.get(&t) {
138                Some(v) => *v,
139                None => {
140                    let new_node_id = self.alloc_node(current);
141                    self.node_pool[current].trans.insert(t, new_node_id);
142                    new_node_id
143                }
144            };
145        });
146        self.node_pool[current].accept = true;
147        current
148    }
149}
150
151impl<TransTable: ConstructiveTransitionTable<KeyType = u8>> Trie<TransTable> {
152    pub fn insert_bytes<S: AsRef<[u8]>>(&mut self, bytes: S) -> TrieNodeID {
153        self.insert(bytes.as_ref().iter().copied())
154    }
155}
156
157impl<TransTable: ConstructiveTransitionTable<KeyType = char>> Trie<TransTable> {
158    pub fn insert_chars<S: AsRef<str>>(&mut self, s: S) -> TrieNodeID {
159        self.insert(s.as_ref().chars())
160    }
161}
162
163impl<TransTable: TransitionTable, TrieRef: Deref<Target = Trie<TransTable>>>
164    TrieState<TransTable, TrieRef>
165{
166    pub fn inner_as_ref(&self) -> TrieState<TransTable, &Trie<TransTable>> {
167        TrieState {
168            trie: &self.trie,
169            node_id: self.node_id,
170        }
171    }
172
173    pub fn is_nil(&self) -> bool {
174        self.node_id == TRIE_NIL_NODE_ID
175    }
176
177    pub fn is_root(&self) -> bool {
178        self.node_id == TRIE_ROOT_NODE_ID
179    }
180
181    pub fn get_node(&self) -> Option<&TrieNode<TransTable>> {
182        self.trie.get_node(self.node_id)
183    }
184
185    pub fn goto_parent(&mut self) {
186        if let Some(node) = self.get_node() {
187            self.node_id = node.parent;
188        } else {
189            self.node_id = TRIE_NIL_NODE_ID;
190        }
191    }
192
193    pub fn goto<K: Borrow<TransTable::KeyType>>(&mut self, t: K) {
194        if let Some(node) = self.get_node() {
195            self.node_id = node
196                .trans
197                .get(t.borrow())
198                .copied()
199                .unwrap_or(TRIE_NIL_NODE_ID)
200        } else {
201            self.node_id = TRIE_NIL_NODE_ID;
202        }
203    }
204
205    pub fn feed<Iter: IntoIterator<Item = TransTable::KeyType>>(&mut self, iter: Iter) {
206        iter.into_iter().for_each(|x| self.goto(&x));
207    }
208
209    pub fn feed_ref<K: Borrow<TransTable::KeyType>, Iter: IntoIterator<Item = K>>(
210        &mut self,
211        iter: Iter,
212    ) {
213        iter.into_iter().for_each(|x| self.goto(x));
214    }
215
216    pub fn feed_slice<S: AsRef<[TransTable::KeyType]>>(&mut self, slice: S) {
217        self.feed_ref(slice.as_ref().iter())
218    }
219}
220
221#[derive(Clone, Debug)]
222pub struct NextTrieStateIter<'s, TransTable: TransitionTable> {
223    trie: &'s Trie<TransTable>,
224    iter: TransTable::IterType<'s>,
225}
226
227impl<'s, TransTable: TransitionTable> TrieNodeAlike
228    for TrieState<TransTable, &'s Trie<TransTable>>
229{
230    type InnerType = TransTable::KeyType;
231    type NextStateIter = NextTrieStateIter<'s, TransTable>;
232
233    fn is_accepting(&self) -> bool {
234        self.get_node().map(|x| x.accept).unwrap_or(false)
235    }
236
237    fn next_states(self) -> Self::NextStateIter {
238        let iter = self.trie.get_node(self.node_id).unwrap().trans.iter();
239        NextTrieStateIter {
240            trie: self.trie,
241            iter,
242        }
243    }
244}
245
246impl<'s, TransTable: TransitionTable> Iterator for NextTrieStateIter<'s, TransTable> {
247    type Item = (
248        TransTable::KeyType,
249        TrieState<TransTable, &'s Trie<TransTable>>,
250    );
251
252    fn next(&mut self) -> Option<Self::Item> {
253        self.iter
254            .next()
255            .map(|(t, next_node_id)| (t.clone(), self.trie.get_state(*next_node_id)))
256    }
257}