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