Normalization by Evaluation written in Beluga

nbe.bel

LF tp : type =
| base : tp
| arr  : tp → tp → tp
;

LF term : tp → type =
| app : term (arr a b) → term a → term b
| lam : (term a → term b) → term (arr a b)
;

% we use old-style definition because
% we need both `neut` and `norm` at the same time
neut : tp → type.
norm : tp → type.
nlam : (neut a → norm b) → norm (arr a b).
rapp : neut (arr a b) → norm a → neut b.
embed : neut base → norm base.

schema ctx = some [a:tp] block x:neut a;

stratified Sem : {g : ctx} [ ⊢ tp] -> ctype =
| Base : [g ⊢ norm base] → Sem [g] [ ⊢ base]
| Arr : {g : ctx} ({h : ctx} {#S: [h ⊢ g]} Sem [h] [ ⊢ A] → Sem [h] [ ⊢ B])
                  → Sem [g] [ ⊢ arr A B]
;

rec sem_wkn : {h : ctx} {g : ctx} {#S : [h ⊢ g]} Sem [g] [ ⊢ A] → Sem [h] [ ⊢ A]  =
  mlam h ⇒ mlam g ⇒ mlam S ⇒ fn e ⇒ case e of 
  | Base [g ⊢ R] ⇒ Base [h ⊢ R[#S]]
  | Arr [g] f ⇒ Arr [h] (mlam h' ⇒ mlam S' ⇒ f [h'] [h' ⊢ #S[#S']])
;

schema tctx = some [t : tp] block x : term t;

typedef Env : {g : tctx} {h : ctx} ctype =
  {T : [ ⊢ tp]} {#p : [g ⊢ term T[]]} Sem [h] [ ⊢ T]
;

rec env_ext : Env [g] [h] → Sem [h] [ ⊢ S] → Env [g, x : term S[]] [h]  =
  fn env ⇒ let env : Env [g] [h] = env
    in fn sem ⇒ mlam T ⇒ mlam p ⇒ case [g, x : term _ ⊢ #p] of 
       | [ g, x : term S ⊢ x ] ⇒ sem
       | [ g, x : term S ⊢ #q[..] ] ⇒ env [ ⊢ T ] [ g ⊢ #q ]
;

rec env_wkn : {h' : ctx} {h : ctx} {#W : [h' ⊢ h]} Env [g] [h] → Env [g] [h']  =
mlam h' ⇒ mlam h ⇒ mlam W ⇒ fn env ⇒ 
  let env : Env [g] [h] = env
  in mlam T ⇒ mlam p ⇒ sem_wkn [h'] [h] [h' ⊢ #W] (env [ ⊢ T] [g ⊢ #p])
;

rec eval : [g ⊢ term S[]] → Env [g] [h] → Sem [h] [ ⊢ S]  =
fn tm ⇒ fn env ⇒
  let env : Env [g] [h] = env
  in case tm of
     | [g ⊢ #p] ⇒ env [ ⊢ _] [g ⊢ #p]
     | [g ⊢ lam (\x. E)] ⇒
       Arr [h] (mlam h' ⇒ mlam W ⇒ fn sem ⇒
                  let extEnv = (env_ext (env_wkn [h'] [h] [h' ⊢ #W] env) sem)
                  in eval [g, x : term _ ⊢ E] extEnv)
     | [g ⊢ app E1 E2] ⇒ let Arr [h] f = eval [g ⊢ E1] env
                           in f [h] [h ⊢ ..] (eval [g ⊢ E2] env)
;

rec app' : (g:ctx) {R : [g ⊢ neut (arr T[] S[])]} [g ⊢ norm T[]] → [g ⊢ neut S[]]  =
mlam R ⇒ fn n ⇒ let [g ⊢ N] = n in [g ⊢ rapp R N];

rec reify : Sem [h] [ ⊢ A] → [h ⊢ norm A[]] =
fn sem ⇒ case sem of
    | Base [h ⊢ R] ⇒ [h ⊢ R]
    | Arr [g] f ⇒ let  [g, x : neut _ ⊢ N] =
                   reify (f [g, x : neut _] [g, x ⊢ ..] (reflect [g, x : neut _ ⊢ x]))
                   in [g ⊢ nlam (\x. N)]
and reflect : [h ⊢ neut A[]] → Sem [h] [ ⊢ A] =
fn n ⇒ let [h ⊢ R] : [h ⊢ neut A[]] = n
        in case [ ⊢ A] of
            | [ ⊢ base ] ⇒ Base [h ⊢ embed R]
            | [ ⊢ arr B C ] ⇒ Arr [h] (mlam h' ⇒ mlam W ⇒ fn tm ⇒
                                           reflect (app' [h' ⊢ R[#W]] (reify tm)))
;

rec nbe : Env [g] [h] → [g ⊢ term A[]] → [h ⊢ norm A[]] =
fn env ⇒ fn tm ⇒ reify (eval tm env)
;