ADT to GADT

main.ml

module Eq =
struct
  type (_, _) refl = Refl : ('a, 'a) refl

  let symm
    : type m n. (m, n) refl -> (n, m) refl
    = fun Refl -> Refl

  let trans
    : type m n p. (m, n) refl -> (n, p) refl -> (m, p) refl
    = fun Refl Refl -> Refl

  module Lift (T : sig type 'a t end) =
  struct
    let lift
      : type a b. (a, b) refl -> (a T.t, b T.t) refl =
      fun Refl -> Refl
  end

  let cast
    : type a b. (a, b) refl -> a -> b
    = fun Refl x -> x

  let imply
    : type a b c d. (a, b) refl -> (c, d) refl option -> (a -> c, b -> d) refl option
    = fun Refl -> function
      | Some Refl -> Some Refl
      | None -> None

  let shift
    : type a b c d. (a, b) refl -> (c, d) refl option -> (a * c, b * d) refl option
    = fun Refl -> function
      | Some Refl -> Some Refl
      | None -> None
end

module A =
struct
  type var =
    | Z
    | S of var

  type typ =
    | Int
    | Arr of typ * typ

  type exp =
    | Con of int
    | Var of var
    | Abs of typ * exp
    | App of exp * exp
    | Add of exp * exp
end

module B =
struct
  type ('e, 'a) var =
    | Z : ('e * 'a, 'a) var
    | S : ('e, 'a) var -> ('e * 'b, 'a) var

  type 'a typ =
    | Int : int typ
    | Arr : 'a typ * 'b typ -> ('a -> 'b) typ

  type ('e, 'a) exp =
    | Con : int -> ('e, int) exp
    | Add : ('e, int) exp * ('e, int) exp -> ('e, int) exp
    | Var : ('e, 'a) var -> ('e, 'a) exp
    | Abs : 'a typ * ('e * 'a, 'b) exp -> ('e, 'a -> 'b) exp
    | App : ('e, 'a -> 'b) exp * ('e, 'a) exp -> ('e, 'b) exp

  type 'e env =
    | Nil : unit env
    | Ext : 'e env * 'a typ -> ('e * 'a) env

  let rec get
    : type e a. (e, a) var -> e -> a
    = fun v e -> match v, e with
      | Z, (_, x) -> x
      | S n, (r, _) -> get n r

  let rec run
    : type e a. (e, a) exp -> e -> a
    = fun x e -> match x with
      | Con i -> i
      | Var x -> get x e
      | Abs (_, r) -> (fun v -> run r (e, v))
      | App (f, a) ->
        (run f e) (run a e)
      | Add (l, r) ->
        (run l e) + (run r e)
end

let option_bind
  : ('a -> 'b option) -> 'a option -> 'b option
  = fun f -> function Some x -> f x | None -> None

let ( >>= )
  : type a b c d. (a, b) Eq.refl option -> ((a, b) Eq.refl -> (c, d) Eq.refl option) -> (a -> c, b -> d) Eq.refl option
  = fun x f -> match x with
    | Some (Eq.Refl as x) -> Eq.imply x (f x)
    | None -> None

let ( >|= )
  : type a b c d. (a, b) Eq.refl option -> ((a, b) Eq.refl -> (c, d) Eq.refl option) -> (a * c, b * d) Eq.refl option
  = fun x f -> match x with
    | Some (Eq.Refl as x) -> Eq.shift x (f x)
    | None -> None

let rec eql_typ
  : type a b. a B.typ -> b B.typ -> (a, b) Eq.refl option
  = fun a b -> match a, b with
    | B.Int, B.Int -> Some Eq.Refl
    | B.Arr (a, a'), B.Arr (b, b') ->
      eql_typ a b >>= fun Eq.Refl -> eql_typ a' b'
    | _, _ -> None

let rec eql_env
  : type a b. a B.env -> b B.env -> (a, b) Eq.refl option
  = fun a b -> match a, b with
    | B.Nil, B.Nil -> Some Eq.Refl
    | B.Ext (a, at), B.Ext (b, bt) ->
      eql_env a b >|= fun Eq.Refl -> eql_typ at bt
    | _ -> None

type wtyp = Wtyp : 'a B.typ -> wtyp

let rec typ
  : A.typ -> wtyp
  = function
    | A.Int -> Wtyp B.Int
    | A.Arr (a, b) ->
      let Wtyp a = typ a in
      let Wtyp b = typ b in
      Wtyp (B.Arr (a, b))

type wenv = Wenv : 'a B.env -> wenv

let rec env
  : A.typ list -> wenv
  = function
    | [] -> Wenv B.Nil
    | x :: r ->
      let Wtyp x = typ x in
      let Wenv r = env r in
      Wenv (B.Ext (r, x))

type wvar = Wvar : ('e, 'a) B.var * 'e B.env * 'a B.typ -> wvar

let rec var
  : A.var -> A.typ list -> wvar
  = fun v e -> match v, e with
    | A.S x, t :: r ->
      let Wvar (x', r', tr) = var x r in
      let Wtyp t' = typ t in

      Wvar (B.S x', B.Ext (r', t'), tr)
    | A.Z, t :: r ->
      let Wtyp t = typ t in
      let Wenv r = env r in

      Wvar (B.Z, B.Ext (r, t), t)
    | _, _ -> assert false

type wexp = Wexp : ('e, 'a) B.exp * 'e B.env * 'a B.typ -> wexp

let rec exp
  : A.exp -> A.typ list -> wexp
  = fun e g -> match e with
    | A.Con i ->
      let Wenv g' = env g in
      Wexp (B.Con i, g', B.Int)
    | A.Var x ->
      let Wvar (x', g', t') = var x g in
      Wexp (B.Var x', g', t')
    | A.Abs (t, e) ->
      let Wtyp t' = typ t in
      (match exp e (t :: g) with
       | Wexp (e', B.Ext (g', t''), tr) ->
         (match eql_typ t' t'' with
          | Some Eq.Refl -> Wexp (B.Abs (t', e'), g', B.Arr (t', tr))
          | None -> assert false)
       | _ -> assert false)
    | A.App (f, a) ->
      (match exp f g, exp a g with
       | Wexp (f', g', (B.Arr (t', u'))),
         Wexp (a', g'', t'') ->
         (match eql_env g' g'', eql_typ t' t'' with
          | Some Eq.Refl, Some Eq.Refl ->
            Wexp (B.App (f', a'), g', u')
          | _, _ -> assert false)
       | _, _ -> assert false)
    | A.Add (l, r) ->
      (match exp l g, exp r g with
       | Wexp (l', g', B.Int),
         Wexp (r', g'', B.Int) ->
         (match eql_env g' g'' with
          | Some Eq.Refl -> Wexp (B.Add (l', r'), g', B.Int)
          | None -> assert false)
       | _, _ -> assert false)

let run
  : A.exp -> int
  = fun m ->
  match exp m [] with
  | Wexp (m', B.Nil, B.Int) -> B.run m' ()
  | _ -> assert false

let () = assert (run (A.Con 4) = 4)