(*********************************************************************
nomst.ml

Prototype implementation of Nominal System T (NST) normalization 
function in Fresh Objective Caml 
<http://www.cl.cam.ac.uk/users/amp12/fresh-ocaml/>.

Tested with Fresh Objective Caml release 3.10.0

(c) Andrew Pitts
    2009-06-28

Provides:
 
- type exp, NST expressions;
 
- function norm : exp -> exp, normalization function;

- examples: capture-avoiding substitution and length.

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

(* NST types *)
type typ =
    Nat              (* type of numbers *)
  | Atm              (* type of names *)
  | Trm              (* type of object-level lambda-terms*)
  | Fun of typ*typ;; (* function type *)

(* infix notion for function types *)
let ( => ) (t : typ) (t' : typ) : typ = Fun(t, t');;

(* Atoms *)
type t and atm  = t name;; 

(* NST variables *)
type t and var = (t name) * typ;;

let newvar (t : typ) : var = (fresh, t);;

(* NST expressions *)
type exp = 
    Var of var                      (* variable *)
  | Z                               (* zero *)
  | S of exp                        (* successor *)
  | Nrec of exp * exp * exp         (* Nat recursion *)
  | Lam of <<var>>exp               (* function abstraction *)
  | App of exp * exp                (* function application *)
  | At of atm                       (* atomic name *)
(* the following is included to make normal forms a subset of expressions *)
  | New                             (* new atomic name *)
  | Nu of <<atm>>exp                (* name restriction *)
  | Eq of exp * exp * typ           (* name quality *)
  | Swap of exp * exp * exp         (* name swapping *)
  | V of exp                        (* object-level variable *)
  | A of exp * exp                  (* application term *)
  | L of <<atm>>exp                 (* lambda-abstraction term *)
  | Trec of exp * exp * exp * exp;; (* Trm recursion *)

let makeL (e :exp) (e' : exp) : exp =
  let a = fresh in L(<<a>>(Swap(At a, e, e')));;

(* Type-checking NST expressions *)
exception NotTypeable;;
let rec typeOf (e : exp) : typ =
  match e with
      Var x -> snd x
    | Z -> Nat
    | S e1 -> if typeOf e1 = Nat then Nat else raise NotTypeable
    | Nrec(e1,e2,e3) -> 
	let t = typeOf e1 in
	  if typeOf e2 = (Nat => (t => t)) && typeOf e3 = Nat
	  then t else raise NotTypeable
    | Lam(<<x>>e1) -> (typeOf(Var x)) => (typeOf e1)
    | App(e1, e2) ->
	(match typeOf e1 with
	     Fun(t, t') -> if typeOf e2 = t then t' else raise NotTypeable
	   | Nat | Atm | Trm -> raise NotTypeable)
    | At a -> Atm
    | New -> Atm
    | Nu(<<a>>e1) -> typeOf e1
    | Eq(e1, e2, t) -> 
	if typeOf e1 = Atm && typeOf e2 = Atm
	then t => (t => t) else raise NotTypeable
    | Swap(e1, e2, e3) ->
	if typeOf e1 = Atm && typeOf e2 = Atm
	then typeOf e3 else raise NotTypeable
    | V e1 -> if typeOf e1 = Atm then Trm else raise NotTypeable
    | A(e1, e2) ->
	if typeOf e1 = Trm && typeOf e2 = Trm
	then Trm else raise NotTypeable
    | L(<<a>>e1) -> if typeOf e1 = Trm then Trm else raise NotTypeable
    | Trec(e1, e2, e3, e4) ->
	(match typeOf e1 with
	     Fun(Atm, t) -> 
	       if typeOf e2 = (Trm => (Trm => (t => (t => t)))) 
		 && typeOf e3 = (Atm => (Trm => (t => t))) 
		 && typeOf e4 = Trm
	       then t else raise NotTypeable
	   | _ -> raise NotTypeable);;

(* Subsets of normal and neutral forms *)
let rec isNormal(e : exp) : bool =
  match e with
      Z -> true 
    | S e1 -> isNormal e1
    | Lam(<<_>>e1) -> isNormal e1
    | At _ -> true
    | New -> true
    | V _ -> true
    | A(e1, e2) -> (isNormal e1) && (isNormal e2)
    | L(<<a>>e1) -> isNormal e1
    | _ -> 
	(match typeOf e with
	    Fun _ -> false
	   | Nat | Atm | Trm -> isNeutral e)
and isNeutral(e : exp) :bool =
  match e with
      Var _ -> true
    | Nrec(e1, e2, e3) -> 
	(isNormal e1) && (isNormal e2) && (isNeutral e3)
    | App(e1, e2) ->
	(isNeutral e1) && (isNormal e2)
    | Eq(e1, e2, t) -> 
	((isNormal e1) && (isNeutral e2)) ||
	  ((isNeutral e1) && (isNormal e2))  
    | Nu(<<a>>e1) ->
	(match typeOf e1 with 
	     Fun _ -> false
	   | Nat | Atm | Trm -> isNeutral e1)
    | Swap(e1, e2, e3) ->
	(match (typeOf e1, typeOf e2) with
	     (Atm, Atm) -> 
	       (isNormal e1) && (isNormal e2) && (isNeutral e3)
	   | _ -> false)
    | Trec(e1, e2, e3, e4) ->
	(isNormal e1) && (isNormal e2) && (isNormal e3) && (isNeutral e4);;

(* Semantic domain *)
type sem =
    Gd of exp            (* normal forms of ground type *)
  | Fn of (sem -> sem);; (* semantics for higher types *)

(* Reification and reflection *)
let rec reify (t : typ) (d : sem) : exp =
  match t with
      Fun(t1, t2) ->
	(match d with 
	     Fn f -> 
	       let x = (fresh, t1) in 
	       let n = reify t2 (f (reflect t1 (Var x))) in
		 (assert (isNormal n)); Lam(<<x>>n)
	   | Gd _  -> assert false)
    | Nat | Atm | Trm -> 
	(match d with
	     Gd n -> (assert ((typeOf n = t) && (isNormal n))); n
	   | Fn _ -> assert false)
and reflect (t : typ) (e : exp) : sem =
  (assert (isNeutral e));
  match t with
      Fun(t1, t2) -> 
	Fn(function d -> reflect t2 (App(e, reify t1 d)))
    | Nat | Atm | Trm -> Gd e;;
  
(* Auxiliary semantic functions *)

let ( @ ) (d1 : sem) (d2 : sem) : sem =
  match d1 with
      Gd _ -> assert false
    | Fn f -> f d2;;

let rec nu (a : atm) (d : sem) : sem =
  match d with
      Gd Z -> d
    | Gd(S e) -> 
	(match nu a (Gd e) with
	     Gd e' -> Gd(S e')
	   | Fn _ -> assert false)
    | Gd(At a') ->
	if a = a' then Gd(New) else d 
    | Gd New -> d
    | Gd(V e) ->
	(match nu a (Gd e) with
	     Gd e' -> Gd(V e')
	   | Fn _ -> assert false) 
    | Gd(A(e1, e2)) ->
	(match (nu a (Gd e1), nu a (Gd e2)) with
	     (Gd e1', Gd e2') -> Gd(A(e1', e2'))
	   | _ -> assert false)
    | Gd(L(<<a'>>e)) ->
	(match nu a (Gd e) with
	     Gd e' -> Gd(L(<<a'>>e'))
	   | Fn _ -> assert false)
    | Gd e -> Gd(Nu(<<a>>e))
    | Fn f -> 
	Fn(function d ->  let a' = fresh in
		 nu a' ((swap a and a' in f) d));;

let semTrue : sem = 
  Fn(function d1 -> Fn(function d2 -> d1));;

let semFalse : sem = 
  Fn(function d1 -> Fn(function d2 -> d2));;

let rec eq (t : typ) (d1 : sem) (d2 : sem) : sem =
  match d1 with
      Gd(At a) ->
	(match d2 with 
	     Gd(At a') -> 
	       if a = a' then semTrue else semFalse
	   | Gd New -> semFalse
	   | Gd e -> 
	       reflect (t => (t => t)) (Eq(At a, e, t)) 
	   | _ -> assert false)
    | Gd(New) ->
	(match d2 with 
	     Gd(At _) -> semFalse
	   | Gd New -> semTrue
	   | Gd e -> 
	       reflect (t => (t => t)) (Eq(New, e, t)) 
	   | _ -> assert false)
    | Gd e ->
	(match d2 with
	     Gd e' ->
	       reflect (t => (t => t)) (Eq(e, e', t))
	   | Fn _ -> assert false)
    | _ -> assert false;;

let ( <-> ) (d1 : sem) (d2 : sem) : sem -> sem =
  let rec sw (d : sem) : sem =
    match d with
	Gd Z -> d
      | Gd(S e) -> 
	  (match sw(Gd e) with
	       Gd e' -> Gd(S e')
	     | Fn _ -> assert false)
      | Gd(At _) ->
	  ((eq Atm d1 d) @ d2) @ (((eq Atm d2 d) @ d1) @ d)
      | Gd New -> d
      | Gd(V e) ->
	  (match sw(Gd e) with
	       Gd e' -> Gd(V e')
	     | Fn _ -> assert false)
      | Gd(A(e1, e2)) ->
	  (match (sw(Gd e1), sw(Gd e2)) with
	       (Gd e1', Gd e2') -> Gd(A(e1', e2'))
	     | _ -> assert false)
      | Gd(L(<<a>>e)) ->
	  (match sw(Gd e) with 
	       Gd e' -> Gd(L(<<a>>e'))
	     | Fn _ -> assert false)
      | Gd e ->
	  (match (d1, d2) with
	       (Gd e1, Gd e2) -> Gd(Swap(e1, e2, e))
	     | _ -> assert false)
      | Fn f -> 
	  Fn(function d -> sw(f(sw d)))
  in sw;;

let  nrec (t : typ) (d1 : sem) (d2 : sem) : sem -> sem =
  let rec r (d : sem) : sem = 
    match d with
     Gd Z -> d1
   | Gd(S e) -> 
       (d2 @ (Gd e)) @ (r(Gd e))
   | Gd e -> 
       reflect t 
	 (Nrec(reify t d1, 
	       reify (Nat => (t => t)) d2, 
	       e))
   | _ -> assert false
  in r;;

let trec (t : typ) (d1 : sem) (d2 : sem) (d3 : sem) : sem -> sem =
  let rec r(d : sem) : sem =
    match d with
	Gd(V e) -> d1 @ (Gd e)
      | Gd(A(e1, e2)) -> 
	  (((d2 @ (Gd e1)) @ (Gd e2)) @ (r(Gd e1))) @ (r(Gd e2))
      | Gd(L(<<a>>e)) ->
	  nu a (((d3 @ (Gd(At a))) @ (Gd e)) @ (r(Gd e)))
      | Gd e -> 
	  reflect t 
	    (Trec(reify (Atm => Trm) d1, 
		  reify (Trm => (Trm => (t => (t => (t => t))))) d2, 
		  reify (Atm => (Trm => (t => t))) d3, 
		  e))
      | _ -> assert false
  in r;;
	
(* Environments *)
type env = var -> sem;;

let ( |-> ) (x : var) (d : sem) : env -> env =
  let up(v : env) : env =
    function x' -> if x' = x then d else v x'
  in up;;

let init : env =
  function x ->
    reflect (typeOf(Var x)) (Var x);;

(* Evaluation *)
let rec eval (e : exp) (v : env) : sem =
  let t = typeOf e in
    match e with
      Var x -> v x
    | Z -> Gd Z
    | S e1 -> 
	(match eval e1 v with
	     Gd e' -> Gd(S e')
	   | Fn _ -> assert false)
  | Nrec(e1, e2, e3) ->
      nrec t (eval e1 v) (eval e2 v) (eval e3 v)
  | Lam(<<x>>e1) -> 
      Fn(function d -> eval e1 ((x |-> d)v)) 
  | App(e1, e2) ->
      (match eval e1 v with 
	   Fn f -> f(eval e2 v)
	 | Gd _ -> assert false)
  | At a -> Gd e
  | New -> Gd New
  | Nu(<<a>>e1) -> nu a (eval e1 v)
  | Eq(e1, e2, t') -> 
      eq t' (eval e1 v) (eval e2 v)
  | Swap(e1, e2, e3) -> 
      ((eval e1 v) <-> (eval e2 v)) (eval e3 v)
  | V e1 -> 
      (match eval e1 v with 
	   Gd e' -> Gd(V e')
	 | Fn _ -> assert false)
  | A(e1, e2) ->
      (match (eval e1 v, eval e2 v) with 
	   (Gd e1', Gd e2') -> Gd(A(e1', e2'))
	 |  _ -> assert false)
  | L(<<a>>e1) ->
      (match eval e1 v with
	   Gd e' -> Gd(L(<<a>>e'))
	 | Fn _ -> assert false)
  | Trec(e1, e2, e3, e4) -> 
      trec t (eval e1 v) (eval e2 v) (eval e3 v) (eval e4 v);;

(* Normalization *)
let norm (e : exp) : exp =
  let t = typeOf e in reify t (eval e init);;

(*********************************************************************
Examples 
*********************************************************************)

let  x : var = newvar Atm;;
let  y : var = newvar Trm;;
let  z : var = newvar Trm;;
let x1 : var = newvar Atm;;
let y1 : var = newvar Trm;;
let y2 : var = newvar Trm;;
let z1 : var = newvar Trm;;
let z2 : var = newvar Trm;;
let  a : atm = fresh;;
let a' : atm = fresh;;

let e1 : exp = 
  (*  "\x1. (x = x1)_Trm y (V x1)" *)
  Lam(<<x1>>(App(App(Eq(Var x, Var x1, Trm), Var y), V(Var x1))));;

let e2 : exp =
(* "\y1.\y2.\z1.\z2. A z1 z2" *)
  Lam(<<y1>>(Lam(<<y2>>(Lam(<<z1>>(Lam(<<z2>>(A(Var z1, Var z2)))))))));;

let e3 : exp =
(* "\x1.\y1.\z1.La.(a x1)*z1" *)
  Lam(<<x1>>(Lam(<<y1>>(Lam(<<z1>>(makeL (Var x1) (Var z1)))))));;

(* Capture-avoiding substitution *)
let sub : exp =
  (* "\x.\y.\z.trec e1 e2 e3 z" *)
  Lam(<<x>>(Lam(<<y>>(Lam(<<z>>(Trec(e1, e2, e3, Var z)))))));;

let testTypeOfSub : bool =
  (typeOf sub) = (Atm => (Trm => (Trm => Trm)));; (* true *)

let subEg : exp =
  (* "sub a' (V a) (La. V a')" *)
  App(App(App(sub, At a'), V(At a)), L(<<a>>(V(At a'))));;

let subEgN : exp =
  (* "La'. V a" *)  
  L(<<a'>>(V(At a)));;

let testNormEg : bool =
  (norm subEg) = subEgN;; (* true *)    

(* Addition *)
let plus : exp =
  (* "\n.\m.nrec n (\n.\m.S m) m" *)
  let n = newvar Nat in
  let m = newvar Nat in
    Lam(<<n>>
	  (Lam(<<m>>
		 (Nrec(Var n, 
		       Lam(<<n>>(Lam(<<m>>(S(Var m))))), 
		       Var m)))));;

(* One plus one equals two *)
let plusTest : bool =
  norm (App(App(plus, S Z), S Z)) = S(S Z);; (* true *)

(* Length of a lambda-term *)
let len : exp =
  (* "\x.trec (\x1.S 0)(\y.\z.\n.\m.plus n m)(\x1.\y.\n.S n) x" *)
  let  x = newvar Trm in
  let x1 = newvar Atm in
  let  y = newvar Trm in
  let  z = newvar Trm in
  let  n = newvar Nat in
  let  m = newvar Nat in
    Lam(<<x>>
	  (Trec(Lam(<<x1>>(S Z)), 
		Lam(<<y>>
		      (Lam(<<z>>
			     (Lam(<<n>>
				    (Lam(<<m>>
					   (App(App(plus, Var n), 
						Var m))))))))),
		Lam(<<x1>>(Lam(<<y>>(Lam(<<n>>(S(Var n))))))),
		Var x)));;

let testTypeOfLen : bool =
  (typeOf len) = (Trm => Nat);; (* true *)

let lenEg =
  (* "len(La. va)" *)
  App(len, L(<<a>>(V(At a))));;

let testLenEg : bool =
  (norm lenEg) = S(S Z);; (* true *)

(* end of nomst.ml *)