<commit_before>
(* This file is free software, part of containers. See file "license" for more details. *)

(** {1 Map specialized for Int keys} *)

(* "Fast Mergeable Integer Maps", Okasaki & Gill.
   We use big-endian trees. *)

(** Masks with exactly one bit active *)
module Bit : sig
  type t = private int
  val highest : int -> t
  val min_int : t
  val equal : t -> t -> bool
  val is_0 : bit:t -> int -> bool
  val is_1 : bit:t -> int -> bool
  val mask : mask:t -> int -> int   (* zeroes the bit, puts all lower bits to 1 *)
  val lt : t -> t -> bool
  val gt : t -> t -> bool
  val equal_int : int -> t -> bool
end = struct
  type t = int

  let min_int = min_int

  let equal : t -> t -> bool = Stdlib.(=)

  let rec highest_bit_naive x m =
    if x=m then m
    else highest_bit_naive (x land (lnot m)) (2*m)

  let mask_20_ = 1 lsl 20
  let mask_40_ = 1 lsl 40

  let highest x =
    if x<0 then min_int
    else if Sys.word_size > 40 && x > mask_40_ then (
      (* remove least significant 40 bits *)
      let x' = x land (lnot (mask_40_ -1)) in
      highest_bit_naive x' mask_40_
    ) else if x> mask_20_ then (
      (* small shortcut: remove least significant 20 bits *)
      let x' = x land (lnot (mask_20_ -1)) in
      highest_bit_naive x' mask_20_
    ) else (
      highest_bit_naive x 1
    )

  let is_0 ~bit x = x land bit = 0
  let is_1 ~bit x = x land bit = bit

  let mask ~mask x = (x lor (mask -1)) land (lnot mask)
  (* low endian: let mask_ x ~mask = x land (mask - 1) *)

  let gt a b = (b != min_int) && (a = min_int || a > b)
  let lt a b = gt b a
  let equal_int = Stdlib.(=)
end

(*$inject
  let highest2 x : int =
    let rec aux i =
      if i=0 then i
      else if 1 = (x lsr i) then 1 lsl i else aux (i-1)
    in
    if x<0 then min_int else aux (Sys.word_size-2)
*)

(*$QR & ~count:1_000
  Q.int (fun x ->
    if Bit.equal_int (highest2 x) (Bit.highest x) then true
    else QCheck.Test.fail_reportf "x=%d, highest=%d, highest2=%d@." x
      (Bit.highest x :> int) (highest2 x))
  *)

(*$inject
  let _list_uniq l = CCList.sort_uniq ~cmp:(fun a b-> Stdlib.compare (fst a)(fst b)) l
*)

type +'a t =
  | E  (* empty *)
  | L of int * 'a  (* leaf *)
  | N of int (* common prefix *) * Bit.t (* bit switch *) * 'a t * 'a t

let empty = E

let is_empty = function
  | E -> true
  | _ -> false

(*$Q
  Q.(small_list (pair int int)) (fun l -> \
    let m = of_list l in \
    is_empty m = (cardinal m = 0))
  *)

let is_prefix_ ~prefix y ~bit =
  prefix = Bit.mask y ~mask:bit

(*$Q
  Q.int (fun i -> \
    let b = Bit.highest i in \
    ((b:>int) land i = (b:>int)) && (i < 0 || ((b:>int) <= i && (i-(b:>int)) < (b:>int))))
  Q.int (fun i -> (Bit.highest i = Bit.min_int) = (i < 0))
  Q.int (fun i -> ((Bit.highest i:>int) < 0) = (Bit.highest i = Bit.min_int))
  Q.int (fun i -> let j = (Bit.highest i :> int) in  j land (j-1) = 0)
*)

(*$T
  (Bit.highest min_int :> int) = min_int
  (Bit.highest 2 :> int) = 2
  (Bit.highest 17 :> int)  = 16
  (Bit.highest 300 :> int) = 256
*)

(* helper:

    let b_of_i i =
      let rec f acc i =
        if i=0 then acc else let q, r = i/2, abs (i mod 2)
      in
      f (r::acc) q in f [] i;;
*)

(* low endian: let branching_bit_ a _ b _ = lowest_bit_ (a lxor b) *)
let branching_bit_ a b = Bit.highest (a lxor b)

(* TODO use hint in branching_bit_ *)

let check_invariants t =
  (* check that keys are prefixed by every node in their path *)
  let rec check_keys path t = match t with
    | E -> true
    | L (k, _) ->
      List.for_all
        (fun (prefix, switch, side) ->
           is_prefix_ ~prefix k ~bit:switch &&
           begin match side with
             | `Left -> Bit.is_0 k ~bit:switch
             | `Right -> Bit.is_1 k ~bit:switch
           end)
        path
    | N (prefix, switch, l, r) ->
      check_keys ((prefix, switch, `Left) :: path) l &&
      check_keys ((prefix, switch, `Right) :: path) r
  in
  check_keys [] t

(*$Q
  Q.(list (pair int bool)) (fun l -> \
    check_invariants (of_list l))
*)

let rec find_exn k t = match t with
  | E -> raise Not_found
  | L (k', v) when k = k' -> v
  | L _ -> raise Not_found
  | N (prefix, m, l, r) ->
    if is_prefix_ ~prefix k ~bit:m then (
      if Bit.is_0 k ~bit:m
      then find_exn k l
      else find_exn k r
    ) else (
      raise Not_found
    )

(* XXX could test with lt_unsigned_? *)

    (*
    if k <= prefix (* search tree *)
    then find_exn k l
    else find_exn k r
       *)

let find k t =
  try Some (find_exn k t)
  with Not_found -> None

(*$Q
  Q.(list (pair int int)) (fun l -> \
    let l = _list_uniq l in \
    let m = of_list l in \
    List.for_all (fun (k,v) -> find k m = Some v) l)
*)

let mem k t =
  try ignore (find_exn k t); true
  with Not_found -> false

(*$Q
  Q.(list (pair int int)) (fun l -> \
    let m = of_list l in \
    List.for_all (fun (k,_) -> mem k m) l)
*)

let mk_node_ prefix switch l r = match l, r with
  | E, o | o, E -> o<commit_msg>small changes in CCIntMap, some cleanup<commit_after>
(* This file is free software, part of containers. See file "license" for more details. *)

(** {1 Map specialized for Int keys} *)

(* "Fast Mergeable Integer Maps", Okasaki & Gill.
   We use big-endian trees. *)

(** Masks with exactly one bit active *)
module Bit : sig
  type t = private int
  val highest : int -> t
  val min_int : t
  val equal : t -> t -> bool
  val is_0 : bit:t -> int -> bool
  val is_1 : bit:t -> int -> bool
  val mask : mask:t -> int -> int   (* zeroes the bit, puts all lower bits to 1 *)
  val lt : t -> t -> bool
  val gt : t -> t -> bool
  val equal_int : int -> t -> bool
end = struct
  type t = int

  let min_int = min_int

  let equal : t -> t -> bool = Stdlib.(=)

  let rec highest_bit_naive x m =
    if x=m then m
    else highest_bit_naive (x land (lnot m)) (2*m)

  let mask_20_ = 1 lsl 20
  let mask_40_ = 1 lsl 40

  let highest x =
    if x<0 then min_int
    else if Sys.word_size > 40 && x > mask_40_ then (
      (* remove least significant 40 bits *)
      let x' = x land (lnot (mask_40_ -1)) in
      highest_bit_naive x' mask_40_
    ) else if x> mask_20_ then (
      (* small shortcut: remove least significant 20 bits *)
      let x' = x land (lnot (mask_20_ -1)) in
      highest_bit_naive x' mask_20_
    ) else (
      highest_bit_naive x 1
    )

  let[@inline] is_0 ~bit x = x land bit = 0
  let[@inline] is_1 ~bit x = x land bit = bit

  let mask ~mask x = (x lor (mask -1)) land (lnot mask)
  (* small endian: let mask_ x ~mask = x land (mask - 1) *)

  let gt a b = (b != min_int) && (a = min_int || a > b)
  let lt a b = gt b a
  let equal_int : int -> int -> bool = Stdlib.(=)
end

(*$inject
  let highest2 x : int =
    let rec aux i =
      if i=0 then i
      else if 1 = (x lsr i) then 1 lsl i else aux (i-1)
    in
    if x<0 then min_int else aux (Sys.word_size-2)
*)

(*$QR & ~count:1_000
  Q.int (fun x ->
    if Bit.equal_int (highest2 x) (Bit.highest x) then true
    else QCheck.Test.fail_reportf "x=%d, highest=%d, highest2=%d@." x
      (Bit.highest x :> int) (highest2 x))
  *)

(*$inject
  let _list_uniq l = CCList.sort_uniq ~cmp:(fun a b-> Stdlib.compare (fst a)(fst b)) l
*)

type +'a t =
  | E  (* empty *)
  | L of int * 'a  (* leaf *)
  | N of int (* common prefix *) * Bit.t (* bit switch *) * 'a t * 'a t

let empty = E

let[@inline] is_empty = function
  | E -> true
  | _ -> false

(*$Q
  Q.(small_list (pair int int)) (fun l -> \
    let m = of_list l in \
    is_empty m = (cardinal m = 0))
  *)

let[@inline] is_prefix_ ~prefix y ~bit =
  prefix = Bit.mask y ~mask:bit

(*$Q
  Q.int (fun i -> \
    let b = Bit.highest i in \
    ((b:>int) land i = (b:>int)) && (i < 0 || ((b:>int) <= i && (i-(b:>int)) < (b:>int))))
  Q.int (fun i -> (Bit.highest i = Bit.min_int) = (i < 0))
  Q.int (fun i -> ((Bit.highest i:>int) < 0) = (Bit.highest i = Bit.min_int))
  Q.int (fun i -> let j = (Bit.highest i :> int) in  j land (j-1) = 0)
*)

(*$T
  (Bit.highest min_int :> int) = min_int
  (Bit.highest 2 :> int) = 2
  (Bit.highest 17 :> int)  = 16
  (Bit.highest 300 :> int) = 256
*)

(* small endian: let branching_bit_ a _ b _ = lowest_bit_ (a lxor b) *)
let branching_bit_ a b = Bit.highest (a lxor b)

(* TODO use hint in branching_bit_ *)

let check_invariants t =
  (* check that keys are prefixed by every node in their path *)
  let rec check_keys path t = match t with
    | E -> true
    | L (k, _) ->
      List.for_all
        (fun (prefix, switch, side) ->
           is_prefix_ ~prefix k ~bit:switch &&
           begin match side with
             | `Left -> Bit.is_0 k ~bit:switch
             | `Right -> Bit.is_1 k ~bit:switch
           end)
        path
    | N (prefix, switch, l, r) ->
      check_keys ((prefix, switch, `Left) :: path) l &&
      check_keys ((prefix, switch, `Right) :: path) r
  in
  check_keys [] t

(*$Q
  Q.(list (pair int bool)) (fun l -> \
    check_invariants (of_list l))
*)

let rec find_exn k t = match t with
  | E -> raise Not_found
  | L (k', v) when k = k' -> v
  | L _ -> raise Not_found
  | N (prefix, m, l, r) ->
    if is_prefix_ ~prefix k ~bit:m then (
      if Bit.is_0 k ~bit:m
      then find_exn k l
      else find_exn k r
    ) else (
      raise Not_found
    )

let find k t =
  try Some (find_exn k t)
  with Not_found -> None

(*$Q
  Q.(list (pair int int)) (fun l -> \
    let l = _list_uniq l in \
    let m = of_list l in \
    List.for_all (fun (k,v) -> find k m = Some v) l)
*)

let mem k t =
  try ignore (find_exn k t); true
  with Not_found -> false

(*$Q
  Q.(list (pair int int)) (fun l -> \
    let m = of_list l in \
    List.for_all (fun (k,_) -> mem k m) l)
*)

let mk_node_ prefix switch l r = match l, r with
  | E, o | o, E -> o
{'fraction_non_alphanumeric': 0.1452991452991453, 'fraction_numerical': 0.016167233034702914, 'mean_word_length': 2.1821756225425952, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 125, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3997164', 'n_tokens_mistral': 3960, 'n_tokens_neox': 3763, 'n_words': 1410}