(**********************************************************************

  plc-nbe.ml

  Type checking and normalisation by evaluation for 
  Polymorphic Lambda Calculus (PLC) in Fresh OCaml
 
  by A.M.Pitts, 2 August 2003

*********************************************************************)


type t;;                       (* some type *)

type tyvar = t name;;          (* type variables, a *)

type ty =                      (* types, t *)
    Tyvar of tyvar             (* a *)
  | Fun of ty * ty             (* t1 -> t2 *)
  | All of <<tyvar>>ty;;       (* All a. t  *)

type t;;                       (* some type *)

type var = t name;;            (* variables, x *)

type expr =                    (* expressions, e *)
    Var of var                 (* x *)
  | Lam of ty * (<<var>>expr)  (* \x:t.e *)
  | App of expr * expr         (* e1 e2 *)
  | Gen of <<tyvar>>expr       (* \a.e *)
  | Spec of expr * ty;;        (* e t *)

type 'a env = (var * 'a)list;; (* environments *)

type sem =                     (* semantics, d *)
    L of ty * (sem -> sem)     
  | G of (ty -> sem)
  | N of neu
and neu =                      (* neutral elements, n *)
    V of var
  | A of neu * sem
  | S of neu * ty;;

(********************************************************************* 
   Type inference:
   typ (tenv : ty env) (e : expr) returns the type of e in the typing
   environment tenv, raising Typ_failure if there is no such type 
*********************************************************************)

let ( |-> ) (a : tyvar) (t : ty) : ty -> ty =
  (* capture-avoiding substitution of t for free occurrences of 
     (Tyvar a) in types *)
  let rec sub (t' : ty) : ty =
    match t' with
	Tyvar a1 -> if a = a1 then t else t'
      | Fun(t1, t2) -> Fun(sub t1, sub t2)
      | All(<<a1>>t1) -> All(<<a1>>(sub t1))
        (* note that a is automatically fresh for t, so we 
	   avoid "capture" in this case! *)
  in sub;;

exception Typ_failure;;

let rec typ (tenv : ty env) (e : expr) : ty =
  match e with 
      Var x -> 
	(try List.assoc x tenv with Not_found -> raise Typ_failure)
    | Lam(t, <<x>>e1) -> Fun(t, typ ((x, t) :: tenv) e1)
    | App(e1, e2) ->
	(match typ tenv e1 with
	    Fun(t1, t2) -> 
	      if t1 = typ tenv e2 then t2 else raise Typ_failure
	  | _ -> raise Typ_failure)
    | Gen(<<a>>e1) -> All(<<a>>(typ tenv e1))
    | Spec(e1, t1) ->
	(match typ tenv e1 with
	    All(<<a>>t) -> (a |-> t1)t
	  | _ -> raise Typ_failure);;

(*********************************************************************
  Normalisation by evaluation:

  nbe (tenv : ty env) (e :expr) returns a pair (t, e'), where t is
  the type of e, given tenv, and e' is the beta-normal
  form of e;  exception Typ_failure is raised if there is no type.

  semeq (tenv : ty env) (e1 : expr) (e2 : expr) returns true iff e1
  and e1 have the same type, given tenv, and the same beta-normal
  form; exception Typ_failure is raised if there is no type.
*********************************************************************)

let ( |=> ) (a : tyvar) (t : ty) : expr -> expr =
  (* capture-avoiding substitution of t for free occurrences of 
     (Tyvar a) in expressions *) 
  let rec sub (e : expr) : expr =
    match e with
	Var x1 -> e
      | Lam(t1, <<x1>>e1) -> Lam((a |-> t)t1, <<x1>>(sub e1))
      | App(e1, e2) -> App(sub e1, sub e2)
      | Gen(<<a1>>e1) -> Gen(<<a1>>(sub e1))
      | Spec(e1, t1) -> Spec(sub e1, (a |-> t)t1)
  in sub;;

(* reification *)
let rec reify (d : sem) : expr =
  match d with
      L(t, f) -> 
	let x = (fresh : var) 
	in Lam(t, <<x>>(reify(f(N(V x)))))
    | G g -> 
	let a = (fresh : tyvar) 
	in Gen(<<a>>(reify(g(Tyvar a))))
    | N n -> reifyn n
and reifyn (n : neu) : expr =
  match n with
      V x -> Var x
    | A(n1, d1) -> App(reifyn n1, reify d1)
    | S(n1, t1) -> Spec(reifyn n1, t1);;

exception Eval_failure;;

(* evaluation *)
let rec eval (senv : sem env) (e : expr) : sem =
  match e with
      Var x -> 
	(try List.assoc x senv with Not_found -> N(V x)) 
    | Lam(t, <<x>>e1) ->
	L(t, function d -> eval ((x, d) :: senv) e1)
    | App(e1, e2) ->
	(match eval senv e1 with
	    L(t, f) -> f(eval senv e2)
	  | G _ -> raise Eval_failure
	      (* this case never occurs for well typed expressions *)
	  | N n -> N(A(n, eval senv e2)))
    | Gen(<<a>>e1) -> 
	G(function t -> eval senv ((a |=> t)e1))
    | Spec(e1, t1) -> 
	(match eval senv e1 with
	    L _ -> raise Eval_failure
	      (* this case never occurs for well typed expressions *)
	  | G g -> g t1
	  | N n -> N(S(n, t1)));;

let tyenv_to_semenv (tenv : ty env) : sem env =
  List.map (function (x, t) -> (x, N(V x))) tenv;;

let nbe (tenv : ty env) (e :expr) : ty * expr =
  (typ tenv e, reify(eval (tyenv_to_semenv tenv) e));;

let semeq (tenv : ty env) (e1 : expr) (e2 : expr) : bool =
  (typ tenv e1 = typ tenv e2) &&
  (let senv = tyenv_to_semenv tenv 
  in (reify(eval senv e1) = reify(eval senv e2)));; 
    

(********************************************************************
  Example: list processing in PLC 
*********************************************************************)

let a, a' = (fresh : tyvar), (fresh : tyvar);;

let a_list =
  (* All a'. a' -> (a -> a' -> a') -> a' *)
  All(<<a'>>(
    Fun(Tyvar a', 
    Fun(Fun(Tyvar a, Fun(Tyvar a', Tyvar a')), Tyvar a'))));;

let x, x', f, l = 
  (fresh : var), (fresh : var), (fresh : var), (fresh : var);;

let plc_nil =
  (* /\a. /\a'. \x' : a'. \f : a -> a' -> a'. x' *)
  Gen(<<a>>(Gen(<<a'>>(Lam(Tyvar a', 
  <<x'>>(Lam(Fun(Tyvar a, Fun(Tyvar a', Tyvar a')), 
  <<f>>(Var x'))))))));;

typ [] plc_nil = All(<<a>>a_list);; 

let plc_cons =
  (*  /\a. \x : a. \l : a_list. /\a'. \x' : a'. \f : a -> a' -> a'. 
      f x (l a' x' f) *)
  Gen(<<a>>(Lam(Tyvar a, <<x>>(Lam(a_list, <<l>>(Gen(<<a'>>(
    Lam(Tyvar a', <<x'>>(Lam(Fun(Tyvar a, Fun(Tyvar a', Tyvar a')),
    <<f>>(App(App(Var f, Var x), App(App(Spec(Var l, Tyvar a'), 
    Var x'), Var f))))))))))))));;

typ [] plc_cons = All(<<a>>(Fun(Tyvar a, Fun(a_list, a_list))));; 

let plc_iter =
  (* /\a. /\a'. \x' : a'. \f : a -> a' -> a'. \l : a_list. 
     l a' x' f *)
  Gen(<<a>>(Gen(<<a'>>(Lam(Tyvar a', <<x'>>(
    Lam(Fun(Tyvar a, Fun(Tyvar a', Tyvar a')), <<f>>(Lam(a_list,
    <<l>>(App(App(Spec(Var l, Tyvar a'), Var x'), Var f)))))))))));;

typ [] plc_iter = 
    All(<<a>>(All(<<a'>>(Fun(Tyvar a', 
    Fun(Fun(Tyvar a, Fun(Tyvar a', Tyvar a')), 
    Fun(a_list, Tyvar a')))))));;

let tenv : ty env  = 
  (* [ x' |-> a', f |-> a -> a' -> a' ] *)
  [(x', Tyvar a'); (f, Fun(Tyvar a, Fun(Tyvar a', Tyvar a')))];;

let nil_lhs =
  (* plc_iter a a' x' f (plc_nil a) 
     This should be beta-equal to x' *)
  App(App(App(Spec(Spec(plc_iter, Tyvar a), 
  Tyvar a'), Var x'), Var f), Spec(plc_nil, Tyvar a));;

typ tenv nil_lhs = Tyvar a';;

semeq tenv nil_lhs (Var x');;

let tenv' : ty env =
  (* [ x |-> a, l |-> a_list, x' |-> a', f |-> a -> a' -> a' ] *)
  (x, Tyvar a) :: (l, a_list) :: tenv;;

let cons_lhs =
  (* plc_iter a a' x' f (plc_cons a x l) *)
  App(App(App(Spec(Spec(plc_iter, Tyvar a), 
  Tyvar a'), Var x'), Var f), 
  App(App(Spec(plc_cons, Tyvar a), Var x), Var l));;

typ tenv' cons_lhs = Tyvar a';;

let cons_rhs =
  (* f x (plc_iter a a' x' f l) *)
  App(App(Var f, Var x), App(App(App(Spec(Spec(plc_iter, Tyvar a), 
  Tyvar a'), Var x'), Var f), Var l));;

typ tenv' cons_rhs = Tyvar a';; (* true! *)

semeq tenv' cons_lhs cons_rhs;; (* true! *)

(*********************************************************************
  end of plc-nbe.ml
*********************************************************************)