Knuth-Morris-Pratt implementation

kmp.rs

//! Knuth-Morris-Pratt implementation
//!
//! See Chapter 17 in Blandy + Orendorff

/// computes the KMP longest prefix-suffix function
pub fn kmp<T: AsRef<[u8]>>(pattern: T) -> Vec<usize> {
    let pattern = pattern.as_ref();
    let n = pattern.len();
    let mut ans = vec![0; n + 1];

    for (i, c) in pattern.iter().enumerate().skip(1) {
        let mut j = ans[i];

        loop {
            if c == &pattern[j] {
                ans[i + 1] = j + 1;
                break;
            }

            if j == 0 {
                break;
            }

            j = ans[j];
        }
    }

    ans
}

/// find pattern in haystack
pub fn find<T: AsRef<[u8]>>(pattern: T, haystack: T) -> Vec<usize> {
    let pattern = pattern.as_ref();
    let haystack = haystack.as_ref();

    let mut ans = Vec::new();

    let f = kmp(pattern);

    let mut i = 0;
    let mut j = 0;
    while i < haystack.len() {
        if haystack[i] == pattern[j] {
            i += 1;
            j += 1;

            if j == pattern.len() {
                ans.push(i - pattern.len());
                j = f[j];
            }
        } else if j > 0 {
            j = f[j];
        } else {
            i += 1;
        }
    }

    ans
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_kmp_fn() {
        let p = super::kmp("ABABAC");
        assert_eq!(p, vec![0, 0, 0, 1, 2, 3, 0]);
    }

    #[test]
    fn test_find() {
        assert_eq!(
            super::find("ABABA", "ABABABABABABABABABABA"),
            vec![0, 2, 4, 6, 8, 10, 12, 14, 16]
        );
    }
}