ocaml/stdlib/string.ml

(**************************************************************************)
(*                                                                        *)
(*                                 OCaml                                  *)
(*                                                                        *)
(*           Damien Doligez, projet Gallium, INRIA Rocquencourt           *)
(*                                                                        *)
(*   Copyright 2014 Institut National de Recherche en Informatique et     *)
(*     en Automatique.                                                    *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

(* String operations, based on byte sequence operations *)

(* WARNING: Some functions in this file are duplicated in bytes.ml for
   efficiency reasons. When you modify the one in this file you need to
   modify its duplicate in bytes.ml.
   These functions have a "duplicated" comment above their definition.
*)

external length : string -> int = "%string_length"
external get : string -> int -> char = "%string_safe_get"
external unsafe_get : string -> int -> char = "%string_unsafe_get"
external unsafe_blit : string -> int ->  bytes -> int -> int -> unit
                     = "caml_blit_string" [@@noalloc]

module B = Bytes

let bts = B.unsafe_to_string
let bos = B.unsafe_of_string

let make n c =
  B.make n c |> bts
let init n f =
  B.init n f |> bts
let empty = ""
let of_char c = make 1 c
let of_bytes = B.to_string
let to_bytes = B.of_string
let sub s ofs len =
  if ofs = 0 && length s = len then s else
  B.sub (bos s) ofs len |> bts
let blit =
  B.blit_string

let ensure_ge (x:int) y = if x >= y then x else invalid_arg "String.concat"

let rec sum_lengths acc seplen = function
  | [] -> acc
  | hd :: [] -> length hd + acc
  | hd :: tl -> sum_lengths (ensure_ge (length hd + seplen + acc) acc) seplen tl

let rec unsafe_blits dst pos sep seplen = function
    [] -> dst
  | hd :: [] ->
    unsafe_blit hd 0 dst pos (length hd); dst
  | hd :: tl ->
    unsafe_blit hd 0 dst pos (length hd);
    unsafe_blit sep 0 dst (pos + length hd) seplen;
    unsafe_blits dst (pos + length hd + seplen) sep seplen tl

let concat sep = function
    [] -> ""
  | [s] -> s
  | l -> let seplen = length sep in bts @@
          unsafe_blits
            (B.create (sum_lengths 0 seplen l))
            0 sep seplen l

let cat = ( ^ )

(* duplicated in bytes.ml *)
let iter f s =
  for i = 0 to length s - 1 do f (unsafe_get s i) done

(* duplicated in bytes.ml *)
let iteri f s =
  for i = 0 to length s - 1 do f i (unsafe_get s i) done

let map f s =
  B.map f (bos s) |> bts
let mapi f s =
  B.mapi f (bos s) |> bts
let fold_right f x a =
  B.fold_right f (bos x) a
let fold_left f a x =
  B.fold_left f a (bos x)
let exists f s =
  B.exists f (bos s)
let for_all f s =
  B.for_all f (bos s)

(* Beware: we cannot use B.trim or B.escape because they always make a
   copy, but String.mli spells out some cases where we are not allowed
   to make a copy. *)

let is_space = function
  | ' ' | '\012' | '\n' | '\r' | '\t' -> true
  | _ -> false

let trim s =
  if s = "" then s
  else if is_space (unsafe_get s 0) || is_space (unsafe_get s (length s - 1))
    then bts (B.trim (bos s))
  else s

let escaped s =
  let b = bos s in
  (* We satisfy [unsafe_escape]'s precondition by passing an
     immutable byte sequence [b]. *)
  let b' = B.unsafe_escape b in
  (* With js_of_ocaml, [bos] and [bts] are not the identity.
     We can avoid a [bts] conversion if [unsafe_escape] returned
     its argument. *)
  if b == b' then s else bts b'

(* Finding indices *)

let invalid_start ~start len =
  let i = string_of_int in
  invalid_arg @@ concat "" ["start: "; i start; " not in range [0;"; i len; "]"]

let find_first_index sat ?(start = 0) s =
  let len = length s in
  if not (0 <= start && start <= len) then invalid_start ~start len else
  let i = ref start in
  while !i < len && not (sat (unsafe_get s !i)) do incr i done;
  if !i < len then Some !i else None

let find_last_index sat ?start s =
  let len = length s in
  let start = match start with None -> len | Some s -> s in
  if not (0 <= start && start <= len) then invalid_start ~start len else
  let i = ref (if start = len then len - 1 else start) in
  while !i >= 0 && not (sat (unsafe_get s !i)) do decr i done;
  if !i < 0 then None else Some !i

(* duplicated in bytes.ml *)
let rec index_rec s lim i c =
  if i >= lim then raise Not_found else
  if unsafe_get s i = c then i else index_rec s lim (i + 1) c

(* duplicated in bytes.ml *)
let index s c = index_rec s (length s) 0 c

(* duplicated in bytes.ml *)
let rec index_rec_opt s lim i c =
  if i >= lim then None else
  if unsafe_get s i = c then Some i else index_rec_opt s lim (i + 1) c

(* duplicated in bytes.ml *)
let index_opt s c = index_rec_opt s (length s) 0 c

(* duplicated in bytes.ml *)
let index_from s i c =
  let l = length s in
  if i < 0 || i > l then invalid_arg "String.index_from / Bytes.index_from" else
    index_rec s l i c

(* duplicated in bytes.ml *)
let index_from_opt s i c =
  let l = length s in
  if i < 0 || i > l then
    invalid_arg "String.index_from_opt / Bytes.index_from_opt"
  else
    index_rec_opt s l i c

(* duplicated in bytes.ml *)
let rec rindex_rec s i c =
  if i < 0 then raise Not_found else
  if unsafe_get s i = c then i else rindex_rec s (i - 1) c

(* duplicated in bytes.ml *)
let rindex s c = rindex_rec s (length s - 1) c

(* duplicated in bytes.ml *)
let rindex_from s i c =
  if i < -1 || i >= length s then
    invalid_arg "String.rindex_from / Bytes.rindex_from"
  else
    rindex_rec s i c

(* duplicated in bytes.ml *)
let rec rindex_rec_opt s i c =
  if i < 0 then None else
  if unsafe_get s i = c then Some i else rindex_rec_opt s (i - 1) c

(* duplicated in bytes.ml *)
let rindex_opt s c = rindex_rec_opt s (length s - 1) c

(* duplicated in bytes.ml *)
let rindex_from_opt s i c =
  if i < -1 || i >= length s then
    invalid_arg "String.rindex_from_opt / Bytes.rindex_from_opt"
  else
    rindex_rec_opt s i c

(* duplicated in bytes.ml *)
let contains_from s i c =
  let l = length s in
  if i < 0 || i > l then
    invalid_arg "String.contains_from / Bytes.contains_from"
  else
    try ignore (index_rec s l i c); true with Not_found -> false

(* duplicated in bytes.ml *)
let contains s c = contains_from s 0 c

(* duplicated in bytes.ml *)
let rcontains_from s i c =
  if i < 0 || i >= length s then
    invalid_arg "String.rcontains_from / Bytes.rcontains_from"
  else
    try ignore (rindex_rec s i c); true with Not_found -> false

(* Finding substrings *)

module Search = struct
  (* Two way string search, see https://doi.org/10.1145/116825.116845 or
     http://www-igm.univ-mlv.fr/~lecroq/string/node26.html#SECTION00260 *)

  let find_maximal_suffix_and_period ~sub =
    let sublen = length sub in
    let i = ref (-1) and j = ref 0 and k = ref 1 and p = ref 1 in
    let[@inline] maximal_suffix ~order =
      while (!j + !k < sublen) do
        let c = order * Char.compare (get sub (!j + !k)) (get sub (!i + !k)) in
        if c < 0 then (j := !j + !k; k := 1; p := !j - !i) else
        if c > 0 then (i := !j; j := !i + 1; k := 1; p := 1) else
        (* c = 0 *)
        if !k = !p then (j := !j + !p; k := 1) else incr k
      done;
    in
    (maximal_suffix[@inlined]) ~order:1;
    let l0 = !i and p0 = !p in
    i := -1; j := 0; k := 1; p := 1;
    (maximal_suffix[@inlined]) ~order:(-1);
    let l1 = !i and p1 = !p in
    if l0 > l1 then (l0, p0) else (l1, p1)

  let is_sub_periodic ~sub ~sub_lp:(l, p) =
    l <= p &&
    let i = ref 0 in
    while !i <= l && Char.equal (get sub !i) (get sub (!i + p))
    do incr i done;
    !i > l

  let find ~start ~sub ~sub_lp:(l, p) ~sub_periodic s = (* -1 on not found *)
    let slen = length s and sublen = length sub in
    if not (0 <= start && start <= slen) then invalid_start ~start slen else
    let smax = slen - sublen in
    let j = ref start in
    try
      if sub_periodic then begin
        let memory = ref (-1) in
        while (!j <= smax) do
          let i = ref (1 + Int.max l !memory) in
          while (!i < sublen && Char.equal (get sub !i) (get s (!i + !j)))
          do incr i done;
          if !i < sublen then (j := !j + (!i - l); memory := -1) else
          begin
            i := l;
            while (!i > !memory && Char.equal (get sub !i) (get s (!i + !j)))
            do decr i done;
            if !i <= !memory then raise_notrace Exit else
            (j := !j + p; memory := sublen - p - 1)
          end
        done;
        -1
      end else begin
        let q = 1 + Int.max (l + 1) (sublen - l - 1) in
        while (!j <= smax) do
          let i = ref (l + 1) in
          while (!i < sublen && Char.equal (get sub !i) (get s (!i + !j)))
          do incr i done;
          if !i < sublen then (j := !j + (!i - l)) else
          begin
            i := l;
            while (!i >= 0 && Char.equal (get sub !i) (get s (!i + !j)))
            do decr i done;
            if !i < 0 then raise_notrace Exit else (j := !j + q)
          end
        done;
        -1
      end
    with Exit -> !j

  (* The following searches from the end of string. Except for changes
     marked with an explicit comment this a cut and paste of the above
     code to search forward but with the [get] function mapping
     indices from the range [0;n-1] to [n-1;0] which we call the
     "reverse space" below *)

  let[@inline] get s i = get s (length s - 1 - i)

  let rfind_maximal_suffix_and_period ~sub =
    let sublen = length sub in
    let i = ref (-1) and j = ref 0 and k = ref 1 and p = ref 1 in
    let[@inline] maximal_suffix ~order =
      while (!j + !k < sublen) do
        let c = order * Char.compare (get sub (!j + !k)) (get sub (!i + !k)) in
        if c < 0 then (j := !j + !k; k := 1; p := !j - !i) else
        if c > 0 then (i := !j; j := !i + 1; k := 1; p := 1) else
        (* c = 0 *)
        if !k = !p then (j := !j + !p; k := 1) else incr k
      done;
    in
    (maximal_suffix[@inlined]) ~order:1;
    let l0 = !i and p0 = !p in
    i := -1; j := 0; k := 1; p := 1;
    (maximal_suffix[@inlined]) ~order:(-1);
    let l1 = !i and p1 = !p in
    if l0 > l1 then (l0, p0) else (l1, p1)

  let ris_sub_periodic ~sub ~rsub_lp:(l, p) =
    l <= p &&
    let i = ref 0 in
    while !i <= l && Char.equal (get sub !i) (get sub (!i + p))
    do incr i done;
    !i > l

  let rfind ~start ~sub ~rsub_lp:(l, p) ~rsub_periodic s = (* -1 on not found *)
    let slen = length s and sublen = length sub in
    if not (0 <= start && start <= slen) then invalid_start ~start slen else
    let start =
      (* In the reverse space we start searches at the index of the end of
         [sub] so we need to adjust start to search from there which is the
         index [start + sublen - 1]. This index is then converted into the
         reverse index space. That may end up negative, e.g. if [start] is
         toward the end and [sub] is large so we clamp to 0. *)
      Int.max 0 (slen - 1 - (start + (sublen - 1)))
    in
    let smax = slen - sublen in
    let j = ref start in
    try
      if rsub_periodic then begin
        let memory = ref (-1) in
        while (!j <= smax) do
          let i = ref (1 + Int.max l !memory) in
          while (!i < sublen && Char.equal (get sub !i) (get s (!i + !j)))
          do incr i done;
          if !i < sublen then (j := !j + (!i - l); memory := -1) else
          begin
            i := l;
            while (!i > !memory && Char.equal (get sub !i) (get s (!i + !j)))
            do decr i done;
            if !i <= !memory then raise_notrace Exit else
            (j := !j + p; memory := sublen - p - 1)
          end
        done;
        -1
      end else begin
        let q = 1 + Int.max (l + 1) (sublen - l - 1) in
        while (!j <= smax) do
          let i = ref (l + 1) in
          while (!i < sublen && Char.equal (get sub !i) (get s (!i + !j)))
          do incr i done;
          if !i < sublen then (j := !j + (!i - l)) else
          begin
            i := l;
            while (!i >= 0 && Char.equal (get sub !i) (get s (!i + !j)))
            do decr i done;
            if !i < 0 then raise_notrace Exit else (j := !j + q)
          end
        done;
        -1
      end
    with Exit ->
      (* This transforms back from the reverse space and compensates
         for the fact that we found the index of the end of [sub]. *)
      slen - 1 - (!j + (sublen - 1))
end

let find_first ~sub =
  let sub_lp = Search.find_maximal_suffix_and_period ~sub in
  let sub_periodic = Search.is_sub_periodic ~sub ~sub_lp in
  fun ?(start = 0) s ->
    match Search.find ~start ~sub ~sub_lp ~sub_periodic s with
    | -1 -> None | i -> Some i

let find_last ~sub =
  let rsub_lp = Search.rfind_maximal_suffix_and_period ~sub in
  let rsub_periodic = Search.ris_sub_periodic ~sub ~rsub_lp in
  fun ?start s ->
    let start = match start with None -> length s | Some s -> s in
    match Search.rfind ~start ~sub ~rsub_lp ~rsub_periodic s with
    | -1 -> None | i -> Some i

let find_all ~sub =
  let sub_lp = Search.find_maximal_suffix_and_period ~sub in
  let sub_periodic = Search.is_sub_periodic ~sub ~sub_lp in
  fun f ?(start = 0) s acc ->
    let rec loop f acc sub sub_lp sub_periodic s ~start ~slen =
      if start > slen then acc else
      match Search.find ~start ~sub ~sub_lp ~sub_periodic s with
      | -1 -> acc
      | i ->
          let start = i + Int.max (length sub) 1 in
          loop f (f i acc) sub sub_lp sub_periodic s ~start ~slen
    in
    let slen = length s in
    if not (0 <= start && start <= slen) then invalid_start ~start slen else
    loop f acc sub sub_lp sub_periodic s ~start ~slen

let rfind_all ~sub =
  let rsub_lp = Search.rfind_maximal_suffix_and_period ~sub in
  let rsub_periodic = Search.ris_sub_periodic ~sub ~rsub_lp in
  fun f ?start s acc ->
    let rec loop f acc sub rsub_lp rsub_periodic s ~start ~slen =
      if start < 0 then acc else
      match Search.rfind ~start ~sub ~rsub_lp ~rsub_periodic s with
      | -1 -> acc
      | i ->
          let start = i - Int.max (length sub) 1 in
          loop f (f i acc) sub rsub_lp rsub_periodic s ~start ~slen
    in
    let slen = length s in
    let start = match start with None -> length s | Some s -> s in
    if not (0 <= start && start <= slen) then invalid_start ~start slen else
    loop f acc sub rsub_lp rsub_periodic s ~start ~slen

let replace_first ~sub:needle =
  let find_first = find_first ~sub:needle in
  fun ~by ?start s ->
    match find_first ?start s with
    | None -> s
    | Some i ->
        let rest_first = i + length needle in
        let rest_len = length s - i - length needle in
        concat by [sub s 0 i; sub s rest_first rest_len]

let replace_last ~sub:needle =
  let find_last = find_last ~sub:needle in
  fun ~by ?start s ->
    match find_last ?start s with
    | None -> s
    | Some i ->
        let rest_first = i + length needle in
        let rest_len = length s - i - length needle in
        concat by [sub s 0 i; sub s rest_first rest_len]

let replace_all ~sub:needle =
  let find_all = find_all ~sub:needle in
  fun ~by ?start s ->
    let chunk_first = ref 0 in
    let add_chunk i acc =
      let acc = sub s !chunk_first (i - !chunk_first) :: acc in
      chunk_first := i + length needle; acc
    in
    match find_all add_chunk ?start s [] with
    | [] -> s
    | chunks ->
        let chunks = sub s !chunk_first (length s - !chunk_first) :: chunks in
        concat by (List.rev chunks)

(* ASCII transforms *)

let uppercase_ascii s =
  B.uppercase_ascii (bos s) |> bts
let lowercase_ascii s =
  B.lowercase_ascii (bos s) |> bts
let capitalize_ascii s =
  B.capitalize_ascii (bos s) |> bts
let uncapitalize_ascii s =
  B.uncapitalize_ascii (bos s) |> bts

(* duplicated in bytes.ml *)
let starts_with ~prefix s =
  let len_s = length s
  and len_pre = length prefix in
  let rec aux i =
    if i = len_pre then true
    else if unsafe_get s i <> unsafe_get prefix i then false
    else aux (i + 1)
  in len_s >= len_pre && aux 0

(* duplicated in bytes.ml *)
let ends_with ~suffix s =
  let len_s = length s
  and len_suf = length suffix in
  let diff = len_s - len_suf in
  let rec aux i =
    if i = len_suf then true
    else if unsafe_get s (diff + i) <> unsafe_get suffix i then false
    else aux (i + 1)
  in diff >= 0 && aux 0

let includes ~affix:sub =
  let sub_lp = Search.find_maximal_suffix_and_period ~sub in
  let sub_periodic = Search.is_sub_periodic ~sub ~sub_lp in
  fun s -> Search.find ~start:0 ~sub ~sub_lp ~sub_periodic s <> -1

external seeded_hash : int -> string -> int = "caml_string_hash" [@@noalloc]
let hash x = seeded_hash 0 x

(* Splitting with magnitudes *)

let[@inline] subrange ?(first = 0) ?(last = max_int) s =
  (* assert (Sys.max_string_length - 1 < max_int) *)
  let max = length s - 1 in
  let first = if first < 0 then 0 else first in
  let last = if last > max then max else last in
  if first > last then "" else sub s first (last - first + 1)

let take_first n s = subrange ~last:(n - 1) s
let drop_first n s = subrange ~first:n s
let cut_first n s = (take_first n s, drop_first n s)
let take_last n s = subrange ~first:(length s - n) s
let drop_last n s = subrange ~last:(length s - n - 1) s
let cut_last n s = (drop_last n s, take_last n s)

let drop_prefix ~prefix s =
  if starts_with ~prefix s
  then Some (drop_first (length prefix) s)
  else None

let drop_suffix ~suffix s =
  if ends_with ~suffix s
  then Some (drop_last (length suffix) s)
  else None

(* Splitting with predicates *)

let take_first_while sat s =
  let len = length s and i = ref 0 in
  while !i < len && sat (unsafe_get s !i) do incr i done;
  if !i = len then s else sub s 0 !i

let drop_first_while sat s =
  let len = length s and i = ref 0 in
  while !i < len && sat (unsafe_get s !i) do incr i done;
  if !i = 0 then s else sub s !i (len - !i)

let cut_first_while sat s =
  let len = length s and i = ref 0 in
  while !i < len && sat (unsafe_get s !i) do incr i done;
  if !i = len then s, "" else
  if !i = 0 then "", s else
  sub s 0 !i, sub s !i (len - !i)

let take_last_while sat s =
  let len = length s in
  let i = ref (len - 1) in
  while !i >= 0 && sat (unsafe_get s !i) do decr i done;
  if !i < 0 then s else
  let j = !i + 1 in
  sub s j (len - j)

let drop_last_while sat s =
  let len = length s in
  let i = ref (len - 1) in
  while !i >= 0 && sat (unsafe_get s !i) do decr i done;
  if !i < 0 then "" else sub s 0 (!i + 1)

let cut_last_while sat s =
  let len = length s in
  let i = ref (len - 1) in
  while !i >= 0 && sat (unsafe_get s !i) do decr i done;
  if !i < 0 then "", s else
  if !i = len - 1 then s, "" else
  let j = !i + 1 in
  sub s 0 j, sub s j (len - j)

(* Splitting with separators *)

let split_first ~sep =
  let find_first = find_first ~sub:sep in
  fun s -> match find_first s with
  | None -> None
  | Some i ->
      Some (subrange ~last:(i - 1) s, subrange ~first:(i + length sep) s)

let split_last ~sep =
  let find_last = find_last ~sub:sep in
  fun s -> match find_last s with
  | None -> None
  | Some i ->
      Some (subrange ~last:(i - 1) s, subrange ~first:(i + length sep) s)

let split_all ~sep =
  let find_all = find_all ~sub:sep in
  fun ?(drop = fun _ -> false) s ->
    let first = ref 0 in
    let add_token i acc =
      let token = subrange ~first:!first ~last:(i - 1) s in
      first := i + length sep;
      if drop token then acc else token :: acc
    in
    let tokens = find_all add_token s [] in
    let last = subrange ~first:!first s in
    List.rev (if drop last then tokens else last :: tokens)

let rsplit_all ~sep =
  let rfind_all = rfind_all ~sub:sep in
  fun ?(drop = fun _ -> false) s ->
    let last = ref (length s - 1) in
    let add_token i acc =
      let token = subrange ~first:(i + length sep) ~last:!last s in
      last := i - 1;
      if drop token then acc else token :: acc
    in
    let tokens = rfind_all add_token s [] in
    let last = subrange ~last:!last s in
    if drop last then tokens else (last :: tokens)

(* duplicated in bytes.ml *)
let split_on_char sep s =
  let r = ref [] in
  let j = ref (length s) in
  for i = length s - 1 downto 0 do
    if unsafe_get s i = sep then begin
      r := sub s (i + 1) (!j - i - 1) :: !r;
      j := i
    end
  done;
  sub s 0 !j :: !r

type t = string

let compare (x: t) (y: t) = Stdlib.compare x y
external equal : string -> string -> bool = "caml_string_equal" [@@noalloc]
let is_empty s = Int.equal (length s) 0

(** {1 Iterators} *)

let to_seq s = bos s |> B.to_seq

let to_seqi s = bos s |> B.to_seqi

let of_seq g = B.of_seq g |> bts

(* UTF decoders and validators *)

let get_utf_8_uchar s i = B.get_utf_8_uchar (bos s) i
let is_valid_utf_8 s = B.is_valid_utf_8 (bos s)

let get_utf_16be_uchar s i = B.get_utf_16be_uchar (bos s) i
let is_valid_utf_16be s = B.is_valid_utf_16be (bos s)

let get_utf_16le_uchar s i = B.get_utf_16le_uchar (bos s) i
let is_valid_utf_16le s = B.is_valid_utf_16le (bos s)

(** {6 Binary encoding/decoding of integers} *)

external get_uint8 : string -> int -> int = "%string_safe_get"
external get_uint16_ne : string -> int -> int = "%caml_string_get16"
external get_int32_ne : string -> int -> int32 = "%caml_string_get32"
external get_int64_ne : string -> int -> int64 = "%caml_string_get64"

let get_int8 s i = B.get_int8 (bos s) i
let get_uint16_le s i = B.get_uint16_le (bos s) i
let get_uint16_be s i = B.get_uint16_be (bos s) i
let get_int16_ne s i = B.get_int16_ne (bos s) i
let get_int16_le s i = B.get_int16_le (bos s) i
let get_int16_be s i = B.get_int16_be (bos s) i
let get_int32_le s i = B.get_int32_le (bos s) i
let get_int32_be s i = B.get_int32_be (bos s) i
let get_int64_le s i = B.get_int64_le (bos s) i
let get_int64_be s i = B.get_int64_be (bos s) i

(* Spellchecking *)

let utf_8_uchar_length s =
  let slen = length s in
  let i = ref 0 and ulen = ref 0 in
  while (!i < slen) do
    let dec_len = Uchar.utf_8_decode_length_of_byte (unsafe_get s !i) in
    i := (!i + if dec_len = 0 then 1 (* count one Uchar.rep *) else dec_len);
    incr ulen;
  done;
  !ulen

let uchar_array_of_utf_8_string s =
  let slen = length s in (* is an upper bound on Uchar.t count *)
  let uchars = Array.make slen Uchar.max in
  let k = ref 0 and i = ref 0 in
  while (!i < slen) do
    let dec = get_utf_8_uchar s !i in
    i := !i + Uchar.utf_decode_length dec;
    uchars.(!k) <- Uchar.utf_decode_uchar dec;
    incr k;
  done;
  uchars, !k

let edit_distance' ?(limit = Int.max_int) s (s0, len0) s1 =
  if limit <= 1 then (if equal s s1 then 0 else limit) else
  let[@inline] minimum a b c = Int.min a (Int.min b c) in
  let s1, len1 = uchar_array_of_utf_8_string s1 in
  let limit = Int.min (Int.max len0 len1) limit in
  if Int.abs (len1 - len0) >= limit then limit else
  let s0, s1 = if len0 > len1 then s0, s1 else s1, s0 in
  let len0, len1 = if len0 > len1 then len0, len1 else len1, len0 in
  let rec loop row_minus2 row_minus1 row i len0 limit s0 s1 =
    if i > len0 then row_minus1.(Array.length row_minus1 - 1) else
    let len1 = Array.length row - 1 in
    let row_min = ref Int.max_int in
    row.(0) <- i;
    let jmax =
      let jmax = Int.min len1 (i + limit - 1) in
      if jmax < 0 then (* overflow *) len1 else jmax
    in
    for j = Int.max 1 (i - limit) to jmax do
      let cost = if Uchar.equal s0.(i-1) s1.(j-1) then 0 else 1 in
      let min = minimum
          (row_minus1.(j-1) + cost) (* substitute *)
          (row_minus1.(j) + 1)      (* delete *)
          (row.(j-1) + 1)           (* insert *)
          (* Note when j = i - limit, the latter [row] read makes a bogus read
             on the value that was in the matrix at d.(i-2).(i - limit - 1).
             Since by induction for all i,j, d.(i).(j) >= abs (i - j),
             (row.(j-1) + 1) is greater or equal to [limit] and thus does
             not affect adversely the minimum computation. *)
      in
      let min =
        if (i > 1 && j > 1 &&
            Uchar.equal s0.(i-1) s1.(j-2) &&
            Uchar.equal s0.(i-2) s1.(j-1))
        then Int.min min (row_minus2.(j-2) + cost) (* transpose *)
        else min
      in
      row.(j) <- min;
      row_min := Int.min !row_min min;
    done;
    if !row_min >= limit then (* can no longer decrease *) limit else
    loop row_minus1 row row_minus2 (i + 1) len0 limit s0 s1
  in
  let ignore =
    (* Value used to make the values around the diagonal stripe ignored
       by the min computations when we have a limit. *)
    limit + 1
  in
  let row_minus2 = Array.make (len1 + 1) ignore in
  let row_minus1 = Array.init (len1 + 1) (fun x -> x) in
  let row = Array.make (len1 + 1) ignore in
  let d = loop row_minus2 row_minus1 row 1 len0 limit s0 s1 in
  if d > limit then limit else d

let edit_distance ?limit s0 s1 =
  let us0 = uchar_array_of_utf_8_string s0 in
  edit_distance' ?limit s0 us0 s1

let default_max_dist s = match utf_8_uchar_length s with
  | 0 | 1 | 2 -> 0
  | 3 | 4 -> 1
  | _ -> 2

let spellcheck ?(max_dist = default_max_dist) iter_dict s =
  let min = ref (max_dist s) in
  let acc = ref [] in
  let select_words s us word =
    let d = edit_distance' ~limit:(!min + 1) s us word in
    if d = !min then (acc := word :: !acc) else
    if d < !min then (min := d; acc := [word]) else ()
  in
  let us = uchar_array_of_utf_8_string s in
  iter_dict (select_words s us);
  List.rev !acc