0 | ||| Fast Knuth-Morris-Pratt search of ByteStrings
  1 | module Data.ByteString.Search.KnuthMorrisPratt
  2 |
  3 | import Data.ByteString.Search.Internal.Utils
  4 |
  5 | import Data.Array.Core
  6 | import Data.Array.Mutable
  7 | import Data.Bits
  8 | import Data.ByteString
  9 | import Data.Linear.Ref1
 10 | import Data.So
 11 |
 12 | %hide Data.Buffer.Core.get
 13 | %hide Data.Buffer.Core.set
 14 |
 15 | %default total
 16 |
 17 | ||| Returns a list of starting positions of a pattern `ByteString`
 18 | ||| (0-based) across the list of target `ByteString`s.
 19 | private
 20 | matcher :  Bool
 21 |         -> ByteString
 22 |         -> List ByteString
 23 |         -> F1 s (List Nat)
 24 | matcher overlap pat chunks t =
 25 |   let bords     # t := kmpBorders pat t
 26 |       searcher' # t := searcher Z Z pat chunks Lin bords overlap t
 27 |     in (searcher' <>> []) # t
 28 |   where
 29 |     mutual
 30 |       findMatch :  (prior : Nat)
 31 |                 -> (pati : Nat)
 32 |                 -> (stri : Nat)
 33 |                 -> (pat : ByteString)
 34 |                 -> (strs : List ByteString)
 35 |                 -> (final : SnocList Nat)
 36 |                 -> (bords : MArray s (S (length pat)) Nat)
 37 |                 -> (overlap : Bool)
 38 |                 -> F1 s (SnocList Nat)
 39 |       findMatch _     _    _    _   []               final _     _       t =
 40 |         final # t
 41 |       findMatch prior pati stri pat strs@(str::rest) final bords overlap t =
 42 |         let patlen := length pat
 43 |           in case pati == patlen of
 44 |                True  =>
 45 |                  case overlap of
 46 |                    True  =>
 47 |                      case tryNatToFin patlen of
 48 |                        Nothing      =>
 49 |                          (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.matcher.findMatch: can't index into ByteString") # t
 50 |                        Just patlen' =>
 51 |                          let final'  := minus (plus prior stri) patlen
 52 |                              ami # t := get bords patlen' t
 53 |                            in case ami == Z of
 54 |                                 True  =>
 55 |                                   let final'' := final :< final'
 56 |                                     in assert_total (checkHead prior stri pat strs final'' bords overlap t)
 57 |                                 False =>
 58 |                                   let final'' := final :< final'
 59 |                                     in assert_total (findMatch prior ami stri pat strs final'' bords overlap t)
 60 |                    False =>
 61 |                      let final'  := minus (plus prior stri) patlen
 62 |                          final'' := final :< final'
 63 |                        in assert_total (checkHead prior stri pat strs final'' bords overlap t)
 64 |                False =>
 65 |                  let strlen := length str
 66 |                    in case stri == strlen of
 67 |                         True  =>
 68 |                           assert_total (searcher (plus prior strlen) pati pat rest final bords overlap t)
 69 |                         False =>
 70 |                           let pati' := index pati pat
 71 |                             in case pati' of
 72 |                                  Nothing     =>
 73 |                                    (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.matcher.findMatch: can't index into ByteString") # t
 74 |                                  Just pati'' =>
 75 |                                    let stri' := index stri str
 76 |                                      in case stri' of
 77 |                                           Nothing     =>
 78 |                                             (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.matcher.findMatch: can't index into ByteString") # t
 79 |                                           Just stri'' =>
 80 |                                             case stri'' == pati'' of
 81 |                                               True  =>
 82 |                                                 assert_total (findMatch prior (S pati) (S stri) pat strs final bords overlap t)
 83 |                                               False =>
 84 |                                                 case tryNatToFin pati of
 85 |                                                   Nothing      =>
 86 |                                                     (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.matcher.findMatch: can't convert Nat to Fin") # t        
 87 |                                                   Just pati''' =>
 88 |                                                     let pati'''' # t := get bords pati''' t
 89 |                                                       in case pati'''' == Z of
 90 |                                                            True  =>
 91 |                                                              assert_total (checkHead prior (S stri) pat strs final bords overlap t)
 92 |                                                            False =>
 93 |                                                              assert_total (findMatch prior pati'''' stri pat strs final bords overlap t)
 94 |       checkHead :  (prior : Nat)
 95 |                 -> (stri : Nat)
 96 |                 -> (pat : ByteString)
 97 |                 -> (strs : List ByteString)
 98 |                 -> (final : SnocList Nat)
 99 |                 -> (bords : MArray s (S (length pat)) Nat)
100 |                 -> (overlap : Bool)
101 |                 -> F1 s (SnocList Nat)
102 |       checkHead _     _    _   []               final _     _       t =
103 |         final # t
104 |       checkHead prior stri pat strs@(str::rest) final bords overlap t =
105 |         let strlen := length str
106 |           in case stri == strlen of
107 |                True  =>
108 |                  assert_total (searcher (plus prior strlen) Z pat rest final bords overlap t)
109 |                False =>
110 |                  let stri' := index stri str
111 |                    in case stri' of
112 |                         Nothing     =>
113 |                           (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.matcher.checkHead: can't index into ByteString") # t
114 |                         Just stri'' =>
115 |                           let patzero := index Z pat
116 |                             in case patzero of
117 |                                  Nothing       =>
118 |                                    (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.matcher.checkHead: can't index into ByteString") # t
119 |                                  Just patzero' =>
120 |                                    case stri'' == patzero' of
121 |                                      True  =>
122 |                                        assert_total (findMatch prior (S Z) (S stri) pat strs final bords overlap t)
123 |                                      False =>
124 |                                        assert_total (checkHead prior (S stri) pat strs final bords overlap t)
125 |       searcher :  (prior : Nat)
126 |                -> (patpos : Nat)
127 |                -> (pat : ByteString)
128 |                -> (strs : List ByteString)
129 |                -> (final : SnocList Nat)
130 |                -> (bords : MArray s (S (length pat)) Nat)
131 |                -> (overlap : Bool)
132 |                -> F1 s (SnocList Nat)
133 |       searcher _     _      _   []   final _     _       t =
134 |         final # t
135 |       searcher prior Z      pat strs final bords overlap t =
136 |         assert_total (checkHead prior Z pat strs final bords overlap t)
137 |       searcher prior patpos pat strs final bords overlap t =
138 |         assert_total (findMatch prior patpos Z pat strs final bords overlap t)
139 |
140 | ||| Performs a Knuth–Morris–Pratt string search on a `ByteString`.
141 | |||
142 | ||| This function finds all (0-based) starting indices of the non-empty pattern `ByteString`
143 | ||| pat within the non-empty target `ByteString`, using the KMP border table
144 | ||| computed by `kmpBorders`.
145 | |||
146 | ||| Example:
147 | |||
148 | ||| | pat  | target     |
149 | ||| | ---- | ---------- |
150 | ||| | "AN" | "ANPANMAN" |
151 | |||
152 | ||| | Start | Substring      | Match? | Explanation                                      |
153 | ||| | ----- | -------------- | ------ | ------------------------------------------------ |
154 | ||| | 0     | **"AN"**PANMAN | Yes    | Full pattern `"AN"` matches starting at index 0. |
155 | ||| | 1     | A**"NP"**ANMAN | No     | Mismatch after the first character.              |
156 | ||| | 2     | AN**"PA"**NMAN | No     | No match — next candidate after suffix shift.    |
157 | ||| | 3     | ANP**"AN"**MAN | Yes    | Match found at index 3.                          |
158 | ||| | 4     | ANPA**"NM"**AN | No     | Mismatch.                                        |
159 | ||| | 5     | ANPAN**"MA"**N | No     | Mismatch.                                        |
160 | ||| | 6     | ANPANM**"AN"** | Yes    | Final match found at index 6.                    |
161 | ||| 
162 | |||
163 | ||| matchKMP "AN" "ANPANMAN" => [0, 3, 6]
164 | |||
165 | export
166 | matchKMP :  (pat : ByteString)
167 |          -> (target : ByteString)
168 |          -> {0 prfpat : So (not $ null pat)}
169 |          -> {0 prftarget : So (not $ null target)}
170 |          -> F1 s (List Nat)
171 | matchKMP pat target {prfpat} {prftarget} t =
172 |   matcher False pat [target] t
173 |
174 | ||| Performs a Knuth–Morris–Pratt string search on a `ByteString`.
175 | |||
176 | ||| This function finds all (0-based) indices (possibly overlapping)
177 | ||| of the non-empty pattern `ByteString` pat
178 | ||| within the non-empty target `ByteString`, using the KMP border table
179 | ||| computed by `kmpBorders`.
180 | |||
181 | ||| Example:
182 | |||
183 | ||| | pat   | target      |
184 | ||| | ----- | ----------- |
185 | ||| | "ABC" | "ABCABCABC" |
186 | |||
187 | ||| | Start | Substring       | Match? | Explanation                                                      |
188 | ||| | ----- | --------------- | ------ | ---------------------------------------------------------------- |
189 | ||| | 0     | **"ABCABC"**ABC | Yes    | Full pattern matches starting at index 0.                        |
190 | ||| | 1     | A**"BCABCA"**BC | No     | Mismatch starts immediately after first letter.                  |
191 | ||| | 2     | AB**"CABCAA"**C | No     | Shift by suffix table → mismatch on 2nd char.                    |
192 | ||| | 3     | ABC**"ABC"**    | Yes    | Overlapping match starting at index 3 (because `"ABC"` repeats). |
193 | ||| 
194 | ||| indicesKMP "ABCABC" "ABCABCABC" => [0, 3]
195 | |||
196 | export
197 | indicesKMP :  (pat : ByteString)
198 |            -> (target : ByteString)
199 |            -> {0 prfpat : So (not $ null pat)}
200 |            -> {0 prftarget : So (not $ null target)}
201 |            -> F1 s (List Nat)
202 | indicesKMP pat target {prfpat} {prftarget} t =
203 |   matcher True pat [target] t
204 |
205 | ||| Splits a ByteString at the first match of pat in target.
206 | |||
207 | ||| This function uses the Knuth-Morris-Pratt matcher (with overlap = False) to
208 | ||| locate the earliest occurrence of pat in target.  If the pattern is
209 | ||| found at index i, the pattern ByteString pat is split at that index,
210 | ||| returning the prefix and suffix as a pair (before, after).
211 | |||
212 | ||| If the pattern does not occur in the target, (pat, empty) is returned.
213 | ||| In other words, the entire pattern becomes the “before” part and the
214 | ||| “after” part is an empty ByteString.
215 | |||
216 | export
217 | breakKMP :  (pat : ByteString)
218 |          -> (target : ByteString)
219 |          -> {0 prfpat : So (not $ null pat)}
220 |          -> {0 prftarget : So (not $ null target)}
221 |          -> {0 prflength : So ((length target) >= (length pat))}
222 |          -> F1 s (ByteString, ByteString)
223 | breakKMP pat target {prfpat} {prftarget} {prflength} t =
224 |    let matcher' # t := matcher False pat [target] t
225 |      in case matcher' of
226 |           []       =>
227 |             (target, empty) # t
228 |           (i :: _) =>
229 |             let target' := splitAt (cast {to=Nat} i) target
230 |               in case target' of
231 |                    Nothing       =>
232 |                      (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.breakKMP: can't split ByteString") # t
233 |                    Just target'' =>
234 |                      target'' # t
235 |
236 | ||| Splits a ByteString after the first match of pat in target.
237 | |||
238 | ||| This function uses the Knuth-Morris-Pratt matcher (with overlap = False) to
239 | ||| find the earliest occurrence of pat in target.  If the pattern is
240 | ||| found at index i, this function splits pat at position i + length pat,
241 | ||| producing a pair (before, after) that places the entire matched region
242 | ||| into the prefix.
243 | |||
244 | ||| If the pattern does not occur in target, the function returns
245 | ||| (pat, empty), the entire pattern is the “before” substring, and the
246 | ||| suffix is empty.
247 | |||
248 | export
249 | breakAfterKMP : (pat : ByteString)
250 |               -> (target : ByteString)
251 |               -> {0 prfpat : So (not $ null pat)}
252 |               -> {0 prftarget : So (not $ null target)}
253 |               -> {0 prflength : So ((length target) >= (length pat))}
254 |               -> F1 s (ByteString, ByteString)
255 | breakAfterKMP pat target {prfpat} {prftarget} {prflength} t =
256 |    let matcher' # t := matcher False pat [target] t
257 |      in case matcher' of
258 |           []       =>
259 |             (target, empty) # t
260 |           (i :: _) =>
261 |             let target' := splitAt (plus (cast {to=Nat} i) (length pat)) target
262 |               in case target' of
263 |                    Nothing    =>
264 |                      (assert_total $ idris_crash "Data.ByteString.Search.KnuthMorrisPratt.breakAfterKMP: can't split ByteString") # t
265 |                    Just target'' =>
266 |                      target'' # t
267 |
268 | ||| Splits a ByteString into a list of pieces according to repeated
269 | ||| matches of target, keeping the matching prefix of pat
270 | ||| at the front of each produced chunk.
271 | |||
272 | ||| This function repeatedly searches target for occurrences of pat
273 | ||| (using the Knuth-Morris-Pratt matcher with overlap = False).  Each time a
274 | ||| match is found at index i, the prefix of pat up to i + length pat
275 | ||| is emitted as the next chunk, and the function continues processing the
276 | ||| remaining suffix of pat.
277 | |||
278 | ||| Unlike breakKMP or breakAfterKMP, this function performs repeated
279 | ||| splitting until the entire pattern has been consumed, producing a
280 | ||| list of ByteStrings.
281 | |||
282 | export
283 | splitKeepFrontKMP :  (pat : ByteString)
284 |                   -> (target : ByteString)
285 |                   -> {0 prfpat : So (not $ null pat)}
286 |                   -> {0 prftarget : So (not $ null target)}
287 |                   -> {0 prflength : So ((length target) >= (length pat))}
288 |                   -> F1 s (List ByteString)
289 | splitKeepFrontKMP pat target {prfpat} {prftarget} {prflength} t =
290 |   let splitter' # t := splitter pat target Lin t
291 |     in (splitter' <>> []) # t
292 |   where
293 |     psSplitter :  (pat : ByteString)
294 |                -> (target : ByteString)
295 |                -> (final : SnocList ByteString)
296 |                -> F1 s (SnocList ByteString)
297 |     psSplitter pat target final t =
298 |       let matcher' # t := matcher False pat [(drop (length pat) target)] t
299 |         in case matcher' of
300 |              []       =>
301 |                let final' := final :< target
302 |                  in final' # t
303 |              (i :: _) =>
304 |                let length' := plus (cast {to=Nat} i) (length pat)
305 |                    final'  := final :< (take length' target)
306 |                  in assert_total (psSplitter pat (drop length' target) final' t)
307 |     splitter :  (pat : ByteString)
308 |              -> (target : ByteString)
309 |              -> (final : SnocList ByteString)
310 |              -> F1 s (SnocList ByteString)
311 |     splitter pat target final t =
312 |       let matcher' # t := matcher False pat [target] t
313 |         in case matcher' of
314 |              []       =>
315 |                let final' := final :< target
316 |                  in final' # t
317 |              (i :: _) =>
318 |                case i == 0 of
319 |                  True  =>
320 |                    assert_total (psSplitter pat target final t)
321 |                  False =>
322 |                    let final' := final :< (take (cast {to=Nat} i) target)
323 |                      in assert_total (psSplitter pat (drop (cast {to=Nat} i) target) final' t)
324 |
325 | ||| Splits a ByteString into a list of pieces according to repeated
326 | ||| matches of pat inside target, keeping the matching
327 | ||| suffix of pat at the end of each produced chunk.
328 | |||
329 | ||| This function repeatedly searches target for occurrences of pat
330 | ||| (using the Knuth-Morris-Pratt matcher with overlap = False).  Each time a
331 | ||| match is found at index i, the next chunk emitted is the prefix of
332 | ||| target of length i + length pat, which includes the entire matched
333 | ||| occurrence of pat at its end.
334 | |||
335 | ||| After emitting this chunk, the function continues splitting the
336 | ||| remainder of target until all input has been consumed.
337 | |||
338 | ||| Unlike splitKeepFrontKMP, which keeps the matched prefix of pat
339 | ||| at the front of each chunk, splitKeepEndKMP ensures the match
340 | ||| appears at the end of each chunk.
341 | |||
342 | ||| If pat does not occur in target, the result is a singleton list
343 | ||| containing the original target.
344 | |||
345 | export
346 | splitKeepEndKMP :  (pat : ByteString)
347 |                 -> (target : ByteString)
348 |                 -> {0 prfpat : So (not $ null pat)}
349 |                 -> {0 prftarget : So (not $ null target)}
350 |                 -> {0 prflength : So ((length target) >= (length pat))}
351 |                 -> F1 s (List ByteString)
352 | splitKeepEndKMP pat target {prfpat} {prftarget} {prflength} t =
353 |   let splitter' # t := splitter pat target Lin t
354 |     in (splitter' <>> []) # t
355 |   where
356 |     splitter :  (pat : ByteString)
357 |              -> (target : ByteString)
358 |              -> (final : SnocList ByteString)
359 |              -> F1 s (SnocList ByteString)
360 |     splitter pat target final t =
361 |       let matcher' # t := matcher False pat [target] t
362 |         in case matcher' of
363 |              []       =>
364 |                let final' := final :< target
365 |                  in final' # t
366 |              (i :: _) =>
367 |                let length' := plus (cast {to=Nat} i) (length pat)
368 |                    final'  := final :< (take length' target)
369 |                  in assert_total (splitter pat (drop length' target) final' t)
370 |
371 | ||| Splits a ByteString into a list of pieces according to repeated
372 | ||| matches of pat inside target, dropping each matched
373 | ||| occurrence from the output entirely.
374 | |||
375 | ||| This function repeatedly searches target for occurrences of pat
376 | ||| (using the Knuth-Morris-Pratt matcher with overlap = False).  Each time a
377 | ||| match is found at index i, the prefix of target of length i
378 | ||| (that is, the portion preceding the match) is emitted as the next
379 | ||| chunk.  The matched substring itself is not included.
380 | |||
381 | ||| After emitting this prefix, the function continues splitting the
382 | ||| remainder of target, skipping over the full match of length
383 | ||| i + length pat.  This process continues until the entire target
384 | ||| has been consumed.
385 | |||
386 | ||| Unlike splitKeepFrontKMP and splitKeepEndKMP, which include the
387 | ||| matched pattern in each emitted chunk, splitDropKMP removes all
388 | ||| occurrences of pat from the output.
389 | |||
390 | ||| If pat does not occur in target, the result is a singleton list
391 | ||| containing the original target.
392 | |||
393 | export
394 | splitDropKMP :  (pat : ByteString)
395 |              -> (target : ByteString)
396 |              -> {0 prfpat : So (not $ null pat)}
397 |              -> {0 prftarget : So (not $ null target)}
398 |              -> {0 prflength : So ((length target) >= (length pat))}
399 |              -> F1 s (List ByteString)
400 | splitDropKMP pat target {prfpat} {prftarget} {prflength} t =
401 |   let splitter' # t := splitter pat target Lin t
402 |     in (splitter' <>> []) # t
403 |   where
404 |     splitter :  (pat : ByteString)
405 |              -> (target : ByteString)
406 |              -> (final : SnocList ByteString)
407 |              -> F1 s (SnocList ByteString)
408 |     splitter pat target final t =
409 |       let matcher' # t := matcher False pat [target] t
410 |         in case matcher' of
411 |              []       =>
412 |                let final' := final :< target
413 |                  in final' # t
414 |              (i :: _) =>
415 |                let length' := plus (cast {to=Nat} i) (length pat)
416 |                    final'  := final :< (take (cast {to=Nat} i) target)
417 |                  in assert_total (splitter pat (drop length' target) final' t)
418 |
419 | ||| Replaces all non-overlapping occurrences of a pattern in a ByteString
420 | ||| using the Knuth-Morris-Pratt matcher.
421 | |||
422 | ||| This function repeatedly searches target for occurrences of pat
423 | ||| (using matcher False). Each time a match is found at index i:
424 | |||
425 | ||| * If i == 0, the match is at the current position. The matched
426 | |||   segment is dropped and sub is appended to the result (unless
427 | |||   sub is empty, in which case nothing is appended).
428 | |||
429 | ||| * If i > 0, the prefix take i target is appended to the result,
430 | |||   followed by sub (unless sub is empty). The matched segment is
431 | |||   then dropped and processing continues on the remaining suffix.
432 | |||
433 | ||| If no further matches are found, the remaining target is appended
434 | ||| unchanged and the result is returned.
435 | |||
436 | ||| The result is accumulated via a `SnocList` and returned as a `List
437 | ||| ByteString`, preserving left-to-right order of the produced chunks.
438 | |||
439 | export
440 | replaceKMP :  (pat : ByteString)
441 |            -> (sub : ByteString)
442 |            -> (target : ByteString)
443 |            -> {0 prfpat : So (not $ null pat)}
444 |            -> {0 prftarget : So (not $ null target)}
445 |            -> {0 prflength : So ((length target) >= (length pat))}
446 |            -> F1 s (List ByteString)
447 | replaceKMP pat sub target {prfpat} {prftarget} {prflength} t =
448 |   let replacer' # t := replacer pat sub target Lin t
449 |     in (replacer' <>> []) # t
450 |   where
451 |     replacer :  (pat : ByteString)
452 |              -> (sub : ByteString)
453 |              -> (target : ByteString)
454 |              -> (final : SnocList ByteString)
455 |              -> F1 s (SnocList ByteString)
456 |     replacer pat sub target final t =
457 |       let matcher' # t := matcher False pat [target] t
458 |         in case matcher' of
459 |              []       =>
460 |                let final' := final :< target
461 |                  in final' # t
462 |              (i :: _) =>
463 |                case i of
464 |                  0 =>
465 |                    case null sub of
466 |                      True  =>
467 |                        assert_total (replacer pat sub (drop (length pat) target) final t)
468 |                      False =>
469 |                        let final' := final :< sub
470 |                          in assert_total (replacer pat sub (drop (length pat) target) final') t
471 |                  _ =>
472 |                    case null sub of
473 |                      True  =>
474 |                        let length' := plus (cast {to=Nat} i) (length pat)
475 |                            final'  := final :< (take (cast {to=Nat} i) target)
476 |                          in assert_total (replacer pat sub (drop length' target) final' t)
477 |                      False =>
478 |                        let length' := plus (cast {to=Nat} i) (length pat)
479 |                            final'  := final :< (take (cast {to=Nat} i) target) :< sub
480 |                          in assert_total (replacer pat sub (drop length' target) final' t)
481 |