(*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

(*****************************************************************************)
(* The heap shared across all the processes.
 *
 * The Heap is not exposed directly to the user (cf shared.mli),
 * because we don't want to mix values of different types. Instead, we want
 * to use a functor.
 *)
(*****************************************************************************)

type config = {
  heap_size: int;
  dep_table_pow: int;
  hash_table_pow: int;
  log_level: int;
}

exception Out_of_shared_memory

exception Hash_table_full

exception Dep_table_full

exception Heap_full

exception Sql_assertion_failure of int

exception C_assertion_failure of string

(*****************************************************************************)
(* Initializes the shared memory. Must be called before forking! *)
(*****************************************************************************)

val init : config -> unit

(*****************************************************************************)
(* Connect a slave to the shared heap *)
(*****************************************************************************)

val connect : unit -> unit

(*****************************************************************************)
(* The shared memory garbage collector. It must be called every time we
 * free data (cf hh_shared.c for the underlying C implementation).
 *)
(*****************************************************************************)

val collect : [ `gentle | `aggressive | `always_TEST ] -> unit

(*****************************************************************************)
(* Must be called after the initialization of the hack server is over.
 * (cf serverInit.ml).
 *)
(*****************************************************************************)

val init_done : unit -> unit

(*****************************************************************************)
(* Serializes the dependency table and writes it to a file *)
(*****************************************************************************)
val save_dep_table_sqlite : string -> string -> int

(*****************************************************************************)
(* Loads the dependency table by reading from a file *)
(*****************************************************************************)
val load_dep_table_sqlite : string -> bool -> int

(*****************************************************************************)
(* Serializes & loads the hash table directly into memory *)
(*****************************************************************************)

val save_table : string -> unit

val load_table : string -> unit

(*****************************************************************************)
(* Serializes the hash table to sqlite *)
(*****************************************************************************)

val save_table_sqlite : string -> int

val save_table_keys_sqlite : string -> string array -> int

(*****************************************************************************)
(* Loads the hash table by reading from a file *)
(*****************************************************************************)

val load_table_sqlite : string -> bool -> int

(*****************************************************************************)
(* Cleans up the artifacts generated by SQL *)
(*****************************************************************************)
val cleanup_sqlite : unit -> unit

(*****************************************************************************)
(* The size of the dynamically allocated shared memory section *)
(*****************************************************************************)
val heap_size : unit -> int

(*****************************************************************************)
(* Part of the heap not reachable from hashtable entries. *)
(*****************************************************************************)
val wasted_heap_size : unit -> int

(*****************************************************************************)
(* Stats of the statically sized hash / dep tables *)
(*****************************************************************************)

type table_stats = {
  nonempty_slots: int;
  used_slots: int;
  slots: int;
}

val dep_stats : unit -> table_stats

val hash_stats : unit -> table_stats

val is_heap_overflow : unit -> bool

(*****************************************************************************)
(* Cache invalidation. *)
(*****************************************************************************)

val invalidate_caches : unit -> unit

(* Size of value in GC heap *)
val value_size : Obj.t -> int

(*****************************************************************************)
(* The signatures of shared memory hashtables
 *
 * Use NoCache/WithCache if you want caching or not. If you do, bear in mind
 * that the cache must be maintained by the caller, so you will have to
 * invalidate the caches yourself.
 *)
(*****************************************************************************)

(*****************************************************************************)
(* The interfaces for keys and values of shared memory tables *)
(*****************************************************************************)

module type KeyType = sig
  type t

  val to_string : t -> string

  val compare : t -> t -> int
end

module type ValueType = sig
  type t

  val prefix : Prefix.t

  val description : string
end

(*****************************************************************************)
(* The common key and value type logic for all shared-memory tables *)
(*****************************************************************************)

module TableTypes : sig
  module type S = sig
    type key

    type value

    module KeySet : Set.S with type elt = key

    module KeyMap : Hack_collections.MyMap.S with type key = key
  end

  module Make : functor (KeyType : KeyType) (Value : ValueType) ->
    S
      with type value = Value.t
       and type key = KeyType.t
       and module KeySet = Set.Make(KeyType)
       and module KeyMap = Hack_collections.MyMap.Make(KeyType)
end

(*****************************************************************************)
(* A shared memory table with no process-local caching of reads *)
(*****************************************************************************)
module NoCache : sig
  module type S = sig
    include TableTypes.S

    (* Add a value to the table. Safe for concurrent writes - the first writer wins, later values
       are discarded. *)
    val add : key -> value -> unit

    val add_old : key -> value -> unit

    (* Api to read and remove from the table *)
    val get : key -> value option

    val get_exn : key -> value

    val mem : key -> bool

    val get_batch : KeySet.t -> value option KeyMap.t

    val remove : key -> unit

    val remove_batch : KeySet.t -> unit

    (* Api to read and remove old data from the table, which lives in a separate hash map. Used in
       situations where we want to know what has changed, for example dependency-tracked tables. *)
    val get_old : key -> value option

    val mem_old : key -> bool

    val get_old_batch : KeySet.t -> value option KeyMap.t

    val remove_old_batch : KeySet.t -> unit

    (* Move keys between the current view of the table and the old-values table *)
    val oldify_batch : KeySet.t -> unit

    val revive_batch : KeySet.t -> unit
  end

  module Make : functor (KeyType : KeyType) (Value : ValueType) ->
    S
      with type value = Value.t
       and type key = KeyType.t
       and module KeySet = Set.Make(KeyType)
       and module KeyMap = Hack_collections.MyMap.Make(KeyType)
end

(*****************************************************************************)
(* A shared memory table with process-local caches of reads. We use the cache
 * to avoid paying the deserialization cost in duplicate reads within process.
 *)
(*****************************************************************************)
module WithCache : sig
  module type S = sig
    include NoCache.S

    val write_around : key -> value -> unit

    val get_no_cache : key -> value option
  end

  module Make : functor (KeyType : KeyType) (Value : ValueType) ->
    S
      with type value = Value.t
       and type key = KeyType.t
       and module KeySet = Set.Make(KeyType)
       and module KeyMap = Hack_collections.MyMap.Make(KeyType)
end

module type CacheType = sig
  type key

  type value

  val add : key -> value -> unit

  val get : key -> value option

  val remove : key -> unit

  val clear : unit -> unit

  val get_size : unit -> int
end

module LocalCache : functor (KeyType : KeyType) (Value : ValueType) ->
  CacheType with type key = KeyType.t and type value = Value.t

module FirstClass : sig
  module NoCache : sig
    module type S = sig
      type t

      val equal : t -> t -> bool

      type key

      type value

      module KeySet : Set.S with type elt = key

      module KeyMap : Hack_collections.MyMap.S with type key = key

      (* The create function must be run on the main ocaml process, and is not thread-safe. *)
      val create : unit -> t

      (* Add a value to the table. Safe for concurrent writes - the first writer wins, later values
         are discarded. *)
      val add : t -> key -> value -> unit

      (* Api to read and remove from the table *)
      val mem : t -> key -> bool

      val get : t -> key -> value option

      val get_batch : t -> KeySet.t -> value option KeyMap.t

      val remove : t -> key -> unit

      val remove_batch : t -> KeySet.t -> unit

      (* Api to read and remove old data from the table, which lives in a separate hash map. Used in
         situations where we want to know what has changed, for example dependency-tracked
         tables. *)
      val mem_old : t -> key -> bool

      val get_old : t -> key -> value option

      val get_old_batch : t -> KeySet.t -> value option KeyMap.t

      val remove_old_batch : t -> KeySet.t -> unit

      (* Move keys between the current view of the table and the old-values table *)
      val oldify_batch : t -> KeySet.t -> unit

      val revive_batch : t -> KeySet.t -> unit
    end

    module Make : functor (KeyType : KeyType) (Value : ValueType) ->
      S
        with type value = Value.t
         and type key = KeyType.t
         and module KeySet = Set.Make(KeyType)
         and module KeyMap = Hack_collections.MyMap.Make(KeyType)
  end

  module WithCache : sig
    module type S = sig
      include NoCache.S

      val write_around : t -> key -> value -> unit

      val get_no_cache : t -> key -> value option
    end

    module Make : functor (KeyType : KeyType) (Value : ValueType) ->
      S
        with type value = Value.t
         and type key = KeyType.t
         and module KeySet = Set.Make(KeyType)
         and module KeyMap = Hack_collections.MyMap.Make(KeyType)
  end
end

module FirstClassWithKeys : sig
  module type S = sig
    type t

    type key

    type value

    module KeySet : Set.S with type elt = key

    val create : unit -> t

    val keys : t -> key list

    val cleanup : clean_old:bool -> t -> unit

    val of_alist_sequential : (key * value) list -> t

    val of_alist_parallel
      :  map_reduce:(map:((key * value) list -> unit) -> inputs:(key * value) list -> unit -> unit) ->
      (key * value) list ->
      t

    val to_alist : t -> (key * value) list

    val merge_with_alist_sequential : t -> f:(value -> value -> value) -> (key * value) list -> t

    val add_alist_sequential : t -> (key * value) list -> t

    val fold_sequential : t -> init:'a -> f:(key:key -> value:value -> 'a -> 'a) -> 'a

    type map_parallel_state

    val map_parallel_keys
      :  t ->
      map_reduce:
        (initial:map_parallel_state ->
        map:(key list -> map_parallel_state) ->
        reduce:(map_parallel_state -> map_parallel_state -> map_parallel_state) ->
        inputs:key list ->
        unit ->
        map_parallel_state) ->
      f:(key:key -> value:value -> value) ->
      keys:key list ->
      t

    val map_parallel
      :  t ->
      map_reduce:
        (initial:map_parallel_state ->
        map:(key list -> map_parallel_state) ->
        reduce:(map_parallel_state -> map_parallel_state -> map_parallel_state) ->
        inputs:key list ->
        unit ->
        map_parallel_state) ->
      f:(key:key -> value:value -> value) ->
      t

    val oldify_batch : t -> KeySet.t -> t

    val remove_old_batch : t -> KeySet.t -> t

    module ReadOnly : sig
      type t

      val get : t -> cache:bool -> key -> value option

      val get_old : t -> key -> value option

      val mem : t -> key -> bool
    end

    val read_only : t -> ReadOnly.t

    module AddOnly : sig
      type t

      val add : t -> key -> value -> t

      val merge_same_handle_disjoint_keys : smaller:t -> larger:t -> t

      val create_empty : t -> t

      val keys : t -> key list
    end

    val add_only : t -> AddOnly.t

    val from_add_only : AddOnly.t -> t

    module PreserveKeyOnly : sig
      type t

      val get : t -> cache:bool -> key -> value option

      val get_old : t -> key -> value option

      val set_new : t -> cache:bool -> key -> value -> t
    end

    val preserve_key_only : t -> PreserveKeyOnly.t
  end

  module Make : functor (KeyType : KeyType) (Value : ValueType) ->
    S with type value = Value.t and type key = KeyType.t and module KeySet = Set.Make(KeyType)
end
