general_sam/
trie_alike.rs

1//! A trait for constructing `GeneralSam` from structures that form a trie,
2//! and some utilities to construct `GeneralSam` from iterators.
3
4use std::collections::VecDeque;
5
6#[derive(Clone, Debug)]
7pub enum TravelEvent<'s, NodeType, ExtraType, KeyType> {
8    PushRoot(NodeType),
9    Push(NodeType, &'s ExtraType, KeyType),
10    Pop(NodeType, ExtraType),
11}
12
13/// This trait provides the essential interfaces required by `GeneralSam`
14/// to construct a suffix automaton from structures that form a trie (prefix tree).
15pub trait TrieNodeAlike {
16    type InnerType;
17    type NextStateIter: Iterator<Item = (Self::InnerType, Self)>;
18    fn is_accepting(&self) -> bool;
19    fn next_states(self) -> Self::NextStateIter;
20
21    fn bfs_travel<
22        ErrorType,
23        ExtraType,
24        F: FnMut(TravelEvent<&Self, ExtraType, Self::InnerType>) -> Result<ExtraType, ErrorType>,
25    >(
26        self,
27        mut callback: F,
28    ) -> Result<(), ErrorType>
29    where
30        Self: Sized,
31    {
32        let mut queue = VecDeque::new();
33
34        let extra = callback(TravelEvent::PushRoot(&self))?;
35        queue.push_back((self, extra));
36
37        while let Some((state, cur_extra)) = queue.pop_front() {
38            let cur_extra = callback(TravelEvent::Pop(&state, cur_extra))?;
39
40            for (t, v) in state.next_states() {
41                let next_extra = callback(TravelEvent::Push(&v, &cur_extra, t))?;
42                queue.push_back((v, next_extra));
43            }
44        }
45        Ok(())
46    }
47
48    fn dfs_travel<
49        ErrorType,
50        ExtraType,
51        F: FnMut(TravelEvent<&Self, ExtraType, Self::InnerType>) -> Result<ExtraType, ErrorType>,
52    >(
53        self,
54        mut callback: F,
55    ) -> Result<(), ErrorType>
56    where
57        Self: Clone,
58    {
59        let mut stack = Vec::new();
60
61        let extra = callback(TravelEvent::PushRoot(&self))?;
62        stack.push((self.clone(), self.next_states(), extra));
63
64        while !stack.is_empty() {
65            if let Some((_, iter, extra)) = stack.last_mut()
66                && let Some((key, next_state)) = iter.next()
67            {
68                let new_extra = callback(TravelEvent::Push(&next_state, extra, key))?;
69                stack.push((next_state.clone(), next_state.next_states(), new_extra));
70                continue;
71            }
72            if let Some((cur, _, extra)) = stack.pop() {
73                callback(TravelEvent::Pop(&cur, extra))?;
74            }
75        }
76        Ok(())
77    }
78}
79
80/// This struct implements `TrieNodeAlike` for any iterator.
81///
82/// It can be used to construct a suffix automaton directly from a sequence.
83#[derive(Clone, Debug)]
84pub struct IterAsChain<Iter: Iterator> {
85    pub iter: Iter,
86    pub val: Option<Iter::Item>,
87}
88
89pub struct IterAsChainNextStateIter<Iter: Iterator> {
90    pub state: Option<IterAsChain<Iter>>,
91}
92
93impl<Iter: IntoIterator> From<Iter> for IterAsChain<Iter::IntoIter> {
94    fn from(iter: Iter) -> Self {
95        let mut iter = iter.into_iter();
96        let val = iter.next();
97        Self { iter, val }
98    }
99}
100
101impl<Iter: Iterator> TrieNodeAlike for IterAsChain<Iter> {
102    type InnerType = Iter::Item;
103    type NextStateIter = IterAsChainNextStateIter<Iter>;
104
105    fn is_accepting(&self) -> bool {
106        self.val.is_none()
107    }
108
109    fn next_states(self) -> Self::NextStateIter {
110        Self::NextStateIter { state: Some(self) }
111    }
112}
113
114impl<Iter: Iterator> Iterator for IterAsChainNextStateIter<Iter> {
115    type Item = (Iter::Item, IterAsChain<Iter>);
116
117    fn next(&mut self) -> Option<Self::Item> {
118        self.state.take().and_then(|mut chain| {
119            if let Some(v) = chain.val {
120                chain.val = chain.iter.next();
121                Some((v, chain))
122            } else {
123                None
124            }
125        })
126    }
127}