general_sam/utils/
suffixwise.rs

1//! Utilities to store suffix-wise data in a suffix automaton.
2
3use std::collections::LinkedList;
4use std::convert::Infallible;
5use std::ops::Deref;
6
7use crate::rope::{Rope, RopeBase, RopeData, RopeUntaggedInner, TreapBasedRopeBase};
8use crate::{
9    GeneralSam, GeneralSamState, SAM_NIL_NODE_ID, SAM_ROOT_NODE_ID, TransitionTable, TravelEvent,
10    TrieNodeAlike,
11};
12
13#[derive(Clone, Default, Debug)]
14pub struct SuffixwiseData<Inner: RopeData + Default> {
15    data: Rope<Inner>,
16    min_suf_len: usize,
17    max_suf_len: usize,
18}
19
20impl<Inner: RopeData + Default> SuffixwiseData<Inner> {
21    pub fn get_rope(&self) -> &Rope<Inner> {
22        &self.data
23    }
24
25    pub fn get_min_suf_len(&self) -> usize {
26        self.min_suf_len
27    }
28
29    pub fn get_max_suf_len(&self) -> usize {
30        self.max_suf_len
31    }
32
33    pub fn map<NewInner: RopeData + Default, F: FnOnce(&Rope<Inner>) -> Rope<NewInner>>(
34        &self,
35        f: F,
36    ) -> SuffixwiseData<NewInner> {
37        SuffixwiseData {
38            data: f(&self.data),
39            min_suf_len: self.min_suf_len,
40            max_suf_len: self.max_suf_len,
41        }
42    }
43
44    pub fn get(&self, suf_len: usize) -> Option<Inner> {
45        if self.data.is_empty()
46            || self.max_suf_len == 0
47            || self.min_suf_len == 0
48            || suf_len < self.min_suf_len
49            || suf_len > self.max_suf_len
50        {
51            return None;
52        }
53        Some(
54            self.data
55                .query(suf_len - self.min_suf_len)
56                .expect("invalid suffixwise data")
57                .as_ref()
58                .deref()
59                .to_owned(),
60        )
61    }
62
63    pub fn build_from_sam<
64        TransTable: TransitionTable,
65        Iter: IntoIterator<Item = (usize, Inner)>,
66        FInit: FnMut(usize) -> Iter,
67    >(
68        sam: &GeneralSam<TransTable>,
69        mut f_init: FInit,
70    ) -> Vec<Self> {
71        let mut res = vec![Self::default(); sam.num_of_nodes()];
72        for node_id in sam.get_topo_and_suf_len_sorted_node_ids().iter().copied() {
73            assert_ne!(node_id, SAM_NIL_NODE_ID);
74
75            let node = sam.get_node(node_id).expect("invalid GeneralSam");
76            let node_data = res
77                .get_mut(node_id)
78                .unwrap_or_else(|| panic!("invalid node id: {}", node_id));
79
80            node_data.max_suf_len = node.max_suffix_len();
81
82            if node_id == SAM_ROOT_NODE_ID {
83                node_data.min_suf_len = 0;
84
85                node_data.data = Rope::new(Inner::default());
86            } else {
87                let parent_id = node.get_suffix_parent_id();
88                let parent = sam.get_node(parent_id).expect("invalid GeneralSam");
89
90                node_data.min_suf_len = parent.max_suffix_len() + 1;
91
92                assert_eq!(
93                    node_data.data.len(),
94                    node_data.max_suf_len - node_data.min_suf_len + 1
95                );
96
97                for (len, data) in f_init(node_id) {
98                    assert!(len >= node_data.min_suf_len && len <= node_data.max_suf_len);
99                    let (left, right) = node_data.data.split(len - node_data.min_suf_len);
100                    let (_, right) = right.split(1);
101                    node_data.data = left.merge(&Rope::new(data)).merge(&right);
102                }
103
104                assert_eq!(
105                    node_data.data.len(),
106                    node_data.max_suf_len - node_data.min_suf_len + 1
107                );
108            }
109
110            node.get_trans()
111                .transitions()
112                .copied()
113                .for_each(|target_id| {
114                    res[target_id].data = res[target_id].data.merge(&res[node_id].data)
115                });
116        }
117        res
118    }
119}
120
121#[derive(Clone, Debug)]
122pub struct SuffixInTrie<Digested: Clone> {
123    pub digested_trie_node: Digested,
124    pub seq_len: usize,
125}
126
127pub type UntaggedSuffixData<Inner> = SuffixwiseData<RopeUntaggedInner<Inner>>;
128pub type SuffixInTrieData<D> = UntaggedSuffixData<Option<SuffixInTrie<D>>>;
129
130impl<Digested: Clone> SuffixInTrieData<Digested> {
131    pub fn build<
132        TransTable: TransitionTable,
133        TN: TrieNodeAlike<InnerType = TransTable::KeyType>,
134        F: FnMut(&TN) -> Digested,
135    >(
136        sam: &GeneralSam<TransTable>,
137        trie_node: TN,
138        mut f: F,
139    ) -> Vec<Self> {
140        let mut sam_to_data = vec![LinkedList::<SuffixInTrie<Digested>>::new(); sam.num_of_nodes()];
141        let callback =
142            |event: TravelEvent<(&GeneralSamState<_, &GeneralSam<_>>, &TN), _, _>| -> Result<_, Infallible> {
143                match event {
144                    crate::TravelEvent::Pop((sam_state, trie_state), len) => {
145                        if trie_state.is_accepting() {
146                            sam_to_data[sam_state.node_id].push_back(SuffixInTrie {
147                                digested_trie_node: f(trie_state),
148                                seq_len: len,
149                            });
150                        }
151                        Ok(len)
152                    }
153                    crate::TravelEvent::PushRoot(_) => Ok(0),
154                    crate::TravelEvent::Push(_, len, _) => Ok(len + 1),
155                }
156            };
157        sam.get_root_state().bfs_along(trie_node, callback).unwrap();
158        Self::build_from_sam(sam, |node_id| {
159            sam_to_data[node_id]
160                .iter()
161                .map(|x| (x.seq_len, Some(x.clone()).into()))
162        })
163    }
164}
general_sam/utils/suffixwise.rs

general_sam/utils/
suffixwise.rs
