gadt_reflection.ml

open Printf

type (_,_) equality = Refl : ('a,'a) equality

module EqTag : sig
  type 'a t
  val create : unit -> 'a t
  val equal : 'a t -> 'b t -> ('a, 'b) equality option
end = struct
  type 'a tag = ..;;
  type 'a polymorphic_equality = {
    eq : 'b . ('a tag -> 'b tag -> ('a, 'b) equality option);
  }
  type 'a t = 'a tag * 'a polymorphic_equality
  let create (type s) () =
    let module M = struct
      type _ tag += Tag : s tag
      let eq : type a b . a tag -> b tag -> (a, b) equality option =
        fun a b ->
          match a, b with
          | Tag, Tag -> Some Refl
          | _, _ -> None
    end in
    M.Tag, {eq = M.eq}
  let equal (tag1, {eq}) (tag2, _) = eq tag1 tag2
end

module Type = struct
  type _ t =
    | Bool : bool t
    | Int : int t
    | Sum : ('s, 'v) sum -> 's t

  and ('s, 'v) sum = {
    sum_tag : 's EqTag.t;
    sum_name : string;
    sum_match : 's -> 's match_result;
    sum_cons : ('s, 'v) con_list;
  }

  and 's match_result =
      Res : ('s, 'con) con * 'con value_list -> 's match_result

  and _ value_list =
    | VNil : unit value_list
    | VCons : 'a * 'b value_list -> ('a -> 'b) value_list

  and 'a type_list =
    | TNil : unit type_list
    | TCons : 'a t * 'b type_list -> ('a -> 'b) type_list

  and ('s, _) con_list =
    | CNil : ('s, unit) con_list
    | CCons : ('s, 'a) con * ('s, 'b) con_list -> ('s, 'a -> 'b) con_list

  and ('s, 'con) con = {
    con_tag : 'con EqTag.t;
    con_name : string;
    con_args : 'con type_list;
    con_build : 'con value_list -> 's option;
  }

  let cast : type a b . (a, b) equality -> b -> a = fun Refl v -> v
  let cast_list : type a b . (a, b) equality -> b value_list -> a value_list = fun Refl v -> v

  let rec equal : type a b . a t -> b t -> (a, b) equality option =
    fun a b ->
    match a, b with
    | Bool, Bool -> Some Refl
    | Int, Int -> Some Refl
    | Sum x, Sum y -> equal_sum x y
    | _, _ -> None

  and equal_sum : type a b x y. (a, x) sum -> (b, y) sum -> (a, b) equality option =
    fun a b ->
      EqTag.equal a.sum_tag b.sum_tag
end

let rec equal : type a . a Type.t -> a -> a -> bool =
  fun ty x y ->
    match ty with
    | Type.Bool -> x = y
    | Type.Int -> x = y
    | Type.Sum sum -> equal_sum sum x y

and equal_sum : type s t . (s, t) Type.sum -> s -> s -> bool =
  let open Type in
  fun sum x y ->
    let Res (xcon, xvalues) = sum.sum_match x in
    let Res (ycon, yvalues) = sum.sum_match y in
    match EqTag.equal xcon.con_tag ycon.con_tag with
    | Some eqwit ->
      equal_list xcon.con_args xvalues (Type.cast_list eqwit yvalues)
    | None -> false

and equal_list : type t . t Type.type_list -> t Type.value_list -> t Type.value_list -> bool =
  fun ty a b ->
    match ty, a, b with
    | Type.TNil, Type.VNil, Type.VNil -> true
    | Type.TCons (ty, trest), Type.VCons (a', arest), Type.VCons (b', brest) ->
      equal ty a' b' && equal_list trest arest brest
    | _, _, _ -> false

let rec print : type a . a Type.t -> a -> unit =
  let open Type in
  fun ty x ->
    match ty with
    | Bool -> printf "%b" x
    | Int -> print_int x
    | Sum sum -> print_sum sum x

and print_sum : type s v . (s, v) Type.sum -> s -> unit =
  let open Type in
  fun sum x ->
    let Res (con, values) = sum.sum_match x in
    let rec loop : type t . bool -> t type_list -> t value_list -> unit =
      fun seen_elt ts vs ->
        match ts, vs with
        | TNil, VNil ->
          if seen_elt then print_char ')'
        | TCons (ty, trest), VCons (v, vrest) ->
          if not seen_elt then print_string " ("
          else print_string ", ";
          print ty v;
          loop true trest vrest
        | _ -> assert false in
    print_string con.con_name;
    loop false con.con_args values

type t =
  | Foo
  | Bar of int * bool
  | Rec of t

let t_tag : t EqTag.t = EqTag.create ()

let rec t_type_sum =
  Type.{
    sum_tag = t_tag;
    sum_name = "t";
    sum_match = (function
      | Foo -> Res (foo_con, VNil)
      | Bar (i, b) -> Res (bar_con, VCons (i, VCons (b, VNil)))
      | Rec t -> Res (rec_con, VCons (t, VNil)));
    sum_cons = CCons (foo_con, CCons (bar_con, CCons (rec_con, CNil)));
  }

and t_type = Type.Sum t_type_sum

and foo_con = Type.{
    con_tag = EqTag.create ();
    con_name = "Foo";
    con_args = TNil;
    con_build = (function
        | VNil -> Some Foo
        | _ -> None);
  }

and bar_con = Type.{
    con_tag = EqTag.create ();
    con_name = "Bar";
    con_args = TCons (Type.Int, TCons (Type.Bool, TNil));
    con_build = (function
        | VCons (i, (VCons (b, VNil))) -> Some (Bar (i, b))
        | _ -> None);
  }

and rec_con = Type.{
    con_tag = EqTag.create ();
    con_name = "Rec";
    con_args = TCons (t_type, TNil);
    con_build = (function
        | VCons (t, VNil) -> Some (Rec t)
        | _ -> None);
  }