general_sam/
trie.rs

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

general_sam/
trie.rs
