Source file table.ml

(**************************************************************************)
(*                                                                        *)
(*                 ACG development toolkit                                *)
(*                                                                        *)
(*                  Copyright 2008-2024 INRIA                             *)
(*                                                                        *)
(*  More information on "https://acg.loria.fr/"                     *)
(*  License: CeCILL, see the LICENSE file or "http://www.cecill.info"     *)
(*  Authors: see the AUTHORS file                                         *)
(*                                                                        *)
(*                                                                        *)
(*                                                                        *)
(*                                                                        *)
(*                                                                        *)
(**************************************************************************)

module type BASE = sig
  val b : int
end

module type TABLE = sig
  type 'a t
  type key

  val empty : 'a t
  val add : ?overwrite:bool -> key -> 'a -> 'a t -> ('a t) option
  val find : key -> 'a t -> 'a option
  val fold : (key -> 'a -> 'b -> 'b) -> 'b -> 'a t -> 'b
  val iter : (key -> 'a -> unit) -> 'a t -> unit

  val pp :
    ?sep:
      ( (Format.formatter -> key * 'a -> unit) -> key * 'a -> unit,
        Format.formatter,
        unit,
        unit,
        unit,
        (Format.formatter -> key * 'a -> unit) -> key * 'a -> unit )
      format6 ->
    (Format.formatter -> key -> 'a -> unit) ->
    Format.formatter ->
    'a t ->
    unit
end

module Make_table (Base : BASE) = struct
  type 'a t = Nil | T of ('a option * 'a t) array
  type key = int

  let create () = T (Array.make Base.b (None, Nil))
  let empty = Nil

  let add ?(overwrite = false) n attr table =
    let rec insert1 n table =
      match table with
      | Nil -> insert1 n (create ())
      | T ar ->
          let r, i = (n / Base.b, n mod Base.b) in
          let a, tb = ar.(i) in
          if r = 0 then (
            match (a, overwrite) with
            | None, _ ->
                ar.(i) <- (Some attr, tb);
                Some (T ar)
            | Some _, false -> None
            | Some _, true ->
                ar.(i) <- (Some attr, tb);
                Some (T ar))
          else (
            Option.map
              (fun o ->
                ar.(i) <- (a, o);
                T ar)
              (insert1 r tb))
    in
    insert1 n table

  let rec find n table =
    match table with
    | Nil -> None
    | T ar ->
        let r, i = (n / Base.b, n mod Base.b) in
        let a, tb = ar.(i) in
        if r = 0 then a
        else find r tb

  let fold f acc table =
    let rec fold_aux q acc = function
      | Nil -> acc
      | T ar ->
          let _, new_acc =
            Array.fold_left
              (fun (i, acc) -> function
                | Some v, _ -> (i + 1, f ((q * Base.b) + i) v acc)
                | None, _ -> (i + 1, acc))
              (0, acc) ar
          in
          snd
            (Array.fold_left
               (fun (i, acc) (_, t) -> (i + 1, fold_aux (q + 1) acc t))
               (0, new_acc) ar)
    in
    fold_aux 0 acc table

  let iter f table =
    let rec iteri_aux q f table =
      match table with
      | Nil -> ()
      | T ar ->
          let () =
            Array.iteri
              (fun i (value, _t) ->
                match value with Some v -> f ((q * Base.b) + i) v | None -> ())
              ar
          in
          Array.iteri (fun q (_value, t) -> iteri_aux (q + 1) f t) ar
    in
    iteri_aux 0 f table

  let pp ?(sep = format_of_string "@,") ppf m t =
    let l_pp m (k, v) = ppf m k v in
    let first = ref true in
    iter
      (fun i v ->
        if !first then
          let () = first := false in
          ppf m i v
        else Format.fprintf m (sep ^^ "%a") l_pp (i, v))
      t
end