LF_dB_PS.ml

open Format


type tm = Var of int
        | App of tm * tm
        | Lam of tm

type ty = U
        | El of tm
        | Pi of ty * ty

(* We need to know the type of the terms for reification! *)
type sub = NilS
         | ConsS of sub * (ty * tm)

(* contexts are of the form
   ConsC(U, ConsC(Pi(El(Var 0),U), ConsC(El(Var 1), NilC))) *)
type con = NilC
         | TyC of ty
         | TmC of ty * tm
         | ConsC of ty * con

let rec pp_tm_ k ppf t =
  match t with
  | Var x -> fprintf ppf "x%d" (k - x)
  | Lam s -> fprintf ppf "@[<1>(λx%d. %a)@]" (k + 1) (pp_tm_ (k + 1)) s
  | App(t,u) -> fprintf ppf "@[<1>(%a %a)@]" (pp_tm_ k) t (pp_tm_ k) u
let pp_tm ppf (t : tm) = (pp_tm_ 0 ppf t)

let rec pp_sub_ k ppf (s : sub) =
  match s with
  | NilS -> fprintf ppf "ε"
  | ConsS(s,(_,t)) -> fprintf ppf "⟨%a , %a⟩" (pp_sub_ k) s (pp_tm_ k) t
let pp_sub ppf s = pp_sub_ 0 ppf s

let rec pp_ty_ k ppf a =
  match a with
  | U -> fprintf ppf "U"
  | El a -> fprintf ppf "(El %a)" (pp_tm_ k) a
  | Pi(a,fam) -> fprintf ppf "(Pi[x%d:%a] %a)" (k + 1) (pp_ty_ k) a (pp_ty_ (k + 1)) fam
let pp_ty ppf (a : ty) = (pp_ty_ 0 ppf a)

let rec con_len (ctx : con) : int =
  match ctx with
  | NilC | TmC _ | TyC _ -> 0
  | ConsC(_,ctx) -> (con_len ctx) + 1

let rec pp_con_ l k ppf (ctx : con) =
  match ctx with
  | NilC -> fprintf ppf "⊢@]"
  | TmC(a,t) -> fprintf ppf "@ ⊢  @[%a : %a@]@]" (pp_tm_ l) t (pp_ty_ l) a
  | TyC a -> fprintf ppf "@ ⊢  @[%a@]@]" (pp_ty_ l) a
  | ConsC(a,ctx) -> fprintf ppf "@ ▹ x%d:%a %a" (k + 1) (pp_ty_ k) a (pp_con_ l (k + 1)) ctx
let pp_con ppf ctx = (fprintf ppf "@[<3>"); (pp_con_ (con_len ctx) 0 ppf ctx)

(* type of normal/neutral terms *)
type nf = NLam of nf   (* Normal terms of type Pi *)
        | NeuU of ne   (* Normal terms of type U *)
        | NeuEl of ne  (* Normal terms of type El *)
and ne = Var_ of int
       | App_ of ne * nf

let rec nf_tm (t : nf) =
  match t with
  | NLam t -> Lam (nf_tm t)
  | NeuU t -> ne_tm t
  | NeuEl t -> ne_tm t
and ne_tm (t : ne) =
  match t with
  | Var_ k -> Var k
  | App_(t,u) -> App(ne_tm t, nf_tm u)

let pp_nf ppf (t : nf) = pp_tm ppf (nf_tm t)
let pp_ne ppf (t : ne) = pp_tm ppf (ne_tm t)


(****************************************************)
(* Type of weakenings                               *)
(****************************************************)
(* These are the morphisms in a category W,
   whose objects are contexts, and whose morphisms are generated by
    W_id : hom(Γ,Γ)
    W1 : hom(Γ, Δ) → hom(Γ×U, Δ)
    W2 : hom(U, Δ) → hom(Γ×U, Δ×U)
   Note that this is for a single base type.
 *)
type wk = W_id
        | W1 of wk
        | W2 of wk

let rec pp_wk ppf (w : wk) =
  match w with
  | W_id -> fprintf ppf "W_id"
  | W1 w -> fprintf ppf "(W1 %a)" pp_wk w
  | W2 w -> fprintf ppf "(W2 %a)" pp_wk w


(* Composition in W *)
(*  wk_o : hom(Γ,Δ) → hom(Δ,Ξ) → hom(Γ,Ξ)  *)
let rec wk_o (w1 : wk) (w2 : wk) : wk =
  match (w1, w2) with
  | W_id,  _ -> w2
  | W1 w1, w2 -> W1 (wk_o w1 w2)
  | W2 w1, W_id -> W2 w1
  | W2 w1, W1 w2 -> W1 (wk_o w1 w2)
  | W2 w1, W2 w2 -> W2 (wk_o w1 w2)


(* Compute the pullback t[w]  *)
let rec wk_nf (w : wk) (t : nf) : nf =
  match (w,t) with
  | W_id, _ -> t
  | _,    NLam s -> NLam (wk_nf (W2 w) s)
  | _,    NeuU t -> NeuU (wk_ne w t)
  | _,    NeuEl t -> NeuEl (wk_ne w t)
and wk_ne (w : wk) (t : ne) : ne =
  match (w,t) with
  | W_id, _ -> t
  | _,    App_(t,u) -> App_(wk_ne w t, wk_nf w u)
  | W1 w, Var_ x -> wk_ne w (Var_ (x + 1))
  | W2 w, Var_ x -> (if x = 0 then (Var_ 0) else (wk_ne w (Var_ (x - 1))))


(****************************************************)
(* Values                                           *)
(****************************************************)
(* Term values *)
type vltm = UD of nf   (* value of type U *)
          | ElD of vltm * nf  (* value of type El a *)
          | PiD of (wk -> (vltm -> vltm)) (* value of type Pi *)
(* Type values *)
type vlty = VU
          | VEl of vltm
          | VPi of vlty * (wk -> (vltm -> vlty))
(* Context values *)
type vlcon = VNilC
           | VTyC of vlty
           | VTmC of vlty * vltm
           | VConsC of vlty * (wk -> (vltm -> vlcon))
(* Substitution values *)
type vlsub = VNilS
           | VConsS of vlsub * (vlty * vltm)



(* Pullback a value *)
let rec wk_vltm (w : wk) (u : vltm) : vltm =
  match u with
  | UD a ->  UD (wk_nf w a)
  | ElD(a,s) -> ElD(wk_vltm w a, wk_nf w s)
  | PiD f -> PiD (fun w' u -> f (wk_o w' w) u)  (* Should this be (wk_o w w') ?? *)

let wk_env (w : wk) (env : vltm list) : vltm list =
  List.map (wk_vltm w) env



(****************************************************)
(* Evaluation/Reification/Reflection                *)
(****************************************************)

let appD (u : vltm) (v : vltm) : vltm =
  match u with
  | PiD f -> f W_id v
  | _ -> failwith "Not a lambda!"


let rec eval_tm (t : tm) (env : vltm list) : vltm =
  match t with
  | Var x -> List.nth env x
  | Lam s -> PiD (fun w u -> eval_tm s (u::(wk_env w env)))
  | App(t,u) -> appD (eval_tm t env) (eval_tm u env)

let rec eval_ty (a : ty) (env : vltm list) : vlty =
  match a with
  | U -> VU
  | El t -> VEl (eval_tm t env)
  | Pi(a,fam) -> VPi(eval_ty a env, (fun w u -> eval_ty fam (u::(wk_env w env))))

let rec eval_con (ctx : con) (env : vltm list) : vlcon =
  match ctx with
  | NilC -> VNilC
  | TyC a -> VTyC(eval_ty a env)
  | TmC (a,t) -> VTmC(eval_ty a env, eval_tm t env)
  | ConsC(a,ctx) -> VConsC(eval_ty a env, (fun w u -> eval_con ctx (u::(wk_env w env))))

let rec eval_sub (g : sub) (env : vltm list) : vlsub =
  match g with
  | NilS -> VNilS
  | ConsS(g,(a,t)) -> VConsS(eval_sub g env, (eval_ty a env, eval_tm t env))

(* Should a be a type? *)
let rec reify_tm (a : vlty) (u : vltm) : nf =
  match (a,u) with
  | VU, UD a -> a
  | VEl _, ElD(_,t) -> t
  | VPi(a,fam), PiD f -> (let v = (reflect_tm a (Var_ 0)) in
                          NLam (reify_tm (fam (W1 W_id) v) (f (W1 W_id) v)))
  | _ -> failwith "Failure in reify!"

and reify_ty (a : vlty) : ty =
  match a with
  | VU -> U
  | VEl a -> El (nf_tm (reify_tm VU a))
  | VPi(a,fam) -> Pi(reify_ty a, (let v = reflect_tm a (Var_ 0) in
                                  (reify_ty (fam (W1 W_id) v))))

and reify_con (ctx : vlcon) : con =
  match ctx with
  | VNilC -> NilC
  | VTyC a -> TyC (reify_ty a)
  (* Shouldn't I need to reify a before passing it to reify_tm? *)
  | VTmC(a,t) -> TmC(reify_ty a, (nf_tm (reify_tm a t)))
  | VConsC(a,ctx) -> ConsC(reify_ty a, (let v = reflect_tm a (Var_ 0) in
                                        (reify_con (ctx (W1 W_id) v))))

and reify_sub (sub : vlsub) : sub =
  match sub with
  | VNilS -> NilS
  | VConsS(sub,(a,t)) -> ConsS(reify_sub sub, (reify_ty a, nf_tm (reify_tm a t)))

and reflect_tm (a : vlty) (t : ne) : vltm =
  match (a,t) with
  | VPi(a,fam), t -> PiD(fun w u -> reflect_tm (fam w u) (App_(wk_ne w t, reify_tm a u)))
  | VU, t -> UD (NeuU t)
  | VEl a, t -> ElD(a,NeuEl t)



let nbe_tm (a : ty) (t : tm) : nf =
  reify_tm (eval_ty a []) (eval_tm t [])

let nbe_ty (a : ty) : ty =
  reify_ty (eval_ty a [])

let nbe_con (ctx : con) : con =
  reify_con (eval_con ctx [])

let nbe_sub (g : sub) : sub =
  reify_sub (eval_sub g [])

(****************************************************************)
(* Tests                                                        *)
(****************************************************************)


let _I = Lam (Var 0)
let _K = Lam (Lam (Var 1))
(* (A -> (B -> C)) -> (A -> B) -> A -> C *)
let _S = Lam (Lam (Lam (App(App(Var 2, Var 0),App(Var 1, Var 0)))))

let _UU = Pi(U,U)

let tests
  = [(_I, Pi(U,U));
     (_K, Pi(U,Pi(U,U)));
     (Lam(App(_I,Var 0)), Pi(U,U));
     (App(_K,_I), Pi(U,Pi(U,U)));
     (App(Lam (Lam (Lam (App(Var 2, Var 0)))), _I), Pi(U,Pi(U,U)));

     (_S, Pi(Pi(U,Pi(U,U)),Pi(Pi(U,U),Pi(U,U))));
     (Lam (* A : U *) (Lam (* x : El A *) (Var 1)), Pi(U,Pi(El (Var 0), U)));
     (Lam (* A : U *) (Lam (* x : El A *) (Var 0)), Pi(U,Pi(El (Var 0), El (Var 1))));
     (Lam (* A : U *) (Lam (* B : El A -> U *) (Var 1)), Pi(U,Pi(Pi(El(Var 0),U),U)));
     (Lam(Lam(App(Var 0, Var 1))), Pi(U,Pi(Pi(U,U), U)));

     (Lam (* A : U *) (Lam (* B : El A -> U *) (Lam (* x : A *) (App(Var 1, Var 0)))), Pi(U,Pi(Pi(El(Var 0),U), Pi(El(Var 1), U))));
     (Lam (* A : U *) (Lam (* B : El A -> U *) (Lam (* C : (x : A) -> (B x) -> U *) (Lam (* x : A *) (Lam (* y : B x *) (App(App(Var 2, Var 1),Var 0)))))),
      Pi(U, Pi(Pi(El(Var 0),U), Pi(Pi(El(Var 1),Pi(El(App(Var 1,Var 0)),U)), Pi(El(Var 2), Pi(El(App(Var 2,Var 0)), U))))));
     (Lam (* A : U *) (Lam (* x : El A *) (Var 0)), Pi(U, Pi(El (App(_I,Var 0)), El (Var 1))));
     (Lam (* A : U *) (Lam (* B : El A -> El A -> U *) (Lam (* x : El A *) (App(App(Var 1,Var 0),Var 0)))),
      Pi(U, Pi(Pi(El(Var 0), Pi(El(Var 1),U)), Pi(El(Var 1), U))));
     (Lam (* A : U *) (Lam (* B : El A -> El A -> U *) (Lam (* x : El A *) (Lam (* y : El A *) (App(App(Var 2,Var 1),Var 0))))),
      Pi(U, Pi(Pi(El(Var 0),Pi(El(Var 1),U)), Pi(El(Var 1), Pi(El(Var 2),U)))));

     (Lam (Lam (Lam (App(App(Var 2,Var 0),Var 0)))), Pi(Pi(U,Pi(U,U)),Pi(U,Pi(U,U))));
     (Lam (* N : U *) (Lam (* 0 : El N*) (Lam (* S : El N -> El N *) (App(Var 0,App(Var 0,App(Var 0,Var 1)))))), Pi(U, Pi(El(Var 0), Pi(Pi(El(Var 1),El(Var 2)), El(Var 2)))))
    ]

let sigs
  = [ConsC(U,NilC);
     ConsC(U, ConsC(El(Var 0), NilC));
     ConsC(U, ConsC(El(Var 0), ConsC(El(Var 1), NilC)));
     ConsC(U, ConsC(El(Var 0), ConsC(El(Var 1), ConsC(El(Var 2), NilC))));
     ConsC(U, ConsC(El(Var 0), ConsC(El(Var 1), ConsC(El(Var 2), ConsC(El(Var 3), NilC)))));
     ConsC(U, ConsC(Pi(El(Var 0),U), ConsC(Pi(El(Var 1),Pi(El(App(Var 1,Var 0)),U)), ConsC(El(Var 2), ConsC(El(App(Var 2,Var 0)), NilC)))));
     ConsC(U,ConsC(El(Var 0),ConsC(Pi(El(Var 1),El(Var 2)), TmC(El(App(_I,Var 2)), App(Var 0,App(Var 0, App(Var 0, Var 1)))))));
     ConsC(U, ConsC(El(App(App(_K,Var 0),_I)), ConsC(El(Var 1), NilC)));
     ConsC(U (* Sig : Type *),
      ConsC(Pi(El(Var 0), U), (* Ty : Sig -> Type *)
      ConsC(Pi(El(Var 1),Pi(El(App(Var 1,Var 0)), U)), (* Tm : (Γ : Sig) -> (Ty Γ) -> Type *)
      ConsC(Pi(El(Var 2), Pi(El(Var 3), U)), (* Sub : Sig -> Sig -> Type *)
      ConsC(El(Var 3), (* · : Sig *)
      ConsC(Pi(El(Var 4),Pi(El(App(Var 4, Var 0)), El(Var 6))), (* ▹ : (Γ : Sig) -> (Ty Γ) -> Sig *)
      ConsC(Pi(El(Var 5),Pi(El(Var 6),Pi(El(App(Var 6,Var 0)),Pi(El(App(App(Var 5,Var 2),Var 1)), El(App(Var 8, Var 3)))))), (* _[_]T *)
      ConsC(Pi(El(Var 6),Pi(El(Var 7),Pi(El(App(Var 7,Var 0)),Pi(El(App(App(Var 7,Var 1),Var 0)),
            Pi(El(App(App(Var 7,Var 3),Var 2)), El(App(App(Var 9,Var 4),App(App(App(App(Var 5,Var 4),Var 3),Var 2),Var 0)))))))), (* _[_]t *)
      ConsC(Pi(El(Var 7),El(App(App(Var 5, Var 0),Var 0))), (* id : (Σ : Sig) -> (Sub Σ Σ) *)
      ConsC(Pi(El(Var 8),Pi(El(Var 9),Pi(El(Var 10), Pi(El(App(App(Var 8,Var 2),Var 1)), Pi(El(App(App(Var 9,Var 2),Var 1)), El(App(App(Var 10, Var 4), Var 2))))))), (* _∘_ *)
      ConsC(Pi(El(Var 9),El(App(App(Var 7, Var 0),Var 6))), (* ε *)
            NilC)))))))))));
    ]


let _ =
  for i=0 to (List.length tests) - 1 do
    (let p = (List.nth tests i) in
     (printf "(%d)@\n%a@ : %a@\n" i pp_tm (fst p) pp_ty (snd p));
     (printf "%a@ : %a@\n" pp_nf (nbe_tm (snd p) (fst p)) pp_ty (nbe_ty (snd p))))
  done

let _ =
  for i=0 to (List.length sigs) - 1 do
    (let p = (List.nth sigs i) in
     (printf "(%d)@\n%a@\n" i pp_con p);
     (printf "%a@\n" pp_con (nbe_con p)))
  done