(************************************************************************)
(*         *      The Rocq Prover / The Rocq Development Team           *)
(*  v      *         Copyright INRIA, CNRS and contributors             *)
(* <O___,, * (see version control and CREDITS file for authors & dates) *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)

(* Type-safe implementation by whitequark *)

(* An extensible variant has an internal representation equivalent
   to the following:

      type constr = {
        name: string,
        id:   int
      }
      type value = (*Object_tag*) constr * v1 * v2...

   and the code generated by the compiler looks like:

      (* type X += Y *)
      let constr_Y = alloc { "Y", %caml_fresh_oo_id () }
      (* match x with Y -> a | _ -> b *)
      if x.0 == constr_Y then a else b

   and the polymorphic comparison function works like:

      let equal = fun (c1, ...) (c2, ...) ->
        c1.id == c2.id

   In every new extension constructor, the name field is a constant
   string and the id field is filled with an unique[1] value returned
   by %caml_fresh_oo_id. Moreover, every value of an extensible variant
   type is allocated as a new block.

   [1]: On 64-bit systems. On 32-bit systems, calling %caml_fresh_oo_id
        2**30 times will result in a wraparound. Note that this does
        not affect soundness because constructors are compared by
        physical equality during matching. See OCaml PR7809 for code
        demonstrating this.

   An extensible variant can be marshalled and unmarshalled, and
   is guaranteed to not be equal to itself after unmarshalling,
   since the id field is filled with another unique value.

   Note that the explanation above is purely informative and we
   do not depend on the exact representation of extensible variants,
   only on the fact that no two constructor representations ever
   alias. In particular, if the definition of constr is replaced with:

      type constr = int

   (where the value is truly unique for every created constructor),
   correctness is preserved.
 *)
type 'a typ = ..

(* Erases the contained type so that the key can be put in a hash table. *)
type boxkey = Box : 'a typ -> boxkey [@@unboxed]

(* Carry the type we just erased with the actual key. *)
type 'a key = 'a typ * boxkey

module EHashtbl = Ephemeron.K1.Make(struct
  type t = boxkey
  let equal = (==)
  let hash  = Hashtbl.hash
end)

type value = { get : 'k. 'k typ -> 'k } [@@unboxed]

let values : value EHashtbl.t =
  EHashtbl.create 1001

let create : type v. v -> v key =
  fun value ->
  let module M = struct
    type _ typ += Typ : v typ

    let get : type k. k typ -> k =
      fun typ ->
      match typ with
      | Typ -> value
      | _ -> assert false

    let boxkey = Box Typ
    let key    = Typ, boxkey
    let value  = { get }
  end in
  EHashtbl.add values M.boxkey M.value;
  M.key

(* Avoid raising Not_found *)
exception InvalidKey
let get (typ, boxkey) =
  try (EHashtbl.find values boxkey).get typ
  with Not_found -> raise InvalidKey

let default (typ, boxkey) default =
  try (EHashtbl.find values boxkey).get typ
  with Not_found -> default

let clean () = EHashtbl.clean values
