391 lines
14 KiB
Rust
391 lines
14 KiB
Rust
/*!
|
|
An implementation of the [Rabin-Karp substring search algorithm][rabinkarp].
|
|
|
|
Rabin-Karp works by creating a hash of the needle provided and then computing
|
|
a rolling hash for each needle sized window in the haystack. When the rolling
|
|
hash matches the hash of the needle, a byte-wise comparison is done to check
|
|
if a match exists. The worst case time complexity of Rabin-Karp is `O(m *
|
|
n)` where `m ~ len(needle)` and `n ~ len(haystack)`. Its worst case space
|
|
complexity is constant.
|
|
|
|
The main utility of Rabin-Karp is that the searcher can be constructed very
|
|
quickly with very little memory. This makes it especially useful when searching
|
|
for small needles in small haystacks, as it might finish its search before a
|
|
beefier algorithm (like Two-Way) even starts.
|
|
|
|
[rabinkarp]: https://en.wikipedia.org/wiki/Rabin%E2%80%93Karp_algorithm
|
|
*/
|
|
|
|
/*
|
|
(This was the comment I wrote for this module originally when it was not
|
|
exposed. The comment still looks useful, but it's a bit in the weeds, so it's
|
|
not public itself.)
|
|
|
|
This module implements the classical Rabin-Karp substring search algorithm,
|
|
with no extra frills. While its use would seem to break our time complexity
|
|
guarantee of O(m+n) (RK's time complexity is O(mn)), we are careful to only
|
|
ever use RK on a constant subset of haystacks. The main point here is that
|
|
RK has good latency properties for small needles/haystacks. It's very quick
|
|
to compute a needle hash and zip through the haystack when compared to
|
|
initializing Two-Way, for example. And this is especially useful for cases
|
|
where the haystack is just too short for vector instructions to do much good.
|
|
|
|
The hashing function used here is the same one recommended by ESMAJ.
|
|
|
|
Another choice instead of Rabin-Karp would be Shift-Or. But its latency
|
|
isn't quite as good since its preprocessing time is a bit more expensive
|
|
(both in practice and in theory). However, perhaps Shift-Or has a place
|
|
somewhere else for short patterns. I think the main problem is that it
|
|
requires space proportional to the alphabet and the needle. If we, for
|
|
example, supported needles up to length 16, then the total table size would be
|
|
len(alphabet)*size_of::<u16>()==512 bytes. Which isn't exactly small, and it's
|
|
probably bad to put that on the stack. So ideally, we'd throw it on the heap,
|
|
but we'd really like to write as much code without using alloc/std as possible.
|
|
But maybe it's worth the special casing. It's a TODO to benchmark.
|
|
|
|
Wikipedia has a decent explanation, if a bit heavy on the theory:
|
|
https://en.wikipedia.org/wiki/Rabin%E2%80%93Karp_algorithm
|
|
|
|
But ESMAJ provides something a bit more concrete:
|
|
http://www-igm.univ-mlv.fr/~lecroq/string/node5.html
|
|
|
|
Finally, aho-corasick uses Rabin-Karp for multiple pattern match in some cases:
|
|
https://github.com/BurntSushi/aho-corasick/blob/3852632f10587db0ff72ef29e88d58bf305a0946/src/packed/rabinkarp.rs
|
|
*/
|
|
|
|
use crate::ext::Pointer;
|
|
|
|
/// A forward substring searcher using the Rabin-Karp algorithm.
|
|
///
|
|
/// Note that, as a lower level API, a `Finder` does not have access to the
|
|
/// needle it was constructed with. For this reason, executing a search
|
|
/// with a `Finder` requires passing both the needle and the haystack,
|
|
/// where the needle is exactly equivalent to the one given to the `Finder`
|
|
/// at construction time. This design was chosen so that callers can have
|
|
/// more precise control over where and how many times a needle is stored.
|
|
/// For example, in cases where Rabin-Karp is just one of several possible
|
|
/// substring search algorithms.
|
|
#[derive(Clone, Debug)]
|
|
pub struct Finder {
|
|
/// The actual hash.
|
|
hash: Hash,
|
|
/// The factor needed to multiply a byte by in order to subtract it from
|
|
/// the hash. It is defined to be 2^(n-1) (using wrapping exponentiation),
|
|
/// where n is the length of the needle. This is how we "remove" a byte
|
|
/// from the hash once the hash window rolls past it.
|
|
hash_2pow: u32,
|
|
}
|
|
|
|
impl Finder {
|
|
/// Create a new Rabin-Karp forward searcher for the given `needle`.
|
|
///
|
|
/// The needle may be empty. The empty needle matches at every byte offset.
|
|
///
|
|
/// Note that callers must pass the same needle to all search calls using
|
|
/// this `Finder`.
|
|
#[inline]
|
|
pub fn new(needle: &[u8]) -> Finder {
|
|
let mut s = Finder { hash: Hash::new(), hash_2pow: 1 };
|
|
let first_byte = match needle.first() {
|
|
None => return s,
|
|
Some(&first_byte) => first_byte,
|
|
};
|
|
s.hash.add(first_byte);
|
|
for b in needle.iter().copied().skip(1) {
|
|
s.hash.add(b);
|
|
s.hash_2pow = s.hash_2pow.wrapping_shl(1);
|
|
}
|
|
s
|
|
}
|
|
|
|
/// Return the first occurrence of the `needle` in the `haystack`
|
|
/// given. If no such occurrence exists, then `None` is returned.
|
|
///
|
|
/// The `needle` provided must match the needle given to this finder at
|
|
/// construction time.
|
|
///
|
|
/// The maximum value this can return is `haystack.len()`, which can only
|
|
/// occur when the needle and haystack both have length zero. Otherwise,
|
|
/// for non-empty haystacks, the maximum value is `haystack.len() - 1`.
|
|
#[inline]
|
|
pub fn find(&self, haystack: &[u8], needle: &[u8]) -> Option<usize> {
|
|
unsafe {
|
|
let hstart = haystack.as_ptr();
|
|
let hend = hstart.add(haystack.len());
|
|
let nstart = needle.as_ptr();
|
|
let nend = nstart.add(needle.len());
|
|
let found = self.find_raw(hstart, hend, nstart, nend)?;
|
|
Some(found.distance(hstart))
|
|
}
|
|
}
|
|
|
|
/// Like `find`, but accepts and returns raw pointers.
|
|
///
|
|
/// When a match is found, the pointer returned is guaranteed to be
|
|
/// `>= start` and `<= end`. The pointer returned is only ever equivalent
|
|
/// to `end` when both the needle and haystack are empty. (That is, the
|
|
/// empty string matches the empty string.)
|
|
///
|
|
/// This routine is useful if you're already using raw pointers and would
|
|
/// like to avoid converting back to a slice before executing a search.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// Note that `start` and `end` below refer to both pairs of pointers given
|
|
/// to this routine. That is, the conditions apply to both `hstart`/`hend`
|
|
/// and `nstart`/`nend`.
|
|
///
|
|
/// * Both `start` and `end` must be valid for reads.
|
|
/// * Both `start` and `end` must point to an initialized value.
|
|
/// * Both `start` and `end` must point to the same allocated object and
|
|
/// must either be in bounds or at most one byte past the end of the
|
|
/// allocated object.
|
|
/// * Both `start` and `end` must be _derived from_ a pointer to the same
|
|
/// object.
|
|
/// * The distance between `start` and `end` must not overflow `isize`.
|
|
/// * The distance being in bounds must not rely on "wrapping around" the
|
|
/// address space.
|
|
/// * It must be the case that `start <= end`.
|
|
#[inline]
|
|
pub unsafe fn find_raw(
|
|
&self,
|
|
hstart: *const u8,
|
|
hend: *const u8,
|
|
nstart: *const u8,
|
|
nend: *const u8,
|
|
) -> Option<*const u8> {
|
|
let hlen = hend.distance(hstart);
|
|
let nlen = nend.distance(nstart);
|
|
if nlen > hlen {
|
|
return None;
|
|
}
|
|
let mut cur = hstart;
|
|
let end = hend.sub(nlen);
|
|
let mut hash = Hash::forward(cur, cur.add(nlen));
|
|
loop {
|
|
if self.hash == hash && is_equal_raw(cur, nstart, nlen) {
|
|
return Some(cur);
|
|
}
|
|
if cur >= end {
|
|
return None;
|
|
}
|
|
hash.roll(self, cur.read(), cur.add(nlen).read());
|
|
cur = cur.add(1);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A reverse substring searcher using the Rabin-Karp algorithm.
|
|
#[derive(Clone, Debug)]
|
|
pub struct FinderRev(Finder);
|
|
|
|
impl FinderRev {
|
|
/// Create a new Rabin-Karp reverse searcher for the given `needle`.
|
|
#[inline]
|
|
pub fn new(needle: &[u8]) -> FinderRev {
|
|
let mut s = FinderRev(Finder { hash: Hash::new(), hash_2pow: 1 });
|
|
let last_byte = match needle.last() {
|
|
None => return s,
|
|
Some(&last_byte) => last_byte,
|
|
};
|
|
s.0.hash.add(last_byte);
|
|
for b in needle.iter().rev().copied().skip(1) {
|
|
s.0.hash.add(b);
|
|
s.0.hash_2pow = s.0.hash_2pow.wrapping_shl(1);
|
|
}
|
|
s
|
|
}
|
|
|
|
/// Return the last occurrence of the `needle` in the `haystack`
|
|
/// given. If no such occurrence exists, then `None` is returned.
|
|
///
|
|
/// The `needle` provided must match the needle given to this finder at
|
|
/// construction time.
|
|
///
|
|
/// The maximum value this can return is `haystack.len()`, which can only
|
|
/// occur when the needle and haystack both have length zero. Otherwise,
|
|
/// for non-empty haystacks, the maximum value is `haystack.len() - 1`.
|
|
#[inline]
|
|
pub fn rfind(&self, haystack: &[u8], needle: &[u8]) -> Option<usize> {
|
|
unsafe {
|
|
let hstart = haystack.as_ptr();
|
|
let hend = hstart.add(haystack.len());
|
|
let nstart = needle.as_ptr();
|
|
let nend = nstart.add(needle.len());
|
|
let found = self.rfind_raw(hstart, hend, nstart, nend)?;
|
|
Some(found.distance(hstart))
|
|
}
|
|
}
|
|
|
|
/// Like `rfind`, but accepts and returns raw pointers.
|
|
///
|
|
/// When a match is found, the pointer returned is guaranteed to be
|
|
/// `>= start` and `<= end`. The pointer returned is only ever equivalent
|
|
/// to `end` when both the needle and haystack are empty. (That is, the
|
|
/// empty string matches the empty string.)
|
|
///
|
|
/// This routine is useful if you're already using raw pointers and would
|
|
/// like to avoid converting back to a slice before executing a search.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// Note that `start` and `end` below refer to both pairs of pointers given
|
|
/// to this routine. That is, the conditions apply to both `hstart`/`hend`
|
|
/// and `nstart`/`nend`.
|
|
///
|
|
/// * Both `start` and `end` must be valid for reads.
|
|
/// * Both `start` and `end` must point to an initialized value.
|
|
/// * Both `start` and `end` must point to the same allocated object and
|
|
/// must either be in bounds or at most one byte past the end of the
|
|
/// allocated object.
|
|
/// * Both `start` and `end` must be _derived from_ a pointer to the same
|
|
/// object.
|
|
/// * The distance between `start` and `end` must not overflow `isize`.
|
|
/// * The distance being in bounds must not rely on "wrapping around" the
|
|
/// address space.
|
|
/// * It must be the case that `start <= end`.
|
|
#[inline]
|
|
pub unsafe fn rfind_raw(
|
|
&self,
|
|
hstart: *const u8,
|
|
hend: *const u8,
|
|
nstart: *const u8,
|
|
nend: *const u8,
|
|
) -> Option<*const u8> {
|
|
let hlen = hend.distance(hstart);
|
|
let nlen = nend.distance(nstart);
|
|
if nlen > hlen {
|
|
return None;
|
|
}
|
|
let mut cur = hend.sub(nlen);
|
|
let start = hstart;
|
|
let mut hash = Hash::reverse(cur, cur.add(nlen));
|
|
loop {
|
|
if self.0.hash == hash && is_equal_raw(cur, nstart, nlen) {
|
|
return Some(cur);
|
|
}
|
|
if cur <= start {
|
|
return None;
|
|
}
|
|
cur = cur.sub(1);
|
|
hash.roll(&self.0, cur.add(nlen).read(), cur.read());
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Whether RK is believed to be very fast for the given needle/haystack.
|
|
#[inline]
|
|
pub(crate) fn is_fast(haystack: &[u8], _needle: &[u8]) -> bool {
|
|
haystack.len() < 16
|
|
}
|
|
|
|
/// A Rabin-Karp hash. This might represent the hash of a needle, or the hash
|
|
/// of a rolling window in the haystack.
|
|
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
|
|
struct Hash(u32);
|
|
|
|
impl Hash {
|
|
/// Create a new hash that represents the empty string.
|
|
#[inline(always)]
|
|
fn new() -> Hash {
|
|
Hash(0)
|
|
}
|
|
|
|
/// Create a new hash from the bytes given for use in forward searches.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// The given pointers must be valid to read from within their range.
|
|
#[inline(always)]
|
|
unsafe fn forward(mut start: *const u8, end: *const u8) -> Hash {
|
|
let mut hash = Hash::new();
|
|
while start < end {
|
|
hash.add(start.read());
|
|
start = start.add(1);
|
|
}
|
|
hash
|
|
}
|
|
|
|
/// Create a new hash from the bytes given for use in reverse searches.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// The given pointers must be valid to read from within their range.
|
|
#[inline(always)]
|
|
unsafe fn reverse(start: *const u8, mut end: *const u8) -> Hash {
|
|
let mut hash = Hash::new();
|
|
while start < end {
|
|
end = end.sub(1);
|
|
hash.add(end.read());
|
|
}
|
|
hash
|
|
}
|
|
|
|
/// Add 'new' and remove 'old' from this hash. The given needle hash should
|
|
/// correspond to the hash computed for the needle being searched for.
|
|
///
|
|
/// This is meant to be used when the rolling window of the haystack is
|
|
/// advanced.
|
|
#[inline(always)]
|
|
fn roll(&mut self, finder: &Finder, old: u8, new: u8) {
|
|
self.del(finder, old);
|
|
self.add(new);
|
|
}
|
|
|
|
/// Add a byte to this hash.
|
|
#[inline(always)]
|
|
fn add(&mut self, byte: u8) {
|
|
self.0 = self.0.wrapping_shl(1).wrapping_add(u32::from(byte));
|
|
}
|
|
|
|
/// Remove a byte from this hash. The given needle hash should correspond
|
|
/// to the hash computed for the needle being searched for.
|
|
#[inline(always)]
|
|
fn del(&mut self, finder: &Finder, byte: u8) {
|
|
let factor = finder.hash_2pow;
|
|
self.0 = self.0.wrapping_sub(u32::from(byte).wrapping_mul(factor));
|
|
}
|
|
}
|
|
|
|
/// Returns true when `x[i] == y[i]` for all `0 <= i < n`.
|
|
///
|
|
/// We forcefully don't inline this to hint at the compiler that it is unlikely
|
|
/// to be called. This causes the inner rabinkarp loop above to be a bit
|
|
/// tighter and leads to some performance improvement. See the
|
|
/// memmem/krate/prebuilt/sliceslice-words/words benchmark.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// Same as `crate::arch::all::is_equal_raw`.
|
|
#[cold]
|
|
#[inline(never)]
|
|
unsafe fn is_equal_raw(x: *const u8, y: *const u8, n: usize) -> bool {
|
|
crate::arch::all::is_equal_raw(x, y, n)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
define_substring_forward_quickcheck!(|h, n| Some(
|
|
Finder::new(n).find(h, n)
|
|
));
|
|
define_substring_reverse_quickcheck!(|h, n| Some(
|
|
FinderRev::new(n).rfind(h, n)
|
|
));
|
|
|
|
#[test]
|
|
fn forward() {
|
|
crate::tests::substring::Runner::new()
|
|
.fwd(|h, n| Some(Finder::new(n).find(h, n)))
|
|
.run();
|
|
}
|
|
|
|
#[test]
|
|
fn reverse() {
|
|
crate::tests::substring::Runner::new()
|
|
.rev(|h, n| Some(FinderRev::new(n).rfind(h, n)))
|
|
.run();
|
|
}
|
|
}
|