gadt-print.ml

module Type = struct
  type _ t =
    | Bool : bool t
    | Int : int t
    | Float : float t
    | Tuple : (_, 't) agg -> 't t
    | Record : (_, 'r) agg * string array * string -> 'r t
    | Fun : 'a t * 'b t -> ('a -> 'b) t
    | Sum : 's sum -> 's t

  and (_, _) agg =
    | Nil : ('t, 't) agg
    | Cons : 'a t * ('t -> 'a) * ('r, 't) agg -> (('a -> 'r), 't) agg

  and 's sum = {
    sum_name : string;
    ctors : 's ctor array;
  }

  and _ ctor = Ctor : (_, 's) ctor_ -> 's ctor

  and ('arg, 's) ctor_ = {
    ctor_name : string;
    arg_type : 'arg t;
    arg_proj : 'a . (('arg -> 'a) -> 's -> 'a option);
  }

  let print ty =
    let rec pr : type a . bool -> a t -> unit =
      fun func_parens -> function
        | Bool -> print_string "bool"
        | Int -> print_string "int"
        | Float -> print_string "float"
        | Record (_, _, name) -> print_string name
        | Tuple Nil -> print_string "()"
        | Tuple (Cons (a, _, rest)) ->
          let rec loop : type a b . (b, a) agg -> unit = function
            | Nil -> ()
            | Cons (a, _, rest) ->
              print_string ", ";
              pr false a;
              loop rest in
          print_char '(';
          pr false a;
          loop rest;
          print_char ')'
        | Fun (a, b) ->
          if func_parens then print_char '(';
          pr true a;
          print_string " -> ";
          pr false b;
          if func_parens then print_char ')'
        | Sum sum -> print_string sum.sum_name in
    pr false ty
end

module Term = struct
  type _ t =
    | Bool : bool -> bool t
    | Int : int -> int t
    | Float : float -> float t
    | Agg : 't agg -> 't t
    | Elt : 't t * ('t -> 'a) -> 'a t

  and 't agg = Agg_ : 'f * ('f, 't) agg_part -> 't agg

  and (_, _) agg_part =
    | Nil : ('t, 't) agg_part
    | Cons : 'a t * ('r, 't) agg_part -> (('a -> 'r), 't) agg_part

  let rec eval : type a . a t -> a = function
    | Int n -> n
    | Bool b -> b
    | Float f -> f
    | Elt (t, f) -> f (eval t)
    | Agg (Agg_ (build, parts)) ->
      let rec loop : type v . v -> (v, a) agg_part -> a =
        fun b parts -> match parts with
          | Nil -> b
          | Cons (term, rest) -> loop (b (eval term)) rest
      in
      loop build parts
end

open Printf

let rec print : type a . a -> a Type.t -> unit =
  fun value -> function
  | Type.Int -> print_int value
  | Type.Bool -> printf "%b" value
  | Type.Float -> printf "%F" value
  | Type.Tuple Type.Nil -> print_string "()"
  | Type.Tuple (Type.Cons (elt_ty, getter, rest)) ->
    let rec loop : type a b . a -> (b, a) Type.agg -> unit =
      fun tuple -> function
        | Type.Nil -> ()
        | Type.Cons (elt_ty, getter, rest) ->
          print_string ", ";
          print (getter tuple) elt_ty;
          loop tuple rest in
    print_char '(';
    print (getter value) elt_ty;
    loop value rest;
    print_char ')'
  | Type.Record (Type.Nil, _, _) -> print_string "{}"
  | Type.Record (Type.Cons (elt_ty, getter, rest), field_names, _) ->
    let rec loop : type a b . a -> int -> (b, a) Type.agg -> unit =
      fun tuple n -> function
        | Type.Nil -> ()
        | Type.Cons (elt_ty, getter, rest) ->
          printf ", %s: " field_names.(n);
          print (getter tuple) elt_ty;
          loop tuple (n + 1) rest in
    printf "{%s: " field_names.(0);
    print (getter value) elt_ty;
    loop value 1 rest;
    print_string "}"
  | Type.Fun _ -> print_string "<fun>"
  | Type.Sum s ->
    let rec loop i bound =
      if i = bound then assert false
      else
        let Type.Ctor ctor = s.Type.ctors.(i) in
        let f arg_value =
          printf "%s " ctor.Type.ctor_name;
          print arg_value ctor.Type.arg_type in
        match ctor.Type.arg_proj f value with
        | Some () -> ()
        | None -> loop (i + 1) bound in
    loop 0 (Array.length s.Type.ctors)

module Env = struct
  type entry = Entry : 'a * 'a Type.t -> entry
  type t = (string, entry) Hashtbl.t

  let find env name =
    Hashtbl.find env name
  let add env name value ty =
    Hashtbl.add env name (Entry (value, ty))
end

type r = {
  x : int;
  y : bool * float;
}

let tuple_ty =
  Type.Tuple (Type.Cons (Type.Bool,
                         fst,
                         Type.Cons (Type.Float, snd, Type.Nil)))

let record_ty =
  Type.Record (Type.Cons (Type.Int,
                          (fun r -> r.x),
                          Type.Cons (tuple_ty,
                                     (fun r -> r.y),
                                     Type.Nil)),
               [|"x"; "y"|],
               "r")

type s =
  | B of bool
  | I of int

let s_proj_b f = function | B b -> Some (f b) | _ -> None
let s_proj_i f = function | I i -> Some (f i) | _ -> None

let sum_ty = Type.(Sum {
    sum_name = "s";
    ctors = [|
      Ctor {ctor_name = "B"; arg_type = Bool; arg_proj = s_proj_b};
      Ctor {ctor_name = "I"; arg_type = Int; arg_proj = s_proj_i};
    |];
  })