1 | /* |
2 | This module implements the classical Rabin-Karp substring search algorithm, |
3 | with no extra frills. While its use would seem to break our time complexity |
4 | guarantee of O(m+n) (RK's time complexity is O(mn)), we are careful to only |
5 | ever use RK on a constant subset of haystacks. The main point here is that |
6 | RK has good latency properties for small needles/haystacks. It's very quick |
7 | to compute a needle hash and zip through the haystack when compared to |
8 | initializing Two-Way, for example. And this is especially useful for cases |
9 | where the haystack is just too short for vector instructions to do much good. |
10 | |
11 | The hashing function used here is the same one recommended by ESMAJ. |
12 | |
13 | Another choice instead of Rabin-Karp would be Shift-Or. But its latency |
14 | isn't quite as good since its preprocessing time is a bit more expensive |
15 | (both in practice and in theory). However, perhaps Shift-Or has a place |
16 | somewhere else for short patterns. I think the main problem is that it |
17 | requires space proportional to the alphabet and the needle. If we, for |
18 | example, supported needles up to length 16, then the total table size would be |
19 | len(alphabet)*size_of::<u16>()==512 bytes. Which isn't exactly small, and it's |
20 | probably bad to put that on the stack. So ideally, we'd throw it on the heap, |
21 | but we'd really like to write as much code without using alloc/std as possible. |
22 | But maybe it's worth the special casing. It's a TODO to benchmark. |
23 | |
24 | Wikipedia has a decent explanation, if a bit heavy on the theory: |
25 | https://en.wikipedia.org/wiki/Rabin%E2%80%93Karp_algorithm |
26 | |
27 | But ESMAJ provides something a bit more concrete: |
28 | http://www-igm.univ-mlv.fr/~lecroq/string/node5.html |
29 | |
30 | Finally, aho-corasick uses Rabin-Karp for multiple pattern match in some cases: |
31 | https://github.com/BurntSushi/aho-corasick/blob/3852632f10587db0ff72ef29e88d58bf305a0946/src/packed/rabinkarp.rs |
32 | */ |
33 | |
34 | /// Whether RK is believed to be very fast for the given needle/haystack. |
35 | pub(crate) fn is_fast(haystack: &[u8], _needle: &[u8]) -> bool { |
36 | haystack.len() < 16 |
37 | } |
38 | |
39 | /// Search for the first occurrence of needle in haystack using Rabin-Karp. |
40 | pub(crate) fn find(haystack: &[u8], needle: &[u8]) -> Option<usize> { |
41 | find_with(&NeedleHash::forward(needle), haystack, needle) |
42 | } |
43 | |
44 | /// Search for the first occurrence of needle in haystack using Rabin-Karp with |
45 | /// a pre-computed needle hash. |
46 | pub(crate) fn find_with( |
47 | nhash: &NeedleHash, |
48 | mut haystack: &[u8], |
49 | needle: &[u8], |
50 | ) -> Option<usize> { |
51 | if haystack.len() < needle.len() { |
52 | return None; |
53 | } |
54 | let start: usize = haystack.as_ptr() as usize; |
55 | let mut hash: Hash = Hash::from_bytes_fwd(&haystack[..needle.len()]); |
56 | // N.B. I've experimented with unrolling this loop, but couldn't realize |
57 | // any obvious gains. |
58 | loop { |
59 | if nhash.eq(hash) && is_prefix(haystack, needle) { |
60 | return Some(haystack.as_ptr() as usize - start); |
61 | } |
62 | if needle.len() >= haystack.len() { |
63 | return None; |
64 | } |
65 | hash.roll(&nhash, old:haystack[0], new:haystack[needle.len()]); |
66 | haystack = &haystack[1..]; |
67 | } |
68 | } |
69 | |
70 | /// Search for the last occurrence of needle in haystack using Rabin-Karp. |
71 | pub(crate) fn rfind(haystack: &[u8], needle: &[u8]) -> Option<usize> { |
72 | rfind_with(&NeedleHash::reverse(needle), haystack, needle) |
73 | } |
74 | |
75 | /// Search for the last occurrence of needle in haystack using Rabin-Karp with |
76 | /// a pre-computed needle hash. |
77 | pub(crate) fn rfind_with( |
78 | nhash: &NeedleHash, |
79 | mut haystack: &[u8], |
80 | needle: &[u8], |
81 | ) -> Option<usize> { |
82 | if haystack.len() < needle.len() { |
83 | return None; |
84 | } |
85 | let mut hash: Hash = |
86 | Hash::from_bytes_rev(&haystack[haystack.len() - needle.len()..]); |
87 | loop { |
88 | if nhash.eq(hash) && is_suffix(haystack, needle) { |
89 | return Some(haystack.len() - needle.len()); |
90 | } |
91 | if needle.len() >= haystack.len() { |
92 | return None; |
93 | } |
94 | hash.roll( |
95 | &nhash, |
96 | old:haystack[haystack.len() - 1], |
97 | new:haystack[haystack.len() - needle.len() - 1], |
98 | ); |
99 | haystack = &haystack[..haystack.len() - 1]; |
100 | } |
101 | } |
102 | |
103 | /// A hash derived from a needle. |
104 | #[derive (Clone, Copy, Debug, Default)] |
105 | pub(crate) struct NeedleHash { |
106 | /// The actual hash. |
107 | hash: Hash, |
108 | /// The factor needed to multiply a byte by in order to subtract it from |
109 | /// the hash. It is defined to be 2^(n-1) (using wrapping exponentiation), |
110 | /// where n is the length of the needle. This is how we "remove" a byte |
111 | /// from the hash once the hash window rolls past it. |
112 | hash_2pow: u32, |
113 | } |
114 | |
115 | impl NeedleHash { |
116 | /// Create a new Rabin-Karp hash for the given needle for use in forward |
117 | /// searching. |
118 | pub(crate) fn forward(needle: &[u8]) -> NeedleHash { |
119 | let mut nh = NeedleHash { hash: Hash::new(), hash_2pow: 1 }; |
120 | if needle.is_empty() { |
121 | return nh; |
122 | } |
123 | nh.hash.add(needle[0]); |
124 | for &b in needle.iter().skip(1) { |
125 | nh.hash.add(b); |
126 | nh.hash_2pow = nh.hash_2pow.wrapping_shl(1); |
127 | } |
128 | nh |
129 | } |
130 | |
131 | /// Create a new Rabin-Karp hash for the given needle for use in reverse |
132 | /// searching. |
133 | pub(crate) fn reverse(needle: &[u8]) -> NeedleHash { |
134 | let mut nh = NeedleHash { hash: Hash::new(), hash_2pow: 1 }; |
135 | if needle.is_empty() { |
136 | return nh; |
137 | } |
138 | nh.hash.add(needle[needle.len() - 1]); |
139 | for &b in needle.iter().rev().skip(1) { |
140 | nh.hash.add(b); |
141 | nh.hash_2pow = nh.hash_2pow.wrapping_shl(1); |
142 | } |
143 | nh |
144 | } |
145 | |
146 | /// Return true if the hashes are equivalent. |
147 | fn eq(&self, hash: Hash) -> bool { |
148 | self.hash == hash |
149 | } |
150 | } |
151 | |
152 | /// A Rabin-Karp hash. This might represent the hash of a needle, or the hash |
153 | /// of a rolling window in the haystack. |
154 | #[derive (Clone, Copy, Debug, Default, Eq, PartialEq)] |
155 | pub(crate) struct Hash(u32); |
156 | |
157 | impl Hash { |
158 | /// Create a new hash that represents the empty string. |
159 | pub(crate) fn new() -> Hash { |
160 | Hash(0) |
161 | } |
162 | |
163 | /// Create a new hash from the bytes given for use in forward searches. |
164 | pub(crate) fn from_bytes_fwd(bytes: &[u8]) -> Hash { |
165 | let mut hash = Hash::new(); |
166 | for &b in bytes { |
167 | hash.add(b); |
168 | } |
169 | hash |
170 | } |
171 | |
172 | /// Create a new hash from the bytes given for use in reverse searches. |
173 | fn from_bytes_rev(bytes: &[u8]) -> Hash { |
174 | let mut hash = Hash::new(); |
175 | for &b in bytes.iter().rev() { |
176 | hash.add(b); |
177 | } |
178 | hash |
179 | } |
180 | |
181 | /// Add 'new' and remove 'old' from this hash. The given needle hash should |
182 | /// correspond to the hash computed for the needle being searched for. |
183 | /// |
184 | /// This is meant to be used when the rolling window of the haystack is |
185 | /// advanced. |
186 | fn roll(&mut self, nhash: &NeedleHash, old: u8, new: u8) { |
187 | self.del(nhash, old); |
188 | self.add(new); |
189 | } |
190 | |
191 | /// Add a byte to this hash. |
192 | fn add(&mut self, byte: u8) { |
193 | self.0 = self.0.wrapping_shl(1).wrapping_add(byte as u32); |
194 | } |
195 | |
196 | /// Remove a byte from this hash. The given needle hash should correspond |
197 | /// to the hash computed for the needle being searched for. |
198 | fn del(&mut self, nhash: &NeedleHash, byte: u8) { |
199 | let factor = nhash.hash_2pow; |
200 | self.0 = self.0.wrapping_sub((byte as u32).wrapping_mul(factor)); |
201 | } |
202 | } |
203 | |
204 | /// Returns true if the given needle is a prefix of the given haystack. |
205 | /// |
206 | /// We forcefully don't inline the is_prefix call and hint at the compiler that |
207 | /// it is unlikely to be called. This causes the inner rabinkarp loop above |
208 | /// to be a bit tighter and leads to some performance improvement. See the |
209 | /// memmem/krate/prebuilt/sliceslice-words/words benchmark. |
210 | #[cold ] |
211 | #[inline (never)] |
212 | fn is_prefix(haystack: &[u8], needle: &[u8]) -> bool { |
213 | crate::memmem::util::is_prefix(haystack, needle) |
214 | } |
215 | |
216 | /// Returns true if the given needle is a suffix of the given haystack. |
217 | /// |
218 | /// See is_prefix for why this is forcefully not inlined. |
219 | #[cold ] |
220 | #[inline (never)] |
221 | fn is_suffix(haystack: &[u8], needle: &[u8]) -> bool { |
222 | crate::memmem::util::is_suffix(haystack, needle) |
223 | } |
224 | |
225 | #[cfg (test)] |
226 | mod simpletests { |
227 | define_memmem_simple_tests!(super::find, super::rfind); |
228 | } |
229 | |
230 | #[cfg (all(test, feature = "std" , not(miri)))] |
231 | mod proptests { |
232 | define_memmem_quickcheck_tests!(super::find, super::rfind); |
233 | } |
234 | |