An abstract machine for NBE

nbe-machine.rkt

#lang typed/racket
(define-type Term (U VAR LAM APP))
(struct: VAR ([x : String]) #:transparent)
(struct: LAM ([x : String] [m0 : Term]) #:transparent)
(struct: APP ([m1 : Term] [m2 : Term]) #:transparent)

(define-type Sem (U Ind Fun))

(define-type Env (U R1 R2))
(struct: R1 ())
(struct: R2 ([x : String] [d : Den] [ρ : Env]))
(define: (apply-env [ρ : Env] [x : String] [k : KS]) : Term
  (match ρ
    [(R1) (apply-ks k (I3 x))]
    [(R2 y d ρ) 
     (if (string=? x y) 
         (apply-den d k)
         (apply-env ρ x k))]))

(define-type Den (U D1 D2))
(struct: D1 ([n : Integer]))
(struct: D2 ([m2 : Term] [ρ : Env]))
(define: (apply-den [d : Den] [k : KS]) : Term
  (match d
    [(D1 n) (apply-ks k (I1 n))]
    [(D2 m2 ρ) (eval m2 ρ k)]))

(define-type Ind (U I1 I2 I3))
(struct: I1 ([n : Integer]))
(struct: I2 ([l : Ind] [m2 : Term] [ρ : Env]))
(struct: I3 ([x : String]))
(define-predicate ind? Ind)
(define: (apply-ind [i : Ind] [n : Integer] [k : KT]) : Term
  (match i
    [(I1 n) (apply-kt k (VAR (format "g~a" n)))]
    [(I2 l m2 ρ) (apply-ind l n (T2 m2 ρ n k))]
    [(I3 x) (apply-kt k (VAR x))]))

(define-type Fun (U F1))
(struct: F1 ([m0 : Term] [x : String] [ρ : Env]))

(define-type KT (U T1 T2 T3 T4))
(struct: T1 ())
(struct: T2 ([m2 : Term] [ρ : Env] [n : Integer] [k : KT]))
(struct: T3 ([k : KT] [n : Term]))
(struct: T4 ([k : KT] [n : Integer]))
(define: (apply-kt [k : KT] [m : Term]) : Term
  (match k
    [(T1) m]
    [(T2 m2 ρ n k) (eval m2 ρ (S2 n k m))]
    [(T3 k n) (apply-kt k (APP n m))]
    [(T4 k n) (apply-kt k (LAM (format "g~a" n) m))]))

(define-type KS (U S1 S2 S3 S4))
(struct: S1 ())
(struct: S2 ([n : Integer] [k : KT] [m : Term]))
(struct: S3 ([n : Integer] [k : KT]))
(struct: S4 ([k : KS] [m2 : Term] [ρ : Env]))
(define: (apply-ks [k : KS] [v : Sem]) : Term
  (match k
    [(S1) (down v 0 (T1))]
    [(S2 n k m) (down v n (T3 k m))]
    [(S3 n k) (down v (add1 n) (T4 k n))]
    [(S4 k m2 ρ)
     (match v
       [(? ind? l) (apply-ks k (I2 l m2 ρ))] 
       [(F1 m0 x ρ*) (eval m0 (R2 x (D2 m2 ρ) ρ*) k)])]))

(define: (down [s : Sem] [n : Integer] [k : KT]) : Term
  (match s   
    [(? ind? l) (apply-ind l n k)]
    [(F1 m0 x ρ) (eval m0 (R2 x (D1 n) ρ) (S3 n k))]))

(define: (eval [m : Term] [ρ : Env] [k : KS]) : Term
  (match m
    [(VAR x) (apply-env ρ x k)]
    [(LAM x m0) (apply-ks k (F1 m0 x ρ))]
    [(APP m1 m2) (eval m1 ρ (S4 k m2 ρ))]))

(define: (norm [m : Term]) : Term
  (eval m (R1) (S1)))

(norm (APP (LAM "s0" (LAM "z0" (APP (VAR "s0") (APP (VAR "s0") (VAR "z0")))))
           (APP (LAM "s1" (LAM "z1" (APP (VAR "s1") (APP (VAR "s1") (VAR "z1")))))
                (LAM "s2" (LAM "z2" (APP (VAR "s2") (APP (VAR "s2") (VAR "z2"))))))))