onemel_clambda.ml

(* (c) Romain Calascibetta *)

[@@@warning "-42"]
[@@@warning "-44"]
[@@@warning "-45"]

type ('l, 'r) either =
  | L of 'l
  | R of 'r

type ('a, 'b) cpl = T of 'a * 'b

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 arrow
    : 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 cpl
    : type a b c d. (a, b) refl -> (c, d) refl option -> ((a, c) cpl, (b, d) cpl) refl option
    = fun Refl -> function
      | Some Refl -> Some Refl
      | None -> None

  let star
    : 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 either
    : type a b c d. (a, b) refl -> (c, d) refl option -> ((a, c) either, (b, d) either) refl option
    = fun Refl -> function
      | Some Refl -> Some Refl
      | None -> None
end

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

  let rec int_of_var = function
    | Z -> 1
    | S n -> 1 + (int_of_var n)

  let rec var_of_int n =
    if n <= 0
    then invalid_arg "A.var_of_int: invalid number (must be upper than 0)"
    else match n with
      | 1 -> Z
      | n -> S (var_of_int (n - 1))

  let pp_var ppf v =
    Fmt.pf ppf "@[$%d@]" (int_of_var v)

  module V =
  struct
    let v = var_of_int
  end

  type typ =
    | Int
    | Bool
    | String
    | Lambda of typ * typ
    | Tuple of typ * typ
    | Either of typ * typ

  module Type =
  struct

    let int = Int
    let bool = Bool
    let string = String
    let lambda a b = Lambda (a, b)
    let tuple a b = Tuple (a, b)
    let either a b = Either (a, b)

    let ( @-> ) = lambda
    let ( * ) = tuple
    let ( or ) = either
  end

  let pp_infix ~infix pp_a pp_b ppf (a, b) =
    Fmt.pf ppf "@[<1>(@[%a@]@ %s @[%a@])@]" pp_a a infix pp_b b

  let rec pp_typ ppf = function
    | Int -> Fmt.string ppf "int"
    | Bool -> Fmt.string ppf "bool"
    | String -> Fmt.string ppf "string"
    | Lambda (a, b) -> Fmt.pf ppf "%a" (pp_infix ~infix:"->" pp_typ pp_typ) (a, b)
    | Tuple (a, b) -> Fmt.pf ppf "%a" (pp_infix ~infix:"*" pp_typ pp_typ) (a, b)
    | Either (a, b) -> Fmt.pf ppf "%a" (pp_infix ~infix:"or" pp_typ pp_typ) (a, b)

  type arithmetic = [ `Add | `Sub | `Mul | `Div ]

  type primitive =
    { name : string
    ; typ  : typ list * typ
    ; exp  : const list -> const }

  and const =
    | VInt of int
    | VBool of bool
    | VString of string
    | VTuple of const * const
    | VEither of typ * (const, const) either
    | VLambda of { typ: typ * typ
                 ; exp: (const -> const) }
  and binop =
    [ arithmetic | `Tuple | `Eq ]
  and unop =
    | Fst | Snd | L of typ | R of typ
  and exp =
    | Con of const
    | Prm of primitive
    | Var of var
    | Abs of typ * exp
    | App of exp * exp
    | Bin of binop * exp * exp
    | Uno of unop * exp
    | Let of typ * exp * exp
    | Swt of { a : exp
             ; b : exp
             ; s : exp }
    | Rec of typ * exp * exp
    | If  of exp * exp * exp

  let rec pp_const ppf = function
    | VInt n -> Fmt.pf ppf "%d" n
    | VBool b -> Fmt.pf ppf "%b" b
    | VString s -> Fmt.pf ppf "%s" s
    | VTuple (a, b) -> Fmt.pf ppf "%a" (Fmt.Dump.pair pp_const pp_const) (a, b)
    | VEither (_, L x) -> Fmt.pf ppf "@[<1>(L@ %a)@]" pp_const x
    | VEither (_, R x) -> Fmt.pf ppf "@[<1>(R@ %a)@]" pp_const x
    | VLambda { typ = (ta, tb); _ } ->
      Fmt.pf ppf "@[<hov>#lambda:%a@]"
        pp_typ (Lambda (ta, tb))

  let rec pp ppf = function
    | Con c -> Fmt.pf ppf "%a" pp_const c
    | Var n -> Fmt.pf ppf "$%d" (int_of_var n)
    | Prm { name; typ = (a, r); _ } ->
      Fmt.pf ppf "(@[<1>#%s@ @[%a@]@])"
        name
        pp_typ (List.fold_right (fun x a -> Lambda (x, a)) a r)
    | Abs (t, e) ->
      Fmt.pf ppf "@[<2>(\\@[%a@]@ @[<2>%a@])@]" pp_typ t pp e
    | App (f, a) ->
      Fmt.pf ppf "@[@[%a@]@ @[%a@]@]" pp f pp a
    | Bin (`Add, a, b) ->
      Fmt.pf ppf "%a" (pp_infix ~infix:"+" pp pp) (a, b)
    | Bin (`Sub, a, b) ->
      Fmt.pf ppf "%a" (pp_infix ~infix:"-" pp pp) (a, b)
    | Bin (`Mul, a, b) ->
      Fmt.pf ppf "%a" (pp_infix ~infix:"*" pp pp) (a, b)
    | Bin (`Div, a, b) ->
      Fmt.pf ppf "%a" (pp_infix ~infix:"/" pp pp) (a, b)
    | Bin (`Eq, a, b) ->
      Fmt.pf ppf "%a" (pp_infix ~infix:"=" pp pp) (a, b)
    | Bin (`Tuple, a, b) ->
      Fmt.pf ppf "%a" (Fmt.Dump.pair pp pp) (a, b)
    | Uno (Fst, a) ->
      Fmt.pf ppf "@[<2>(fst@ @[%a@])@]" pp a
    | Uno (Snd, a) ->
      Fmt.pf ppf "@[<2>(snd@ @[%a@])@]" pp a
    | Uno (L _, a) ->
      Fmt.pf ppf "@[<2>(L@ @[%a@])@]" pp a
    | Uno (R _, a) ->
      Fmt.pf ppf "@[<2>(R@ @[%a@])@]" pp a
    | Let (t, v, f) ->
      Fmt.pf ppf "@[<2>let @[$1 : %a@] =@ @[%a@] in@ @[%a@]@]" (Fmt.hvbox pp_typ) t (Fmt.hvbox pp) v (Fmt.hvbox pp) f
    | Swt { a; b; s; } ->
      Fmt.pf ppf "@[<2>@[match %a with@\n| @[<2>@[L $1@] ->@ @[%a@]@]@\n| @[<2>@[R $1@] ->@ @[%a@]@]@]@]"
        pp s pp a pp b
    | Rec (_, z, s) ->
      Fmt.pf ppf "@[<2>fix@\n@[z:%a@]@\n@[s:%a@]@]" pp z pp s
    | If (t, a, b) ->
      Fmt.pf ppf "@[if %a@\nthen @[%a@]@\nelse @[%a@]@]"
        pp t pp a pp b

  let int n = Con (VInt n)
  let string s = Con (VString s)
  let tuple (a, b) = Bin (`Tuple, a, b)
  let lambda t f = Abs (t, f)
  let apply f a = App (f, a)
  let fix ~typ ~init s = Rec (typ, init, s)
  let var n = Var V.(v n)

  let ( = ) a b = Bin (`Eq, a, b)
  let ( + ) a b = Bin (`Add, a, b)
  let ( * ) a b = Bin (`Mul, a, b)
  let ( - ) a b = Bin (`Sub, a, b)

  let si t a b = If (t, a, b)
  let left ~rtyp x = Uno (L rtyp, x)
  let right ~ltyp x = Uno (R ltyp, x)
  let fst x = Uno (Fst, x)
  let snd x = Uno (Snd, x)
end

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

  let rec int_of_var : type e a. (e, a) var -> int = function
    | Z -> 1
    | S n -> 1 + (int_of_var n)

  module V =
  struct
    let o = Z
    let x = ()

    let ( $ ) x () = S x

    (* let v = o$x$x$x$x$x = 5 (church style) *)
  end

  type 'a typ =
    | Int    : int typ
    | Bool   : bool typ
    | String : string typ
    | Lambda : 'a typ * 'b typ -> ('a -> 'b) typ
    | Tuple  : 'a typ * 'b typ -> (('a, 'b) cpl) typ
    | Either : 'a typ * 'b typ -> (('a, 'b) either) typ

  let rec pp_typ : type a. a typ Fmt.t = fun ppf -> function
    | Int -> Fmt.pf ppf "int"
    | Bool -> Fmt.pf ppf "bool"
    | String -> Fmt.pf ppf "string"
    | Lambda (a, b) -> Fmt.pf ppf "%a" (A.pp_infix ~infix:"->" pp_typ pp_typ) (a, b)
    | Tuple (a, b) -> Fmt.pf ppf "%a" (A.pp_infix ~infix:"*" pp_typ pp_typ) (a, b)
    | Either (a, b) -> Fmt.pf ppf "%a" (A.pp_infix ~infix:"or" pp_typ pp_typ) (a, b)

  module Type =
  struct
    let int = Int
    let bool = Bool
    let string = String
    let lambda a b = Lambda (a, b)
    let tuple a b = Tuple (a, b)
    let either a b = Either (a, b)

    let ( @-> ) = lambda
    let ( * ) = tuple
    let ( or ) = either
  end

  type ('l, 'r, 'v) binop =
    | Add : (int, int, int) binop
    | Mul : (int, int, int) binop
    | Sub : (int, int, int) binop
    | Div : (int, int, int) binop
    | Cpl : ('a, 'b, ('a, 'b) cpl) binop
    | Eq  : ('a, 'a, bool) binop

  type ('u, 'v) unop =
    | Fst : (('a, 'b) cpl, 'a) unop
    | Snd : (('a, 'b) cpl, 'b) unop
    | L : ('a, ('a, 'b) either) unop
    | R : ('b, ('a, 'b) either) unop

  type ('r, 'e) prim =
    { name : string
    ; exp  : (env:'e -> 'r)
    ; env  : 'e env
    ; typ  : 'r typ }

  and ('e, 'a) exp =
    | Con : 'a -> ('e, 'a) exp
    | Prm : ('r, 'e) prim -> ('e, 'r) exp
    | Uno : ('a, 'res) unop * ('e, 'a) exp -> ('e, 'res) exp
    | Bin : ('a, 'b, 'res) binop * ('e, 'a) exp * ('e, 'b) exp -> ('e, 'res) 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
    | Let : 'a typ * ('e, 'a) exp * ('e * 'a, 'b) exp -> ('e, 'b) exp
    | Rec : { z : ('e, 'a) exp
            ; s : ('e * 'a, ('a, 'r) either) exp } -> ('e, 'r) exp
    | Swt : { s : ('e, ('l, 'r) either) exp
            ; a : ('e * 'l, 'a) exp
            ; b : ('e * 'r, 'a) exp } -> ('e, 'a) exp
    | If  : ('e, bool) exp * ('e, 'a) exp * ('e, 'a) exp -> ('e, 'a) exp

  and 'e env =
    | [] : unit env
    | (::) : 'a typ * 'e env -> ('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
      | Prm { exp; _ } -> exp ~env:e
      | Uno (L, x) -> L (run x e)
      | Uno (R, x) -> R (run x e)
      | Uno (Fst, x) -> let T (a, _) = run x e in a
      | Uno (Snd, x) -> let T (_, b) = run x e in b
      | Var x -> get x e
      | Abs (_, r) -> (fun v -> run r (e, v))
      | App (f, a) -> (run f e) (run a e)
      | Bin (Add, l, r) -> (run l e) + (run r e)
      | Bin (Sub, l, r) -> (run l e) - (run r e)
      | Bin (Mul, l, r) -> (run l e) * (run r e)
      | Bin (Div, l, r) -> (run l e) / (run r e)
      | Bin (Cpl, l, r) -> T (run l e, run r e)
      | Bin (Eq, l, r) -> (run l e) = (run r e)
      | Let (_, a, f) ->
        run f (e, (run a e))
      | Swt { s; a; b; } ->
        (match run s e with
         | L l -> run a (e, l)
         | R r -> run b (e, r))
      | Rec { z; s; } ->
        let z' = run z e in
        let rec loop a = match run s (e, a) with
          | L a -> loop a
          | R r -> r
        in
        loop z'
      | If (t, a, b) ->
        (match run t e with
         | true -> run a e
         | false -> run b e)

  let int (x:int) = Con x
  let string (x:string) = Con x
  let tuple (a, b) = Bin (Cpl, a, b)
  let lambda t e = Abs (t, e)
  let apply f a = App (f, a)

  let fix ~z ~s = Rec { z; s; }
  let var x = Var x

  let ( = ) a b = Bin (Eq, a, b)
  let ( + ) a b = Bin (Add, a, b)
  let ( * ) a b = Bin (Mul, a, b)
  let ( - ) a b = Bin (Sub, a, b)

  let si t a b = If (t, a, b)

  let left e = Uno (L, e)
  let right e = Uno (R, e)
  let fst e = Uno (Fst, e)
  let snd e = Uno (Snd, 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.arrow 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) cpl, (b, d) cpl) Eq.refl option

  = fun x f -> match x with
    | Some (Eq.Refl as x) -> Eq.cpl 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.star 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) either, (b, d) either) Eq.refl option

  = fun x f -> match x with
    | Some (Eq.Refl as x) -> Eq.either 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.Bool,           B.Bool           -> Some Eq.Refl
    | B.String,         B.String         -> Some Eq.Refl
    | B.Lambda (a, a'), B.Lambda (b, b') -> eql_typ a b >>= fun Eq.Refl -> eql_typ a' b'
    | B.Either (a, a'), B.Either (b, b') -> eql_typ a b >?= fun Eq.Refl -> eql_typ a' b'
    | B.Tuple  (a, a'), B.Tuple  (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.[], B.[] -> Some Eq.Refl
    | B.(at :: a), B.(bt :: b) ->
      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.Bool -> Wtyp B.Bool
    | A.String -> Wtyp B.String
    | A.Lambda (a, b) ->
      let Wtyp a = typ a in
      let Wtyp b = typ b in
      Wtyp (B.Lambda (a, b))
    | A.Tuple (a, b) ->
      let Wtyp a = typ a in
      let Wtyp b = typ b in
      Wtyp (B.Tuple (a, b))
    | A.Either (a, b) ->
      let Wtyp a = typ a in
      let Wtyp b = typ b in
      Wtyp (B.Either (a, b))

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

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

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.(t' :: r'), tr)
    | A.Z, t :: r ->
      let Wtyp t = typ t in
      let Wenv r = env r in

      Wvar (B.Z, B.(t :: r), t)
    | v, g -> invalid_arg (Fmt.strf "Variable %a unbound in %a."
                             A.pp_var v Fmt.(Dump.list A.pp_typ) g)

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

type error =
  | EnvMismatch of { g  : wenv
                   ; g' : wenv }
  | TypMismatch of { a  : wtyp
                   ; b  : wtyp }
  | ExpectedTuple of wexp
  | ExpectedLambda of wexp
  | ExpectedEither of wexp
  | UnexpectedAnonymousLambda

exception Error of A.exp * A.typ list * error list

let rec unsafe_cast_const
  : type a. A.const -> a B.typ -> (unit, a) B.exp
  = fun v t ->
    let open B in

    match v, t with
    | A.VInt n, B.Int -> Con n
    | A.VBool b, B.Bool -> Con b
    | A.VString s, B.String -> Con s
    | A.VTuple (a, b), B.Tuple (ta, tb) ->
      (match unsafe_cast_const a ta, unsafe_cast_const b tb with
       | Con a, Con b -> Con (T (a, b))
       | _, _ -> assert false)
    (* XXX(dinosaure); [unsafe_cast_const] returns only a [B.Con _]. *)
    | A.VEither (tr, L x), B.Either (tl, tr') ->
      let Wtyp tr = typ tr in
      (match unsafe_cast_const x tl, eql_typ tr tr' with
       | Con x, Some Eq.Refl -> Con (L x)
       | _, Some Eq.Refl -> assert false
       | _, None -> invalid_arg (Fmt.strf "Unable to cast %a with type %a."
                                   A.pp_const v B.pp_typ t))
    | A.VEither (tl, R x), B.Either (tl', tr) ->
      let Wtyp tl = typ tl in
      (match unsafe_cast_const x tr, eql_typ tl tl' with
       | Con x, Some Eq.Refl -> Con (R x)
       | _, Some Eq.Refl -> assert false
       | _, None -> invalid_arg (Fmt.strf "Unable to cast %a with type %a."
                                   A.pp_const v B.pp_typ t))
    | A.VLambda { typ = (uta, utr); exp; }, B.Lambda (ta, tr) ->
      let Wtyp ta' = typ uta in
      let Wtyp tr' = typ utr in

      let exp = fun a ->
        let ua = safe_cast_const ta a in
        let ur = exp ua in
        match unsafe_cast_const ur tr, eql_typ ta ta', eql_typ tr tr' with
        | Con r, Some Eq.Refl, Some Eq.Refl -> r
        | _, None, _ | _, _, None ->
          invalid_arg (Fmt.strf "Unable to cast %a with type %a."
                         A.pp_const v B.pp_typ t)
        | _ -> assert false
      in Con exp
    | v, t ->
      invalid_arg (Fmt.strf "Unable to cast %a with type %a."
                     A.pp_const v B.pp_typ t)

and safe_cast_typ
  : type a. a B.typ -> A.typ
  = function
    | B.Int -> A.Int
    | B.Bool -> A.Bool
    | B.String -> A.String
    | B.Tuple (a, b) -> A.Tuple (safe_cast_typ a, safe_cast_typ b)
    | B.Either (a, b) -> A.Either (safe_cast_typ a, safe_cast_typ b)
    | B.Lambda (a, b) -> A.Lambda (safe_cast_typ a, safe_cast_typ b)

and safe_cast_const
  : type a. a B.typ -> a -> A.const
  = fun t v -> match t, v with
    | B.Int, n -> A.VInt n
    | B.Bool, b -> A.VBool b
    | B.String, s -> A.VString s
    | B.Tuple (ta, tb), T (a, b) -> A.VTuple (safe_cast_const ta a, safe_cast_const tb b)
    | B.Either (tl, tr), L l -> A.VEither (safe_cast_typ tr, L (safe_cast_const tl l))
    | B.Either (tl, tr), R r -> A.VEither (safe_cast_typ tl, R (safe_cast_const tr r))
    | B.Lambda (ta, tr), f ->
      let exp = (fun ua ->
          match unsafe_cast_const ua ta with
          | B.Con a -> safe_cast_const tr (f a)
          | _ -> assert false)
          (* XXX(dinosaure); [unsafe_cast_const] returns only a [B.Con _]. *)
      in
      A.VLambda { typ = (safe_cast_typ ta, safe_cast_typ tr)
                ; exp }

let to_index =
  let rec go acc index = function
    | [] -> List.rev acc
    | _ :: rest -> go (index :: acc) (index + 1) rest
  in go [] 1

let env_from_args args =
  List.fold_left
    (fun (Wenv e) t ->
       let Wtyp t' = typ t in
       Wenv B.(t' :: e))
    (Wenv B.[]) args

let uncurry args =
  let open B in
  List.fold_right
    (fun t -> function
       | Wexp (a, ta :: e, tr) ->
         let Wtyp t' = typ t in
         (match eql_typ ta t' with
          | Some Eq.Refl ->
            Wexp (Abs (t', a), e, Lambda (t', tr))
          | None ->
            invalid_arg (Fmt.strf "Type mismatch: %a <> %a."
                           B.pp_typ ta B.pp_typ t'))
       | Wexp _ -> invalid_arg "Impossible to uncurry, \
                                environment mismatch")
    args

let var_of_int
  : A.typ list -> int -> wvar
  = fun e n ->
    let v = A.var_of_int n in
    var v e

let env_to_list
  : type e. e B.env -> e -> A.const list
  = fun g e ->
    let rec go
      : type e. A.const list -> e B.env -> e -> A.const list
      = fun a g e -> match g, e with
      | B.(t :: tr), (vr, v) ->
        let c = safe_cast_const t v in
        go (c :: a) tr vr
      | B.([]), () -> a
    in go [] g e

let cast_primitive ~name g args rettype unsafe_primitive =
  let rec chop n l = match n, l with
    | 0, l -> l
    | n, _ :: r -> chop (n - 1) r
    | _, _ -> assert false in
  let miss = List.length g in
  let Wenv g' = env List.(rev_append args g) in
  Wexp (B.Prm { B.name
              ; env = g'
              ; typ = rettype
              ; exp = (fun ~env ->
                    let ua = env_to_list g' env in
                    let uc = unsafe_primitive (chop miss ua) in
                    match unsafe_cast_const uc rettype with
                    | B.Con v -> v
                    | _ -> assert false) },
        g', rettype)

let rec exp
  : A.exp -> A.typ list -> wexp
  = fun e g ->
    match e with
    | A.Con (A.VInt i) ->
      let Wenv g' = env g in
      Wexp (B.Con i, g', B.Int)
    | A.Con (A.VBool b) ->
      let Wenv g' = env g in
      Wexp (B.Con b, g', B.Bool)
    | A.Con (A.VString s) ->
      let Wenv g' = env g in
      Wexp (B.Con s, g', B.String)
    | A.Con (A.VEither (tr, L x)) ->
      let Wenv g' = env g in
      let Wtyp tr' = typ tr in
      let Wexp (x', g'', tx') = exp (A.Con x) g in

      (match x', eql_env g' g'' with
       | B.Con x', Some Eq.Refl ->
         Wexp (B.Con (L x'), g', B.Either (tx', tr'))
       | _ ->
         raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' } ])))
    | A.Con (A.VEither (tl, R x)) ->
      let Wenv g' = env g in
      let Wtyp tl' = typ tl in
      let Wexp (x', g'', tx') = exp (A.Con x) g in

      (match x', eql_env g' g'' with
       | B.Con x', Some Eq.Refl ->
         Wexp (B.Con (R x'), g', B.Either (tl', tx'))
       | _ ->
         raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' } ])))
    | A.Con (A.VLambda _) ->
      raise (Error (e, g, [ UnexpectedAnonymousLambda ]))
    | A.Prm { A.name; typ = (args, ret); exp; } ->
      let Wtyp rettype = typ ret in

      (match cast_primitive ~name g args rettype exp with
       | Wexp (B.Prm _, _, _) as wexp ->
         uncurry args wexp
       | _ -> assert false)
    (* XXX(dinosaure): [p] returns only a [B.Prm _]. *)
    | A.Uno (A.L tr, x) ->
      let Wenv g' = env g in
      let Wtyp tr' = typ tr in
      let Wexp (x', g'', tx') = exp x g in

      (match eql_env g' g'' with
       | Some Eq.Refl ->
         Wexp (B.Uno (B.L, x'), g', B.Either (tx', tr'))
       | _ ->
         raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' } ])))
    | A.Uno (A.R tl, x) ->
      let Wenv g' = env g in
      let Wtyp tl' = typ tl in
      let Wexp (x', g'', tx') = exp x g in

      (match eql_env g' g'' with
       | Some Eq.Refl ->
         Wexp (B.Uno (B.R, x'), g', B.Either (tl', tx'))
       | _ ->
         raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' } ])))
    | A.Uno (A.Fst, x) ->
      let Wenv g' = env g in
      (match exp x g with
       | Wexp (x', g'', B.Tuple (ta', _)) ->
         (match eql_env g' g'' with
          | Some Eq.Refl ->
            Wexp (B.Uno (B.Fst, x'), g', ta')
          | _ ->
            raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' } ])))
       | wexp -> raise (Error (e, g, [ ExpectedTuple wexp ])))
    | A.Uno (A.Snd, x) ->
      let Wenv g' = env g in
      (match exp x g with
       | Wexp (x', g'', B.Tuple (_, tb')) ->
         (match eql_env g' g'' with
          | Some Eq.Refl ->
            Wexp (B.Uno (B.Snd, x'), g', tb')
          | _ ->
            raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' } ])))
       | wexp -> raise (Error (e, g, [ ExpectedTuple wexp ])))
    | A.Con (A.VTuple (a, b)) ->
      let Wenv g' = env g in
      let Wexp (a', g'', ta') = exp (A.Con a) g in
      let Wexp (b', g''', tb') = exp (A.Con b) g in
      (match a', b', eql_env g' g'', eql_env g'' g''' with
       | B.Con a', B.Con b', Some (Eq.Refl as g1), Some (Eq.Refl as g2) ->
         (match Eq.trans g1 g2 with Eq.Refl -> Wexp (B.Con (T (a', b')), g', B.Tuple (ta', tb')))
       | _ ->
         raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' }
                             ; EnvMismatch { g = Wenv g''; g' = Wenv g''' } ])))
    | A.Bin (`Tuple, a, b) ->
      let Wenv g' = env g in
      let Wexp (a', g'', ta') = exp a g in
      let Wexp (b', g''', tb') = exp b g in
      (match eql_env g' g'', eql_env g'' g''' with
       | Some (Eq.Refl as g1), Some (Eq.Refl as g2) ->
         (match Eq.trans g1 g2 with Eq.Refl -> Wexp (B.Bin (B.Cpl, a', b'), g', B.Tuple (ta', tb')))
       | _, _ ->
         raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' }
                             ; EnvMismatch { g = Wenv g''; g' = Wenv g''' } ])))
    | 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
      (let open B in match exp e (t :: g) with
        | Wexp (e', t'' :: g', tr) ->
          (match eql_typ t' t'' with
           | Some Eq.Refl -> Wexp (B.Abs (t', e'), g', B.Lambda (t', tr))
           | None -> raise (Error (e, g, [ TypMismatch { a = Wtyp t'; b = Wtyp t'' } ])))
        | Wexp (_, g'', _) -> raise (Error (e, g, [ EnvMismatch { g = Wenv B.(t' :: g''); g' = Wenv g'' } ])))
    | A.App (f, a) ->
      (match exp f g, exp a g with
       | Wexp (f', g', (B.Lambda (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')
          | _, _ -> raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' }
                                        ; TypMismatch { a = Wtyp t'; b = Wtyp t'' } ])))
       | wexp, _ ->
         raise (Error (e, g, [ ExpectedLambda wexp ])))
    | A.Bin (#A.arithmetic as binop, l, r) ->
      let binop = match binop with
        | `Add -> B.Add
        | `Sub -> B.Sub
        | `Mul -> B.Mul
        | `Div -> B.Div
      in
      (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.Bin (binop, l', r'), g', B.Int)
          | None -> raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv g'' } ])))
       | Wexp (_, _, ta), Wexp (_, _, tb) ->
         raise (Error (e, g, [ TypMismatch { a = Wtyp B.Int; b = Wtyp ta }
                             ; TypMismatch { a = Wtyp B.Int; b = Wtyp tb } ])))
    | A.Let (t, a, f) ->
      let Wenv g' = env g in

      (match typ t, exp a g, exp f g with
       | Wtyp t', Wexp (a', g'', t''), Wexp (f', g''', r') ->
         (match eql_typ t' t'' , eql_env g' g'', eql_env B.(t'' :: g'') g''' with
          | Some Eq.Refl, Some Eq.Refl, Some Eq.Refl ->
            Wexp (B.Let (t', a', f'), g', r')
          | _, _, _ ->
            raise (Error (e, g, [ TypMismatch { a = Wtyp t'; b = Wtyp t'' }
                                ; EnvMismatch { g = Wenv g'; g' = Wenv g'' }
                                ; EnvMismatch { g = Wenv B.(t'' :: g''); g' = Wenv g''' } ]))))
    | A.If (t, a, b) ->
      (match exp t g, exp a g, exp b g with
       | Wexp (t', g', B.Bool),
         Wexp (a', g'', u'),
         Wexp (b', g''', v') ->
         (match eql_typ u' v', eql_env g' g'', eql_env g'' g''' with
          | Some Eq.Refl, Some (Eq.Refl as ae), Some (Eq.Refl as be) ->
            (match Eq.trans ae be with Eq.Refl -> Wexp (B.If (t', a', b'), g', u'))
          | _, _, _ ->
            raise (Error (e, g, [ TypMismatch { a = Wtyp u'; b = Wtyp v' }
                                ; EnvMismatch { g = Wenv g'; g' = Wenv g'' }
                                ; EnvMismatch { g = Wenv g''; g' = Wenv g''' } ])))
       | Wexp (_, _, t'), _, _ ->
         raise (Error (e, g, [ TypMismatch { a = Wtyp B.Bool; b = Wtyp t' } ])))
    | A.Bin (`Eq, a, b) ->
      let Wexp (a', g', ta') = exp a g in
      let Wexp (b', g'', tb') = exp b g in
      (match eql_typ ta' tb', eql_env g' g'' with
       | Some Eq.Refl, Some Eq.Refl -> Wexp (B.Bin (B.Eq, a', b'), g', B.Bool)
       | _, _ -> assert false)
    | A.Swt { a; b; s; } ->
      (match exp s g with
       | Wexp (s', g'', B.Either (l', r')) ->
         let l = safe_cast_typ l' in
         let r = safe_cast_typ r' in

         (let open B in match exp a (l :: g), exp b (r :: g) with
           | Wexp (a', l'' :: ag'', ar),
             Wexp (b', r'' :: bg'', br) ->
             (match eql_typ l' l'', eql_typ r' r'',
                    eql_env ag'' g'', eql_env bg'' g'',
                    eql_typ ar br with
             | Some Eq.Refl, Some Eq.Refl,
               Some Eq.Refl, Some Eq.Refl,
               Some Eq.Refl ->
               Wexp (B.Swt { s = s'; a = a'; b = b'; }, g'', ar)
             | _, _, _, _, _ ->
               raise (Error (e, g, [ TypMismatch { a = Wtyp l'; b = Wtyp l'' }
                                   ; TypMismatch { a = Wtyp r'; b = Wtyp r'' }
                                   ; EnvMismatch { g = Wenv ag''; g' = Wenv g'' }
                                   ; EnvMismatch { g = Wenv bg''; g' = Wenv g'' }
                                   ; TypMismatch { a = Wtyp ar; b = Wtyp br } ])))
           | Wexp (_, ag'', _),
             Wexp (_, bg'', _) ->
             raise (Error (e, g, [ EnvMismatch { g = Wenv ag''; g' = Wenv B.(l' :: ag'') }
                                 ; EnvMismatch { g = Wenv bg''; g' = Wenv B.(r' :: bg'') } ])))
       | wexp -> raise (Error (e, g, [ ExpectedEither wexp ])))
    | A.Rec (u, z, s) ->
      let Wenv g' = env g in
      let Wtyp u' = typ u in

      let gz = g in
      let gs = (u :: g) in

      let Wenv gz' = env gz in
      let Wenv gs' = env gs in
      let Wexp (z', gz'', tz') = exp z gz in
      let Wexp (s', gs'', ts') = exp s gs in

      (match eql_env g' gz',
             eql_env B.(tz' :: g') gs' with
      | Some (Eq.Refl as f1), Some (Eq.Refl as e1)->
        (match eql_typ tz' u',
               eql_typ ts' (B.Either (tz', tz')),
               eql_env gz' gz'',
               eql_env gs' gs'' with
        | Some Eq.Refl, Some Eq.Refl, Some (Eq.Refl as f2), Some (Eq.Refl as e2) ->
          (match Eq.trans f1 f2, Eq.trans e1 e2 with
           | Eq.Refl, Eq.Refl -> Wexp (B.Rec { z = z'; s = s'; }, g', u'))
        | _, _, _, _ ->
          raise (Error (e, g, [ TypMismatch { a = Wtyp tz'; b = Wtyp u' }
                              ; TypMismatch { a = Wtyp ts'; b = Wtyp (B.Either (tz', tz')) }
                              ; EnvMismatch { g = Wenv gz'; g' = Wenv gz'' }
                              ; EnvMismatch { g = Wenv gs'; g' = Wenv gs'' } ])))
      | _, _ ->
        raise (Error (e, g, [ EnvMismatch { g = Wenv g'; g' = Wenv gz' }
                            ; EnvMismatch { g = Wenv B.(tz' :: g'); g' = Wenv gs' } ])))

(* RUN **********************************************************************************)

let run
  : type a. A.exp -> a B.typ -> a
  = fun m t ->
  let open B in match exp m [] with
  | Wexp (m', [], t') ->
    let v = B.run m' () in
    (match eql_typ t t' with Some Eq.Refl -> v | None -> raise (Error (m, [], [ TypMismatch { a = Wtyp t; b = Wtyp t' } ])))
  | Wexp (_, g', _) ->
    raise (Error (m, [], [ EnvMismatch { g = Wenv []; g' = Wenv g' } ]))

(* EXAMPLE ******************************************************************************)

let () = assert (run (A.If (A.Con (A.VBool true), A.Con (A.VInt 42), A.Con (A.VInt 21))) B.Int = 42)

let switch_example =
  A.Swt
    { a = A.Var A.Z
    ; b = A.Con (A.VString "<int>")
    ; s = A.Uno (A.L A.Int, A.Con (A.VString "Hello World!")) }

let k =
  A.App
    ((A.App
        ((A.Abs (A.Int,
                 A.Abs (A.Int, A.Var A.(S Z)))),
         (A.Con (A.VInt 42)))),
     (A.Con (A.VInt 21)))

let () = assert (run switch_example B.String = "Hello World!")
let () = assert (run k B.Int = 42)

let fix_example =
  let open B in

  fix
    ~z:(var V.o)
    ~s:(si (int 123 = var V.o)
           (left (var V.o + var V.o))
           (right (var V.o)))

let fact_example =
  let code =
    let rec fact v =
      if 0 = (snd v)
      then (fst v, 0)
      else fact ((fun x -> (fst v) x * (snd v)), (snd v) - 1)
    in
    let init x = ((fun _ -> 1), x) in

    fun x -> (fst ((fun x -> fact (init x)) x)) 0
  in

  assert (code 5 = 120);

  let open B in

  let main =
    fix
      ~z:(tuple (lambda Type.int (int 1), var V.o))
      ~s:(si (int 0 = (snd (var V.o)))
            (right (tuple (fst (var V.o), int 0)))
            (left  (tuple (lambda Type.int ((apply (fst (var V.(o$x))) (var V.o)) * (snd (var V.(o$x)))),
                           ((snd (var V.o)) - int 1)))))
  in

  lambda Type.int (apply (fst (apply (lambda Type.int main) (var V.o))) (int 0))

let unsafe_fact =
  let open A in

  let main =
    let typ = Type.(tuple (int @-> int) int) in
    fix
      ~typ
      ~init:(tuple (lambda Type.int (int 1), var 1))
      (si
        (snd (var 1) = (int 0))
        (right
          ~ltyp:typ
          (tuple (fst (var 1), int 0)))
        (left
          ~rtyp:typ
          (tuple
            (lambda Type.int ((apply (fst (var 2)) (var 1)) * (snd (var 2))),
             ((snd (var 1)) - (int 1))))))
  in

  lambda Type.int (apply (fst (apply (lambda Type.int main) (var 1))) (int 0))

let () = assert (run unsafe_fact B.Type.(int @-> int) 5 = B.run fact_example () 5)

let padd =
  { A.name = "%add"
  ; typ = ([A.Int; A.Int], A.Int)
  ; exp = (function
        | A.VInt a :: A.VInt b :: [] -> A.VInt (a + b)
        | args ->
          invalid_arg (Fmt.strf "%%add: receive %a."
                         Fmt.(Dump.list A.pp_const) args)) }
  |> fun p -> A.Prm p

let padebool =
  { A.name = "%add-bool"
  ; typ = ([A.Int; A.Bool], A.Int)
  ; exp = (function
        | A.VInt a :: A.VBool true :: [] -> A.VInt (a + 1)
        | A.VInt a :: A.VBool false :: [] -> A.VInt a
        | args ->
          invalid_arg (Fmt.strf "%%add-bool: receive %a."
                         Fmt.(Dump.list A.pp_const) args)) }
  |> fun p -> A.Prm p

let () = assert (run padd B.Type.(int @-> int @-> int) 21 21 = 42)
let () = assert (run padebool B.Type.(int @-> bool @-> int) 0 true = 1)

let env_and_prim =
  let open A in

  (lambda Type.string (apply (apply padd (int 21)) (int 21)))

let () = assert (run env_and_prim B.Type.(string @-> int) "Hello World!" = 42)
let () = assert (run A.(lambda Type.int env_and_prim) B.Type.(int @-> string @-> int) 0 "Hello World!" = 42)