compcol 0.5.0

//! Burrows–Wheeler transform — forward (encoder) and inverse (decoder).
//!
//! ## Forward
//!
//! We build the suffix array of `input` (sorted indices of all
//! rotations of the string treated as cyclic), then the BWT column L is
//! the byte immediately preceding each suffix in cyclic order:
//!
//! ```text
//! L[i] = input[(sa[i] - 1) mod n]
//! ```
//!
//! The "origin" or "pointer" returned alongside is `pos` such that
//! `sa[pos] == 0` — i.e. the row in the sorted matrix that corresponds
//! to the original (unrotated) string. The decoder needs this to know
//! where the original string sits in the reconstructed matrix.
//!
//! Cyclic-rotation sort is reduced to ordinary suffix sort over the
//! **doubled** input `T + T` followed by a sentinel that is strictly
//! smaller than any byte. Suffix `i` of `T+T+$` for `0 <= i < n` begins
//! with the i-th cyclic rotation of `T` (the first `n` characters
//! match exactly); for non-periodic `T` this means the suffix-sort
//! order over those `n` positions equals the cyclic-rotation order.
//! For periodic `T` two rotations may be equal, in which case any
//! consistent tie-break produces a valid BWT (the L bytes are the same
//! anyway). The sentinel ensures the SA-IS recursion has a strictly
//! smallest character, and any entries `SA[i] >= n` are filtered out
//! (they represent suffixes of the duplicated half).
//!
//! The suffix array itself is built with **SA-IS** (Suffix Array by
//! Induced Sorting; Nong, Zhang & Chan, IEEE TC 2009), a linear-time
//! O(n) algorithm. The classification, induced-sort, and recursion on
//! LMS substrings give us a single-digit-millisecond build on the
//! 900 KB bzip2 block, replacing the previous Θ(n²·log n) naive sort.
//!
//! ## Inverse
//!
//! Classic O(n) algorithm: build the "next-row" permutation from L
//! using a stable rank-by-value, then walk it starting at `origin` for
//! `n` steps, emitting L at each step. See the rust comments in
//! [`bwt_inverse`] for details.

extern crate alloc;
use alloc::vec;
use alloc::vec::Vec;

/// Compute the forward Burrows–Wheeler transform.
///
/// Returns `(L, origin)` where `L` is the BWT output (same length as
/// `input`) and `origin` is the row index of the original (unrotated)
/// string in the sorted rotation matrix.
///
/// Panics: never (no asserts under release builds; debug ones only).
pub(crate) fn bwt_forward(input: &[u8]) -> (Vec<u8>, u32) {
    let n = input.len();
    if n == 0 {
        return (Vec::new(), 0);
    }

    // Build suffix array of the **doubled** input `T + T + sentinel`.
    // Bytes are encoded as values 1..=256 and the sentinel as 0 so the
    // SA-IS routine can use an unsigned alphabet of size 257.
    //
    // n is bounded by ~9·100_000 for bzip2 (well within i32, even
    // doubled).
    let mut text: Vec<i32> = Vec::with_capacity(2 * n + 1);
    for &b in input {
        text.push(b as i32 + 1);
    }
    for &b in input {
        text.push(b as i32 + 1);
    }
    text.push(0); // sentinel

    let sa = sa_is(&text, 257);

    debug_assert_eq!(sa.len(), 2 * n + 1);
    debug_assert_eq!(sa[0] as usize, 2 * n); // sentinel-only suffix sorts first.

    // Filter `sa` to entries `< n` (cyclic rotations of original `T`);
    // the others are suffixes of the duplicated second half and are
    // dominated by the corresponding rotation in the first half.
    let mut l = Vec::with_capacity(n);
    let mut origin: u32 = 0;
    let mut out_i: u32 = 0;
    for &s32 in sa.iter() {
        let s = s32 as usize;
        if s >= n {
            continue;
        }
        let prev = if s == 0 { n - 1 } else { s - 1 };
        l.push(input[prev]);
        if s == 0 {
            origin = out_i;
        }
        out_i += 1;
    }
    debug_assert_eq!(l.len(), n);
    (l, origin)
}

// ─── SA-IS suffix array construction ──────────────────────────────────

/// Build the suffix array of `text` over an integer alphabet whose
/// largest symbol is `< alphabet_size`. The text MUST end with a unique
/// sentinel symbol that is strictly smaller than every other symbol
/// (we use 0; real bytes are shifted into 1..=256). The sentinel
/// guarantees the last suffix is S-type and unique.
///
/// Returns an array of length `text.len()`, where `sa[i]` is the start
/// index of the i-th smallest suffix.
///
/// Linear time, linear extra space. Pure safe Rust.
fn sa_is(text: &[i32], alphabet_size: usize) -> Vec<i32> {
    let n = text.len();
    let mut sa = vec![-1i32; n];
    sa_is_inner(text, &mut sa, alphabet_size);
    sa
}

/// SA-IS core. Writes the suffix array into `sa` (length must equal
/// `text.len()`; entries get overwritten).
fn sa_is_inner(text: &[i32], sa: &mut [i32], alphabet_size: usize) {
    let n = text.len();
    debug_assert_eq!(sa.len(), n);

    if n == 0 {
        return;
    }
    if n == 1 {
        sa[0] = 0;
        return;
    }
    if n == 2 {
        // Two suffixes; just compare.
        if text[0] < text[1] {
            sa[0] = 0;
            sa[1] = 1;
        } else {
            sa[0] = 1;
            sa[1] = 0;
        }
        return;
    }

    // 1. Classify each suffix as S-type or L-type. By convention the
    //    last suffix (just the sentinel) is S-type. A suffix i is
    //    S-type iff text[i] < text[i+1], or text[i] == text[i+1] and
    //    suffix i+1 is S-type. Otherwise L-type.
    //
    //    `t[i] == true` ⇒ S-type.
    let mut t = vec![false; n];
    t[n - 1] = true;
    for i in (0..n - 1).rev() {
        t[i] = if text[i] < text[i + 1] {
            true
        } else if text[i] == text[i + 1] {
            t[i + 1]
        } else {
            false
        };
    }

    // 2. Compute bucket sizes (count of each symbol in `text`).
    let mut counts = vec![0i32; alphabet_size];
    for &c in text {
        counts[c as usize] += 1;
    }

    // 3. Step A: place LMS suffixes at the END of their buckets in `sa`.
    sa.fill(-1);
    let mut ends = bucket_ends(&counts);
    for (i, &c_i) in text.iter().enumerate().take(n).skip(1) {
        if is_lms(&t, i) {
            let c = c_i as usize;
            ends[c] -= 1;
            sa[ends[c] as usize] = i as i32;
        }
    }

    // 4. Induced sort of L-suffixes (left-to-right pass).
    induce_sort_l(text, sa, &t, &counts);

    // 5. Induced sort of S-suffixes (right-to-left pass).
    induce_sort_s(text, sa, &t, &counts);

    // 6. Compact LMS suffixes to the front of SA (preserving the order
    //    we just induced) and name them by their LMS-substring identity.
    let mut n1 = 0usize;
    for i in 0..n {
        if sa[i] >= 0 && is_lms(&t, sa[i] as usize) {
            sa[n1] = sa[i];
            n1 += 1;
        }
    }
    // Clear the rest as a workspace for naming.
    for slot in sa.iter_mut().take(n).skip(n1) {
        *slot = -1;
    }

    // Name LMS substrings.
    let mut name: i32 = 0;
    let mut prev: i32 = -1;
    for i in 0..n1 {
        let pos = sa[i] as usize;
        let mut diff = false;
        if prev == -1 {
            diff = true;
        } else {
            let p = prev as usize;
            // Compare LMS substrings starting at `pos` and `p`.
            let mut d = 0usize;
            loop {
                if pos + d >= n || p + d >= n {
                    diff = true;
                    break;
                }
                if text[pos + d] != text[p + d] || t[pos + d] != t[p + d] {
                    diff = true;
                    break;
                }
                if d > 0 && (is_lms(&t, pos + d) || is_lms(&t, p + d)) {
                    // Reached the next LMS; substrings agreed.
                    if is_lms(&t, pos + d) != is_lms(&t, p + d) {
                        diff = true;
                    }
                    break;
                }
                d += 1;
            }
        }
        if diff {
            name += 1;
            prev = pos as i32;
        }
        // Store name at position pos/2 in the back half.
        sa[n1 + pos / 2] = name - 1;
    }
    // Compact names into the back half (which had -1 in unused slots).
    let mut j = n - 1;
    for i in (n1..n).rev() {
        if sa[i] >= 0 {
            sa[j] = sa[i];
            j -= 1;
        }
    }
    // Now sa[n - n1..n] holds the reduced text of length n1.

    // 7. Solve the reduced problem.
    let new_alpha = (name as usize) + 1;
    let (sa1_area, t1_area) = sa.split_at_mut(n - n1);
    // sa1_area is sa[..n - n1]; t1_area is sa[n - n1..] = reduced text.
    if (name as usize) == n1 {
        // All names distinct ⇒ SA is directly computable from names.
        for (i, &name_of_pos) in t1_area.iter().enumerate() {
            sa1_area[name_of_pos as usize] = i as i32;
        }
    } else {
        // Recurse on the reduced text. We need a slice of length n1 for
        // sa1, and the reduced text occupies the trailing n1 cells.
        let sa1 = &mut sa1_area[..n1];
        // The reduced text is t1_area[..n1] but we want to pass it as
        // an immutable slice. We must copy to avoid aliasing.
        let mut reduced_text: Vec<i32> = Vec::with_capacity(n1);
        reduced_text.extend_from_slice(&t1_area[..n1]);
        sa_is_inner(&reduced_text, sa1, new_alpha);
    }

    // 8. Recover positions of LMS suffixes in the original text.
    //    The trailing region currently holds the reduced text; we need
    //    to translate name-indices in sa1 back to original positions.
    //    We rebuild a list of LMS positions in left-to-right order.
    let mut lms_positions: Vec<i32> = Vec::with_capacity(n1);
    for (i, &is_s) in t.iter().enumerate().take(n).skip(1) {
        if is_s && !t[i - 1] {
            lms_positions.push(i as i32);
        }
    }
    debug_assert_eq!(lms_positions.len(), n1);

    // Translate: sa1[i] = index-of-LMS in original. Reuse trailing area
    // as scratch for translated positions, then place them at bucket
    // ends.
    for slot in sa.iter_mut().take(n1) {
        let idx = *slot as usize; // recursive SA gave us the LMS index in left-to-right order.
        *slot = lms_positions[idx];
    }
    // Clear the rest.
    for slot in sa.iter_mut().take(n).skip(n1) {
        *slot = -1;
    }

    // 9. Place sorted LMS suffixes at the ENDS of their buckets in SA,
    //    in the order produced by the recursive call.
    let mut ends = bucket_ends(&counts);
    // Move them from positions 0..n1 to bucket-end positions, going
    // right-to-left to preserve relative order within each bucket.
    let mut lms_sorted: Vec<i32> = Vec::with_capacity(n1);
    lms_sorted.extend_from_slice(&sa[..n1]);
    for slot in sa.iter_mut().take(n) {
        *slot = -1;
    }
    for &pos in lms_sorted.iter().rev() {
        let c = text[pos as usize] as usize;
        ends[c] -= 1;
        sa[ends[c] as usize] = pos;
    }

    // 10. Final induced sorts: L then S.
    induce_sort_l(text, sa, &t, &counts);
    induce_sort_s(text, sa, &t, &counts);
}

/// `true` iff suffix `i` is S-type AND suffix `i-1` is L-type (left-
/// most S in a run). Suffix 0 is never an LMS in our convention.
fn is_lms(t: &[bool], i: usize) -> bool {
    i > 0 && t[i] && !t[i - 1]
}

/// Compute exclusive prefix sums giving the *start* index of each
/// bucket in SA.
fn bucket_starts(counts: &[i32]) -> Vec<i32> {
    let mut s = Vec::with_capacity(counts.len());
    let mut acc = 0i32;
    for &c in counts {
        s.push(acc);
        acc += c;
    }
    s
}

/// Compute the *end* (one-past-last) index of each bucket in SA.
fn bucket_ends(counts: &[i32]) -> Vec<i32> {
    let mut e = Vec::with_capacity(counts.len());
    let mut acc = 0i32;
    for &c in counts {
        acc += c;
        e.push(acc);
    }
    e
}

/// Induced sort of L-type suffixes. Scans `sa` left-to-right; for each
/// non-negative entry `sa[i] = j`, if `j > 0` and suffix `j-1` is
/// L-type, place `j-1` at the next free slot at the START of bucket
/// `text[j-1]`.
fn induce_sort_l(text: &[i32], sa: &mut [i32], t: &[bool], counts: &[i32]) {
    let n = text.len();
    let mut starts = bucket_starts(counts);
    for i in 0..n {
        if sa[i] <= 0 {
            continue; // -1 or 0 — we handle 0 by not predecessing.
        }
        let j = (sa[i] as usize) - 1;
        if !t[j] {
            // L-type.
            let c = text[j] as usize;
            let slot = starts[c] as usize;
            sa[slot] = j as i32;
            starts[c] += 1;
        }
    }
}

/// Induced sort of S-type suffixes. Scans `sa` right-to-left; for each
/// non-negative entry `sa[i] = j`, if `j > 0` and suffix `j-1` is
/// S-type, place `j-1` at the next free slot at the END of bucket
/// `text[j-1]`.
fn induce_sort_s(text: &[i32], sa: &mut [i32], t: &[bool], counts: &[i32]) {
    let n = text.len();
    let mut ends = bucket_ends(counts);
    for i in (0..n).rev() {
        if sa[i] <= 0 {
            continue;
        }
        let j = (sa[i] as usize) - 1;
        if t[j] {
            // S-type.
            let c = text[j] as usize;
            ends[c] -= 1;
            sa[ends[c] as usize] = j as i32;
        }
    }
}

// ─── inverse BWT (unchanged from previous implementation) ─────────────

/// Compute the inverse BWT given the L column and the origin row.
///
/// The algorithm:
/// 1. Build a frequency table of bytes in L.
/// 2. From it, compute `start[c]` = first position in the sorted L
///    (which is the F column) where byte `c` appears.
/// 3. Build the "next" permutation: for each i (in L order),
///    `next[i] = start[L[i]] + occ_so_far[L[i]]`; `occ_so_far`
///    counts how many `L[i]` values we've seen so far.
/// 4. Walk: `i = origin; for _ in 0..n: emit L[i]; i = next[i]`.
///
/// Equivalently many texts express this as walking `F` and using
/// `prev` — both produce the same byte stream.
pub(crate) fn bwt_inverse(l: &[u8], origin: u32) -> Vec<u8> {
    let n = l.len();
    if n == 0 {
        return Vec::new();
    }
    debug_assert!((origin as usize) < n);

    // Step 1: byte frequency in L.
    let mut count = [0u32; 256];
    for &b in l {
        count[b as usize] += 1;
    }

    // Step 2: `start[c]` = sum of count[0..c] = first position of c in
    // the sorted F column.
    let mut start = [0u32; 256];
    let mut s: u32 = 0;
    for c in 0..256 {
        start[c] = s;
        s += count[c];
    }

    // Step 3: build `next` permutation. We reuse `start` as the
    // running write-cursor: for each i, the next slot for byte L[i] is
    // `start[L[i]]`, which we then increment.
    let mut next = vec![0u32; n];
    let mut cursor = start;
    for (i, &b) in l.iter().enumerate() {
        let c = b as usize;
        next[i] = cursor[c];
        cursor[c] += 1;
    }

    // Step 4: walk. The LF-mapping walk from `origin` yields the
    // original string in REVERSE order: at each step we move "one
    // character backward" in the original. So we walk and place the
    // emitted byte into the buffer from the end backward.
    //
    // (Equivalently we could compute the inverse permutation and walk
    // it forward, but that's more code and constant-factor slower.)
    let mut out = vec![0u8; n];
    let mut i = origin as usize;
    for k in (0..n).rev() {
        out[k] = l[i];
        i = next[i] as usize;
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::vec;

    #[test]
    fn round_trip_short() {
        let input = b"BANANA";
        let (l, origin) = bwt_forward(input);
        let back = bwt_inverse(&l, origin);
        assert_eq!(back, input);
    }

    #[test]
    fn round_trip_single() {
        let input = b"a";
        let (l, origin) = bwt_forward(input);
        assert_eq!(l, b"a");
        assert_eq!(origin, 0);
        let back = bwt_inverse(&l, origin);
        assert_eq!(back, input);
    }

    #[test]
    fn round_trip_empty() {
        let (l, origin) = bwt_forward(&[]);
        assert!(l.is_empty());
        assert_eq!(origin, 0);
        let back = bwt_inverse(&l, origin);
        assert!(back.is_empty());
    }

    #[test]
    fn round_trip_longer() {
        let input = b"the quick brown fox jumps over the lazy dog";
        let (l, origin) = bwt_forward(input);
        let back = bwt_inverse(&l, origin);
        assert_eq!(back, input);
    }

    #[test]
    fn round_trip_repeated_bytes() {
        let input = vec![b'a'; 50];
        let (l, origin) = bwt_forward(&input);
        let back = bwt_inverse(&l, origin);
        assert_eq!(back, input);
    }

    #[test]
    fn round_trip_two_bytes() {
        for (a, b) in [(0u8, 0u8), (0, 255), (255, 0), (1, 2), (2, 1), (5, 5)] {
            let input = [a, b];
            let (l, origin) = bwt_forward(&input);
            let back = bwt_inverse(&l, origin);
            assert_eq!(back, input);
        }
    }

    #[test]
    fn round_trip_three_bytes() {
        for a in 0..3u8 {
            for b in 0..3u8 {
                for c in 0..3u8 {
                    let input = [a, b, c];
                    let (l, origin) = bwt_forward(&input);
                    let back = bwt_inverse(&l, origin);
                    assert_eq!(back, input);
                }
            }
        }
    }

    #[test]
    fn round_trip_with_zero_bytes() {
        // Make sure we handle 0 bytes correctly (sentinel handling).
        let input: Vec<u8> = (0u8..=255).collect();
        let (l, origin) = bwt_forward(&input);
        let back = bwt_inverse(&l, origin);
        assert_eq!(back, input);
    }

    #[test]
    fn round_trip_many_zeros() {
        let input = vec![0u8; 500];
        let (l, origin) = bwt_forward(&input);
        let back = bwt_inverse(&l, origin);
        assert_eq!(back, input);
    }

    #[test]
    fn round_trip_pseudo_random_4k() {
        let mut data = Vec::with_capacity(4096);
        let mut state: u32 = 0xDEAD_BEEF;
        for _ in 0..4096 {
            state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            data.push((state >> 16) as u8);
        }
        let (l, origin) = bwt_forward(&data);
        let back = bwt_inverse(&l, origin);
        assert_eq!(back, data);
    }

    #[test]
    fn matches_naive_on_small_inputs() {
        // Cross-check SA-IS output against a naive cyclic sort for
        // several small inputs.
        fn naive(input: &[u8]) -> (Vec<u8>, u32) {
            let n = input.len();
            if n == 0 {
                return (Vec::new(), 0);
            }
            let mut sa: Vec<usize> = (0..n).collect();
            sa.sort_by(|&a, &b| {
                for k in 0..n {
                    let ai = input[(a + k) % n];
                    let bi = input[(b + k) % n];
                    if ai != bi {
                        return ai.cmp(&bi);
                    }
                }
                core::cmp::Ordering::Equal
            });
            let mut l = Vec::with_capacity(n);
            let mut origin = 0u32;
            for (i, &s) in sa.iter().enumerate() {
                let prev = if s == 0 { n - 1 } else { s - 1 };
                l.push(input[prev]);
                if s == 0 {
                    origin = i as u32;
                }
            }
            (l, origin)
        }

        let cases: &[&[u8]] = &[
            b"",
            b"a",
            b"ab",
            b"ba",
            b"abc",
            b"cba",
            b"banana",
            b"mississippi",
            b"the quick brown fox jumps over the lazy dog",
            b"\0\0\0",
            b"\xff\xff\xff",
            b"\x00\xff\x00\xff\x00",
        ];
        for &case in cases {
            let (sa_l, sa_o) = bwt_forward(case);
            // For inputs containing repeated cyclic-equivalent strings,
            // naive cyclic sort may pick a different origin among tied
            // rows (BWT is not unique then). To compare safely, run a
            // round trip and check we recover the input — that's the
            // contract we need.
            let back = bwt_inverse(&sa_l, sa_o);
            assert_eq!(back.as_slice(), case);
            // Sanity: naive also recovers itself.
            let (nl, no) = naive(case);
            let nback = bwt_inverse(&nl, no);
            assert_eq!(nback.as_slice(), case);
        }
    }
}