1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
/*
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
*/

/// Whether RK is believed to be very fast for the given needle/haystack.
pub(crate) fn is_fast(haystack: &[u8], _needle: &[u8]) -> bool {
    haystack.len() < 16
}

/// Search for the first occurrence of needle in haystack using Rabin-Karp.
pub(crate) fn find(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    find_with(&NeedleHash::forward(needle), haystack, needle)
}

/// Search for the first occurrence of needle in haystack using Rabin-Karp with
/// a pre-computed needle hash.
pub(crate) fn find_with(
    nhash: &NeedleHash,
    mut haystack: &[u8],
    needle: &[u8],
) -> Option<usize> {
    if haystack.len() < needle.len() {
        return None;
    }
    let start = haystack.as_ptr() as usize;
    let mut hash = Hash::from_bytes_fwd(&haystack[..needle.len()]);
    // N.B. I've experimented with unrolling this loop, but couldn't realize
    // any obvious gains.
    loop {
        if nhash.eq(hash) && is_prefix(haystack, needle) {
            return Some(haystack.as_ptr() as usize - start);
        }
        if needle.len() >= haystack.len() {
            return None;
        }
        hash.roll(&nhash, haystack[0], haystack[needle.len()]);
        haystack = &haystack[1..];
    }
}

/// Search for the last occurrence of needle in haystack using Rabin-Karp.
pub(crate) fn rfind(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    rfind_with(&NeedleHash::reverse(needle), haystack, needle)
}

/// Search for the last occurrence of needle in haystack using Rabin-Karp with
/// a pre-computed needle hash.
pub(crate) fn rfind_with(
    nhash: &NeedleHash,
    mut haystack: &[u8],
    needle: &[u8],
) -> Option<usize> {
    if haystack.len() < needle.len() {
        return None;
    }
    let mut hash =
        Hash::from_bytes_rev(&haystack[haystack.len() - needle.len()..]);
    loop {
        if nhash.eq(hash) && is_suffix(haystack, needle) {
            return Some(haystack.len() - needle.len());
        }
        if needle.len() >= haystack.len() {
            return None;
        }
        hash.roll(
            &nhash,
            haystack[haystack.len() - 1],
            haystack[haystack.len() - needle.len() - 1],
        );
        haystack = &haystack[..haystack.len() - 1];
    }
}

/// A hash derived from a needle.
#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct NeedleHash {
    /// 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 NeedleHash {
    /// Create a new Rabin-Karp hash for the given needle for use in forward
    /// searching.
    pub(crate) fn forward(needle: &[u8]) -> NeedleHash {
        let mut nh = NeedleHash { hash: Hash::new(), hash_2pow: 1 };
        if needle.is_empty() {
            return nh;
        }
        nh.hash.add(needle[0]);
        for &b in needle.iter().skip(1) {
            nh.hash.add(b);
            nh.hash_2pow = nh.hash_2pow.wrapping_shl(1);
        }
        nh
    }

    /// Create a new Rabin-Karp hash for the given needle for use in reverse
    /// searching.
    pub(crate) fn reverse(needle: &[u8]) -> NeedleHash {
        let mut nh = NeedleHash { hash: Hash::new(), hash_2pow: 1 };
        if needle.is_empty() {
            return nh;
        }
        nh.hash.add(needle[needle.len() - 1]);
        for &b in needle.iter().rev().skip(1) {
            nh.hash.add(b);
            nh.hash_2pow = nh.hash_2pow.wrapping_shl(1);
        }
        nh
    }

    /// Return true if the hashes are equivalent.
    fn eq(&self, hash: Hash) -> bool {
        self.hash == hash
    }
}

/// 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)]
pub(crate) struct Hash(u32);

impl Hash {
    /// Create a new hash that represents the empty string.
    pub(crate) fn new() -> Hash {
        Hash(0)
    }

    /// Create a new hash from the bytes given for use in forward searches.
    pub(crate) fn from_bytes_fwd(bytes: &[u8]) -> Hash {
        let mut hash = Hash::new();
        for &b in bytes {
            hash.add(b);
        }
        hash
    }

    /// Create a new hash from the bytes given for use in reverse searches.
    fn from_bytes_rev(bytes: &[u8]) -> Hash {
        let mut hash = Hash::new();
        for &b in bytes.iter().rev() {
            hash.add(b);
        }
        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.
    fn roll(&mut self, nhash: &NeedleHash, old: u8, new: u8) {
        self.del(nhash, old);
        self.add(new);
    }

    /// Add a byte to this hash.
    fn add(&mut self, byte: u8) {
        self.0 = self.0.wrapping_shl(1).wrapping_add(byte as u32);
    }

    /// Remove a byte from this hash. The given needle hash should correspond
    /// to the hash computed for the needle being searched for.
    fn del(&mut self, nhash: &NeedleHash, byte: u8) {
        let factor = nhash.hash_2pow;
        self.0 = self.0.wrapping_sub((byte as u32).wrapping_mul(factor));
    }
}

/// Returns true if the given needle is a prefix of the given haystack.
///
/// We forcefully don't inline the is_prefix call and 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.
#[cold]
#[inline(never)]
fn is_prefix(haystack: &[u8], needle: &[u8]) -> bool {
    crate::memmem::util::is_prefix(haystack, needle)
}

/// Returns true if the given needle is a suffix of the given haystack.
///
/// See is_prefix for why this is forcefully not inlined.
#[cold]
#[inline(never)]
fn is_suffix(haystack: &[u8], needle: &[u8]) -> bool {
    crate::memmem::util::is_suffix(haystack, needle)
}

#[cfg(test)]
mod simpletests {
    define_memmem_simple_tests!(super::find, super::rfind);
}

#[cfg(all(test, feature = "std", not(miri)))]
mod proptests {
    define_memmem_quickcheck_tests!(super::find, super::rfind);
}