optimizations for KMP

benchs/run_benchs.ml

@@ -1636,6 +1636,159 @@ module Str = struct
     assert (CCList.equal CCInt.equal (mk_naive ()) (mk_current ()));
     B.throughputN 3 ~repeat [ "naive", mk_naive, (); "current", mk_current, () ]
 
+  (* Inlined reference implementations to have a fixpoint to compae
+      CCString with.
+
+     - [Classic_kmp] is a copy of CCString's old KMP: the
+       basic ("weak") failure table and the window-shift search loop.
+     - [New_kmp] is a bit optimized. *)
+  module Classic_kmp = struct
+    type t = {
+      failure: int array;
+      pat: string;
+    }
+
+    let compile pat =
+      let len = String.length pat in
+      match len with
+      | 0 -> { failure = [||]; pat }
+      | 1 -> { failure = [| -1 |]; pat }
+      | _ ->
+        let failure = Array.make len 0 in
+        failure.(0) <- -1;
+        let i = ref 2 and j = ref 0 in
+        while !i < len do
+          if CCChar.equal pat.[!i - 1] pat.[!j] then (
+            incr j;
+            failure.(!i) <- !j;
+            incr i
+          ) else if !j = 0 then (
+            failure.(!i) <- 0;
+            incr i
+          ) else
+            j := failure.(!j)
+        done;
+        { failure; pat }
+
+    let find { failure; pat } s =
+      let len = String.length s in
+      let i = ref 0 and j = ref 0 in
+      let pat_len = String.length pat in
+      while !j < pat_len && !i + !j < len do
+        if CCChar.equal s.[!i + !j] pat.[!j] then
+          incr j
+        else (
+          let fail_offset = failure.(!j) in
+          if fail_offset >= 0 then (
+            assert (fail_offset < !j);
+            i := !i + !j - fail_offset;
+            j := fail_offset
+          ) else (
+            j := 0;
+            incr i
+          )
+        )
+      done;
+      if !j = pat_len then
+        !i
+      else
+        -1
+  end
+
+  module New_kmp = struct
+    type t = {
+      failure: int array;
+      pat: string;
+    }
+
+    let compile pat =
+      let len = String.length pat in
+      match len with
+      | 0 -> { failure = [||]; pat }
+      | 1 -> { failure = [| -1 |]; pat }
+      | _ ->
+        let failure = Array.make len 0 in
+        failure.(0) <- -1;
+        let i = ref 2 and j = ref 0 in
+        while !i < len do
+          if CCChar.equal pat.[!i - 1] pat.[!j] then (
+            incr j;
+            failure.(!i) <- !j;
+            incr i
+          ) else if !j = 0 then (
+            failure.(!i) <- 0;
+            incr i
+          ) else
+            j := failure.(!j)
+        done;
+        (* strengthen: if the char after the border equals the one we'd be
+           retrying, skip straight to that border's (already-strengthened)
+           target instead of re-comparing it and failing again *)
+        for k = 1 to len - 1 do
+          let b = failure.(k) in
+          if b >= 0 && CCChar.equal pat.[k] pat.[b] then
+            failure.(k) <- failure.(b)
+        done;
+        { failure; pat }
+
+    let find { failure; pat } s =
+      let len = String.length s in
+      let i = ref 0 and j = ref 0 in
+      let pat_len = String.length pat in
+      while !j < pat_len && !i + !j < len do
+        if CCChar.equal s.[!i + !j] pat.[!j] then
+          incr j
+        else (
+          let fail_offset = failure.(!j) in
+          if fail_offset >= 0 then (
+            assert (fail_offset < !j);
+            i := !i + !j - fail_offset;
+            j := fail_offset
+          ) else (
+            (* no usable border: advance past the mismatch (NOT [incr i]) *)
+            i := !i + !j + 1;
+            j := 0
+          )
+        )
+      done;
+      if !j = pat_len then
+        !i
+      else
+        -1
+  end
+
+  (* a worst case test to show the difference between classic and new KMP,
+     with a needle [a^size] searched in a haystack made of [a^(size-1)b]
+     blocks. *)
+  let bench_find_special ~size n =
+    let needle = String.make size 'a' in
+    let block = String.make (size - 1) 'a' ^ "b" in
+    let haystack = CCString.repeat block (n / size) in
+    pp_pb needle haystack;
+    let classic = Classic_kmp.compile needle in
+    let strong = New_kmp.compile needle in
+    let cur = CCString.find ~start:0 ~sub:needle in
+    let mk_naive () = find ~sub:needle haystack
+    and mk_classic () = Classic_kmp.find classic haystack
+    and mk_new () = New_kmp.find strong haystack
+    and mk_cur () = cur haystack in
+    assert (mk_naive () = -1);
+    assert (mk_classic () = -1);
+    assert (mk_new () = -1);
+    assert (mk_cur () = -1);
+    (* positive control: a haystack that does contain the needle *)
+    let with_match = haystack ^ needle in
+    let expected = CCString.find ~sub:needle with_match in
+    assert (Classic_kmp.find classic with_match = expected);
+    assert (New_kmp.find strong with_match = expected);
+    B.throughputN 2 ~repeat
+      [
+        "naive", mk_naive, ();
+        "classic", mk_classic, ();
+        "new", mk_new, ();
+        "cur", mk_cur, ();
+      ]
+
   let bench_find = bench_find_ ~dir:`Direct
   let bench_rfind = bench_find_ ~dir:`Reverse
 
@@ -1757,6 +1910,20 @@ module Str = struct
                     @>> app_ints (bench_find ~size:50) [ 100; 100_000; 500_000 ];
                     "500"
                     @>> app_ints (bench_find ~size:500) [ 100_000; 500_000 ];
+                    (* short haystack, long needle: needle (500) longer than
+                       haystack, so the lazy failure table is never built *)
+                    "500_short" @>> app_ints (bench_find ~size:500) [ 50; 100 ];
+                  ];
+             "find_special"
+             @>>> [
+                    "100"
+                    @>> app_ints
+                          (bench_find_special ~size:100)
+                          [ 200_000; 1_000_000 ];
+                    "1000"
+                    @>> app_ints
+                          (bench_find_special ~size:1000)
+                          [ 200_000; 1_000_000 ];
                   ];
              "find_all"
              @>>> [

src/core/CCString.ml

@@ -45,11 +45,12 @@ type _ direction =
 (* we follow https://en.wikipedia.org/wiki/Knuth–Morris–Pratt_algorithm *)
 module Find = struct
   type 'a kmp_pattern = {
-    failure: int array;
+    mutable failure: int array; (* empty=not computed yet *)
     str: string;
+    dir: 'a direction;
   }
-  (* invariant: [length failure = length str].
-     We use a phantom type to avoid mixing the directions. *)
+  (* invariant: [failure=[||] || length failure = length str].
+     The failure table is built lazily, we bypass it if len(haystack)<len(needle) *)
 
   let kmp_pattern_length p = String.length p.str
 
@@ -60,14 +61,16 @@ module Find = struct
     | Direct -> String.get
     | Reverse -> fun s i -> s.[String.length s - i - 1]
 
-  let kmp_compile_ : type a. dir:a direction -> string -> a kmp_pattern =
+  (* build the KMP failure table for [str], whose elements are read
+     according to [dir] *)
+  let kmp_failure_ : type a. dir:a direction -> string -> int array =
    fun ~dir str ->
     let len = length str in
     let get = get_ ~dir in
     (* how to read elements of the string *)
     match len with
-    | 0 -> { failure = [||]; str }
-    | 1 -> { failure = [| -1 |]; str }
+    | 0 -> [||]
+    | 1 -> [| -1 |]
     | _ ->
       (* at least 2 elements, the algorithm can work *)
       let failure = Array.make len 0 in
@@ -92,8 +95,27 @@ module Find = struct
           assert (!j > 0);
           j := failure.(!j)
       done;
+
       (* Format.printf "{@[failure:%a, str:%s@]}@." CCFormat.(array int) failure str; *)
-      { failure; str }
+
+      (* strengthen the failure function: if falling back to a border would
+         re-compare a character equal to the one that just mismatched, skip
+         straight to that border's (already-strengthened) target instead.
+         Uses the direction-aware [get] so it works for both directions. *)
+      for k = 1 to len - 1 do
+        let b = failure.(k) in
+        if b >= 0 && CCChar.equal (get str k) (get str b) then
+          failure.(k) <- failure.(b)
+      done;
+      failure
+
+  let kmp_compile_ : type a. dir:a direction -> string -> a kmp_pattern =
+   fun ~dir str -> { str; failure = [||]; dir }
+
+  let[@inline] kmp_get_ (type a) (f : a kmp_pattern) : int array =
+    if Array.length f.failure = 0 then
+      f.failure <- kmp_failure_ ~dir:f.dir f.str;
+    f.failure
 
   let kmp_compile s = kmp_compile_ ~dir:Direct s
   let kmp_rcompile s = kmp_compile_ ~dir:Reverse s
@@ -107,23 +129,24 @@ module Find = struct
     let i = ref idx in
     let j = ref 0 in
     let pat_len = kmp_pattern_length pattern in
+    let failure = kmp_get_ pattern in
     while !j < pat_len && !i + !j < len do
       let c = String.get s (!i + !j) in
       let expected = String.get pattern.str !j in
       if CCChar.equal c expected then
         (* char matches *)
         incr j
       else (
-        let fail_offset = pattern.failure.(!j) in
+        let fail_offset = failure.(!j) in
         if fail_offset >= 0 then (
           assert (fail_offset < !j);
           (* follow the failure link *)
           i := !i + !j - fail_offset;
           j := fail_offset
         ) else (
-          (* beginning of pattern *)
-          j := 0;
-          incr i
+          (* no usable border: advance past the mismatch *)
+          i := !i + !j + 1;
+          j := 0
         )
       )
     done;
@@ -141,6 +164,7 @@ module Find = struct
     let i = ref (len - idx - 1) in
     let j = ref 0 in
     let pat_len = kmp_pattern_length pattern in
+    let failure = kmp_get_ pattern in
     while !j < pat_len && !i + !j < len do
       let c = String.get s (len - !i - !j - 1) in
       let expected =
@@ -150,16 +174,16 @@ module Find = struct
         (* char matches *)
         incr j
       else (
-        let fail_offset = pattern.failure.(!j) in
+        let fail_offset = failure.(!j) in
         if fail_offset >= 0 then (
           assert (fail_offset < !j);
           (* follow the failure link *)
           i := !i + !j - fail_offset;
           j := fail_offset
         ) else (
-          (* beginning of pattern *)
-          j := 0;
-          incr i
+          (* no usable border: advance past the mismatch *)
+          i := !i + !j + 1;
+          j := 0
         )
       )
     done;
@@ -181,12 +205,14 @@ module Find = struct
     | P_KMP p -> kmp_pattern_length p
 
   let compile sub : [ `Direct ] pattern =
+    if length sub = 0 then invalid_arg "CCString.Find.compile: empty pattern";
     if length sub = 1 then
       P_char sub.[0]
     else
       P_KMP (kmp_compile sub)
 
   let rcompile sub : [ `Reverse ] pattern =
+    if length sub = 0 then invalid_arg "CCString.Find.rcompile: empty pattern";
     if length sub = 1 then
       P_char sub.[0]
     else
@@ -195,7 +221,12 @@ module Find = struct
   let find ?(start = 0) ~(pattern : [ `Direct ] pattern) s =
     match pattern with
     | P_char c -> (try String.index_from s start c with Not_found -> -1)
-    | P_KMP pattern -> kmp_find ~pattern s start
+    | P_KMP pattern ->
+      (* haystack too short, bail out without forcing kmp state *)
+      if String.length s - start < kmp_pattern_length pattern then
+        -1
+      else
+        kmp_find ~pattern s start
 
   let rfind ?start ~(pattern : [ `Reverse ] pattern) s =
     let start =
@@ -205,7 +236,12 @@ module Find = struct
     in
     match pattern with
     | P_char c -> (try String.rindex_from s start c with Not_found -> -1)
-    | P_KMP pattern -> kmp_rfind ~pattern s start
+    | P_KMP pattern ->
+      (* haystack too short, bail out without forcing kmp state *)
+      if start + 1 < kmp_pattern_length pattern then
+        -1
+      else
+        kmp_rfind ~pattern s start
 end
 
 let find ?(start = 0) ~sub =