Skip to main content

lzcomplexity_core/
suffix_array.rs

1//! Suffix array + LCP construction.
2//!
3//! The C++ library has its own CaPS algorithm (a parallel SA construction).
4//! We build the suffix array in linear time with `cdivsufsort` (a byte-level
5//! divsufsort) and compute LCP with Kasai's algorithm. The suffix array of a
6//! string is unique, so this produces the same order as a comparison sort would
7//! — bit-identical downstream results — but in O(n) instead of O(n² log n).
8
9/// Suffix array + LCP array for a text of length `n`.
10#[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
29/// Build a suffix array for the given text.
30///
31/// The output `sa` does **not** contain a sentinel suffix — its length is
32/// exactly `text.len()`, matching what the C++ `CaPS_SA` produces.
33pub fn build_suffix_array(text: &[u8]) -> Vec<u32> {
34    if text.is_empty() {
35        return Vec::new();
36    }
37    // Both paths produce the (unique) suffix array, so results are identical.
38    // For tiny inputs a plain suffix sort has a much lower constant than
39    // divsufsort's setup, and even its worst case is sub-millisecond at this
40    // size; above the cutoff we switch to the byte-level linear-time divsufsort
41    // so repetitive/large inputs never hit the O(n² log n) sort.
42    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
51/// Comparison suffix sort — low constant, but O(n² log n) on repetitive data,
52/// so it is only used below `SORT_CUTOFF` (and as the fuzz-test reference).
53fn 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
60/// Kasai's algorithm — computes the LCP array from a text and its suffix array.
61///
62/// `lcp[i]` is the longest common prefix length between `sa[i-1]` and `sa[i]`
63/// (with `lcp[0] = 0`). Same convention as `lz::utils::LZ_SuffixArray::LCP`.
64pub 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
94/// Convenience: build both SA and LCP in one go.
95pub 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        // The fast (divsufsort) SA must equal the reference comparison sort for
124        // every input — including non-UTF-8 bytes and degenerate repeats.
125        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            // Span both branches of build_suffix_array: below the sort cutoff
136            // (sort path) and above it (divsufsort path, vs the sort reference).
137            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        // degenerate: all-equal bytes
147        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        // Standard Kasai LCP for "banana" (lcp[0] = 0):
156        // SA = [5,3,1,0,4,2] for suffixes a,ana,anana,banana,na,nana
157        // lcp[0..] = [0, 1, 3, 0, 0, 2]
158        assert_eq!(lcp, vec![0, 1, 3, 0, 0, 2]);
159    }
160}