1use 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}