lzcomplexity_core/
suffix_array.rs1#[derive(Clone, Debug, Default, PartialEq, Eq)]
11pub struct SuffixArray {
12 pub sa: Vec<u32>,
13 pub lcp: Vec<u32>,
14 pub n: u32,
15}
16
17impl SuffixArray {
18 pub fn new() -> Self {
19 Self::default()
20 }
21
22 pub fn clear(&mut self) {
23 self.sa.clear();
24 self.lcp.clear();
25 self.n = 0;
26 }
27}
28
29pub fn build_suffix_array(text: &[u8]) -> Vec<u32> {
34 if text.is_empty() {
35 return Vec::new();
36 }
37 const SORT_CUTOFF: usize = 2048;
43 if text.len() < SORT_CUTOFF {
44 sort_based_sa(text)
45 } else {
46 let (_, sa) = cdivsufsort::sort(text).into_parts();
47 sa.into_iter().map(|x| x as u32).collect()
48 }
49}
50
51fn sort_based_sa(text: &[u8]) -> Vec<u32> {
54 let n = text.len();
55 let mut idx: Vec<u32> = (0..n as u32).collect();
56 idx.sort_unstable_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..]));
57 idx
58}
59
60pub fn build_lcp(text: &[u8], sa: &[u32]) -> Vec<u32> {
65 let n = sa.len();
66 if n == 0 {
67 return Vec::new();
68 }
69 let mut rank = vec![0u32; n];
70 for i in 0..n {
71 rank[sa[i] as usize] = i as u32;
72 }
73 let mut lcp = vec![0u32; n];
74 let mut h: u32 = 0;
75 for i in 0..n {
76 let r = rank[i] as usize;
77 if r > 0 {
78 let j = sa[r - 1] as usize;
79 while i + (h as usize) < n
80 && j + (h as usize) < n
81 && text[i + (h as usize)] == text[j + (h as usize)]
82 {
83 h += 1;
84 }
85 lcp[r] = h;
86 h = h.saturating_sub(1);
87 } else {
88 h = 0;
89 }
90 }
91 lcp
92}
93
94pub fn build(text: &[u8]) -> SuffixArray {
96 let sa = build_suffix_array(text);
97 let lcp = build_lcp(text, &sa);
98 SuffixArray {
99 sa,
100 lcp,
101 n: text.len() as u32,
102 }
103}
104
105#[cfg(test)]
106mod tests {
107 use super::*;
108
109 #[test]
110 fn banana_sa() {
111 let sa = build_suffix_array(b"banana");
112 assert_eq!(sa, vec![5, 3, 1, 0, 4, 2]);
113 }
114
115 #[test]
116 fn test_text_sa() {
117 let sa = build_suffix_array(b"test text");
118 assert_eq!(sa, vec![4, 1, 6, 2, 8, 3, 0, 5, 7]);
119 }
120
121 #[test]
122 fn fuzz_sa_matches_reference() {
123 let mut state: u64 = 0x1234_5678_9abc_def0;
126 let mut next = || {
127 state ^= state << 13;
128 state ^= state >> 7;
129 state ^= state << 17;
130 state
131 };
132 let alphabets: [&[u8]; 5] = [b"01", b"012", b"ACGT", &[0u8, 1, 255], b"abcdefghij"];
133 for _ in 0..300 {
134 let ab = alphabets[(next() as usize) % alphabets.len()];
135 let n = 1 + (next() as usize) % 5000;
138 let text: Vec<u8> = (0..n).map(|_| ab[(next() as usize) % ab.len()]).collect();
139 assert_eq!(
140 build_suffix_array(&text),
141 sort_based_sa(&text),
142 "SA mismatch for prefix {:?}",
143 &text[..text.len().min(24)]
144 );
145 }
146 assert_eq!(build_suffix_array(&[7u8; 500]), sort_based_sa(&[7u8; 500]));
148 assert_eq!(build_suffix_array(b"a"), sort_based_sa(b"a"));
149 }
150
151 #[test]
152 fn lcp_banana() {
153 let sa = build_suffix_array(b"banana");
154 let lcp = build_lcp(b"banana", &sa);
155 assert_eq!(lcp, vec![0, 1, 3, 0, 0, 2]);
159 }
160}