corresponding.ml

(* Sigh, GADTs are allergic to uninhabited types with params  *)
type nullary = { nullary : 'a . 'a }
type 'a unary = { unary : 'a . 'a }
type ('a, 'b) binary = { binary : 'a . 'a }

type symbol = string

module Env = Map.Make (struct
    type t = symbol
    let compare = compare
  end)

type binary_op = Add
type cmp_op = Lt

type _ term =
  | Unit : nullary term
  | Bool : bool -> nullary term
  | Int : int -> nullary term
  | BinOp : binary_op * 'a term * 'b term -> ('a, 'b) binary term
  | Cmp : cmp_op * 'a term * 'b term -> ('a, 'b) binary term
  | Fun : symbol * ty * 'a term -> 'a unary term
  | App : 'a term * 'b term -> ('a, 'b) binary term
  | Var : symbol -> nullary term

and ty =
  | TyUnit
  | TyBool
  | TyInt
  | TyPair of ty * ty
  | TyFun of ty * ty

module Annot = struct
  type ('a, _) t =
    | Nullary : 'a -> ('a, nullary) t
    | Unary : 'a * ('a, 'b) t -> ('a, 'b unary) t
    | Binary : 'a * ('a, 'b) t * ('a, 'c) t -> ('a, ('b, 'c) binary) t

  let elt : type b . ('a, b) t -> 'a = function
    | Nullary x -> x
    | Unary (x, _) -> x
    | Binary (x, _, _) -> x
end

type typed_term = Typed : 'a term * (ty, 'a) Annot.t -> typed_term

type type_error =
  | Unbound of symbol
  | Mismatch of ty * ty
  | NonFunctionApplied of ty

exception TypeError of type_error

let type_error err = raise (TypeError err)

let rec match_types a b =
  match a, b with
  | TyUnit, TyUnit
  | TyBool, TyBool
  | TyInt, TyInt -> ()
  | TyFun (x1, y1), TyFun (x2, y2) ->
    match_types x1 x2;
    match_types y1 y2
  | _, _ -> type_error (Mismatch (a, b))

let rec type_annot : type a . ty Env.t -> a term -> (ty, a) Annot.t =
  fun env term ->
    match term with
    | Unit -> Annot.Nullary TyUnit
    | Bool _ -> Annot.Nullary TyBool
    | Int _ -> Annot.Nullary TyInt
    | BinOp (_, a, b) ->
      let a_annot = type_annot env a in
      let b_annot = type_annot env b in
      match_types TyInt (Annot.elt a_annot);
      match_types TyInt (Annot.elt b_annot);
      Annot.Binary (TyInt, a_annot, b_annot)
    | Cmp (_, a, b) ->
      let a_annot = type_annot env a in
      let b_annot = type_annot env b in
      match_types (Annot.elt a_annot) (Annot.elt b_annot);
      Annot.Binary (TyBool, a_annot, b_annot)
    | Fun (arg, arg_ty, body) ->
      let body_annot = type_annot (Env.add arg arg_ty env) body in
      Annot.Unary (TyFun (arg_ty, Annot.elt body_annot),
                   body_annot)
    | App (f, arg) ->
      let f_annot = type_annot env f in
      let arg_annot = type_annot env arg in
      let arg_ty = Annot.elt arg_annot in
      (match Annot.elt f_annot with
       | TyFun (formal_ty, result_ty) ->
         match_types formal_ty arg_ty;
         Annot.Binary (result_ty, f_annot, arg_annot)
       | t -> type_error (NonFunctionApplied t))
    | Var name ->
      Annot.Nullary (try Env.find name env
                     with Not_found -> type_error (Unbound name))

let typed env term =
  Typed (term, type_annot env term)