<commit_before>open BitUtils
open Monomorphic.Int

module type Config = sig
  (* Number of bits taken from the hashed key at every step *)
  val shift_step : int

  (* Maximum size of a BitmapIndexedNode*)
  val bmnode_max : int

  (* Minimum size of an ArrayNode (must be lesser than bmnode_max) *)
  val arrnode_min : int
end

module StdConfig : Config = struct
  let shift_step = 5
  let bmnode_max = 16
  let arrnode_min = 8
end

module StdConfig32 : Config = struct
  let shift_step = 4
  let bmnode_max = 8
  let arrnode_min = 4
end

module type S = sig
  type key
  type 'a t

  val empty : 'a t
  val is_empty : 'a t -> bool
  val singleton : key -> 'a -> 'a t
  val cardinal : 'a t -> int
  val length : 'a t -> int

  val alter : key -> ('a option -> 'a option) -> 'a t -> 'a t
  val add : key -> 'a -> 'a t -> 'a t
  val add_carry : key -> 'a -> 'a t -> 'a t * 'a option
  val remove : key -> 'a t -> 'a t
  val extract : key -> 'a t -> 'a * 'a t
  val update : key -> ('a -> 'a option) -> 'a t -> 'a t
  val modify : key -> ('a -> 'a) -> 'a t -> 'a t
  val modify_def : 'a -> key -> ('a -> 'a) -> 'a t -> 'a t

  val find_exn : key -> 'a t -> 'a
  val mem : key -> 'a t -> bool
  val choose : 'a t -> key * 'a
  val pop : 'a t -> (key * 'a) * 'a t

  val keys : 'a t -> key list
  val values : 'a t -> 'a list
  val bindings : 'a t -> (key * 'a) list
  val iter : (key -> 'a -> unit) -> 'a t -> unit
  val map : ('a -> 'b) -> 'a t -> 'b t
  val mapi : (key -> 'a -> 'b) -> 'a t -> 'b t
  val filterv : ('a -> bool) -> 'a t -> 'a t
  val filter : (key -> 'a -> bool) -> 'a t -> 'a t
  val filter_map : (key -> 'a -> 'b option) -> 'a t -> 'b t
  val foldv : ('a -> 'b -> 'b) -> 'a t -> 'b -> 'b
  val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b

  val for_all : (key -> 'a -> bool) -> 'a t -> bool
  val exists : (key -> 'a -> bool) -> 'a t -> bool
  val partition : (key -> 'a -> bool) -> 'a t -> 'a t * 'a t

  val intersecti : (key -> 'a -> 'b -> 'c) -> 'a t -> 'b t -> 'c t
  val intersect : ('a -> 'b -> 'c) -> 'a t -> 'b t -> 'c t
  val merge :
    (key -> 'a option -> 'b option -> 'c option) ->
    'a t -> 'b t -> 'c t
  val union : 'a t -> 'a t -> 'a t
  val union_f : (key -> 'a -> 'a -> 'a) -> 'a t -> 'a t -> 'a t

  module Import : sig
    module type Foldable = sig
      type key
      type 'a t
      val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
    end

    module Make (M : Foldable with type key = key) : sig
      val add_from : 'a M.t -> 'a t -> 'a t
      val from : 'a M.t -> 'a t
    end

    module AssocList : sig
      val add_from : (key * 'a) list -> 'a t -> 'a t
      val from : (key * 'a) list -> 'a t
    end
  end

  module ExceptionLess : sig
    val extract : key -> 'a t -> 'a option * 'a t
    val update : key -> ('a -> 'a option) -> 'a t -> 'a t
    val modify : key -> ('a -> 'a) -> 'a t -> 'a t
    val find : key -> 'a t -> 'a option
    val choose : 'a t -> (key * 'a) option
  end

  module Infix : sig
    val ( --> ) : 'a t -> key -> 'a
    val ( <-- ) : 'a t -> key * 'a -> 'a t
  end
end

module Make (Config : Config) (Key : Hashtbl.HashedType) : S with type key = Key.t = struct
  module Key = struct
    include Key
    let (=) = Key.equal
  end
  open Config
  (* Branching factor of the structure *)
  let chunk = 1 lsl shift_step
  (* Mask of shift_step `1` bits *)
  let mask = pred chunk

  type key = Key.t
  let hash = Key.hash

  type 'a t =
    | Empty
    | Leaf of int * key * 'a
    | HashCollision of int * (key * 'a) list
    | BitmapIndexedNode of int * 'a t array
    | ArrayNode of int * 'a t array

  let empty = Empty

  let leaf h k v = Leaf (h, k, v)

  let node_children = function
    | BitmapIndexedNode (_, base) -> base
    | ArrayNode (_, children) -> children
    | _ -> failwith "children"

  let singleton k v = Leaf (hash k, k, v)

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

  let rec cardinal = function
    | Empty -> 0
    | Leaf (_, _, _) -> 1
    | HashCollision (_, pairs) -> List.length pairs
    | BitmapIndexedNode (_, base) ->
      Array.fold_left (fun acc child -> acc + cardinal child) 0 base
    | ArrayNode (_, children) ->
      Array.fold_left (fun acc child -> acc + cardinal child) 0 children

  let length = cardinal

  let is_tip_node = function
    | Empty | Leaf (_, _, _) | HashCollision (_, _) -> true
    | _ -> false

  let node_hash = function
    | Leaf (h, _, _) -> h
    | HashCollision (h, _) -> h
    | _ -> failwith "node_hash"

  let hash_fragment shift h = (h asr shift) land mask

  let option default f = function
    | None -> default
    | Some x -> f x

  let remove tab ix =
    let tab' = Array.make (Array.length tab - 1) Empty in<commit_msg>small style change<commit_after>open BitUtils
open Monomorphic.Int

module type Config = sig
  (* Number of bits taken from the hashed key at every step *)
  val shift_step : int

  (* Maximum size of a BitmapIndexedNode*)
  val bmnode_max : int

  (* Minimum size of an ArrayNode (must be lesser than bmnode_max) *)
  val arrnode_min : int
end

module StdConfig : Config = struct
  let shift_step = 5
  let bmnode_max = 16
  let arrnode_min = 8
end

module StdConfig32 : Config = struct
  let shift_step = 4
  let bmnode_max = 8
  let arrnode_min = 4
end

module type S = sig
  type key
  type 'a t

  val empty : 'a t
  val is_empty : 'a t -> bool
  val singleton : key -> 'a -> 'a t
  val cardinal : 'a t -> int
  val length : 'a t -> int

  val alter : key -> ('a option -> 'a option) -> 'a t -> 'a t
  val add : key -> 'a -> 'a t -> 'a t
  val add_carry : key -> 'a -> 'a t -> 'a t * 'a option
  val remove : key -> 'a t -> 'a t
  val extract : key -> 'a t -> 'a * 'a t
  val update : key -> ('a -> 'a option) -> 'a t -> 'a t
  val modify : key -> ('a -> 'a) -> 'a t -> 'a t
  val modify_def : 'a -> key -> ('a -> 'a) -> 'a t -> 'a t

  val find_exn : key -> 'a t -> 'a
  val mem : key -> 'a t -> bool
  val choose : 'a t -> key * 'a
  val pop : 'a t -> (key * 'a) * 'a t

  val keys : 'a t -> key list
  val values : 'a t -> 'a list
  val bindings : 'a t -> (key * 'a) list
  val iter : (key -> 'a -> unit) -> 'a t -> unit
  val map : ('a -> 'b) -> 'a t -> 'b t
  val mapi : (key -> 'a -> 'b) -> 'a t -> 'b t
  val filterv : ('a -> bool) -> 'a t -> 'a t
  val filter : (key -> 'a -> bool) -> 'a t -> 'a t
  val filter_map : (key -> 'a -> 'b option) -> 'a t -> 'b t
  val foldv : ('a -> 'b -> 'b) -> 'a t -> 'b -> 'b
  val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b

  val for_all : (key -> 'a -> bool) -> 'a t -> bool
  val exists : (key -> 'a -> bool) -> 'a t -> bool
  val partition : (key -> 'a -> bool) -> 'a t -> 'a t * 'a t

  val intersecti : (key -> 'a -> 'b -> 'c) -> 'a t -> 'b t -> 'c t
  val intersect : ('a -> 'b -> 'c) -> 'a t -> 'b t -> 'c t
  val merge :
    (key -> 'a option -> 'b option -> 'c option) ->
    'a t -> 'b t -> 'c t
  val union : 'a t -> 'a t -> 'a t
  val union_f : (key -> 'a -> 'a -> 'a) -> 'a t -> 'a t -> 'a t

  module Import : sig
    module type Foldable = sig
      type key
      type 'a t
      val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
    end

    module Make (M : Foldable with type key = key) : sig
      val add_from : 'a M.t -> 'a t -> 'a t
      val from : 'a M.t -> 'a t
    end

    module AssocList : sig
      val add_from : (key * 'a) list -> 'a t -> 'a t
      val from : (key * 'a) list -> 'a t
    end
  end

  module ExceptionLess : sig
    val extract : key -> 'a t -> 'a option * 'a t
    val update : key -> ('a -> 'a option) -> 'a t -> 'a t
    val modify : key -> ('a -> 'a) -> 'a t -> 'a t
    val find : key -> 'a t -> 'a option
    val choose : 'a t -> (key * 'a) option
  end

  module Infix : sig
    val ( --> ) : 'a t -> key -> 'a
    val ( <-- ) : 'a t -> key * 'a -> 'a t
  end
end

module Make (Config : Config) (Key : Hashtbl.HashedType) : S with type key = Key.t = struct
  module Key = struct
    include Key
    let (=) = Key.equal
  end
  open Config
  (* Branching factor of the structure *)
  let chunk = 1 lsl shift_step
  (* Mask of shift_step `1` bits *)
  let mask = pred chunk

  type key = Key.t
  let hash = Key.hash

  type 'a t =
    | Empty
    | Leaf of int * key * 'a
    | HashCollision of int * (key * 'a) list
    | BitmapIndexedNode of int * 'a t array
    | ArrayNode of int * 'a t array

  let empty = Empty

  let leaf h k v = Leaf (h, k, v)

  let node_children = function
    | BitmapIndexedNode (_, base) -> base
    | ArrayNode (_, children) -> children
    | _ -> failwith "children"

  let singleton k v = Leaf (hash k, k, v)

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

  let rec cardinal = function
    | Empty -> 0
    | Leaf (_, _, _) -> 1
    | HashCollision (_, pairs) -> List.length pairs
    | BitmapIndexedNode (_, base) ->
      Array.fold_left (fun acc child -> acc + cardinal child) 0 base
    | ArrayNode (_, children) ->
      Array.fold_left (fun acc child -> acc + cardinal child) 0 children

  let length = cardinal

  let is_tip_node = function
    | Empty | Leaf (_, _, _) | HashCollision (_, _) -> true
    | _ -> false

  let node_hash = function
    | Leaf (h, _, _)
    | HashCollision (h, _) -> h
    | _ -> failwith "node_hash"

  let hash_fragment shift h = (h asr shift) land mask

  let option default f = function
    | None -> default
    | Some x -> f x

  let remove tab ix =
    let tab' = Array.make (Array.length tab - 1) Empty in
{'fraction_non_alphanumeric': 0.16638406104281475, 'fraction_numerical': 0.003391267486222976, 'mean_word_length': 1.9036923076923078, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 322, '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': '3334550', 'n_tokens_mistral': 3615, 'n_tokens_neox': 3468, 'n_words': 1422}