Evaluator for STLC in Polarity, using HOAS in values

eval_stlc.pol

use "../std/data/bool.pol"
use "../std/data/nat.pol"

data Ty {
    FunTy(a b: Ty),
    NatTy,
    BoolTy,
}

data Ctx {
    Empty,
    Extend(ctx: Ctx, a: Ty),
}

data Index(ctx: Ctx, a: Ty) {
    Stop(ctx: Ctx, a: Ty): Index(Extend(ctx, a), a),
    Pop(ctx: Ctx, a b: Ty, i: Index(ctx, a)): Index(Extend(ctx, b), a),
}

data Expr(ctx: Ctx, a: Ty) {
    Let(ctx: Ctx, a b: Ty, defn: Expr(ctx, a), body: Expr(Extend(ctx, a), b)): Expr(ctx, b),
    Var(ctx: Ctx, a: Ty, i: Index(ctx, a)): Expr(ctx, a),
    FunLit(ctx: Ctx, a b: Ty, body: Expr(Extend(ctx, a), b)): Expr(ctx, FunTy(a, b)),
    FunApp(ctx: Ctx, a b: Ty, fn: Expr(ctx, FunTy(a, b)), arg: Expr(ctx, a)): Expr(ctx, b),
    NatLit(ctx: Ctx, n: Nat): Expr(ctx, NatTy),
    BoolLit(ctx: Ctx, b: Bool): Expr(ctx, BoolTy),
    PrimExpr(ctx: Ctx, a: Ty, prim: Prim(a)): Expr(ctx, a),
}

data Prim(a: Ty) {
    NatAdd: Prim(FunTy(NatTy, FunTy(NatTy, NatTy))),
    NatFact: Prim(FunTy(NatTy, NatTy)),
}

data Value(a: Ty) {
    VFunLit(a b: Ty, clos : Clos(a, b)): Value(FunTy(a, b)),
    VNatLit(n: Nat): Value(NatTy),
    VBoolLit(b: Bool): Value(BoolTy),
}

codata Clos(a b: Ty) {
    Clos(a, b).closAp(a b: Ty, arg: Value(a)): Value(b),
}

def Value(FunTy(a, b)).asFunLit(a b: Ty): Clos(a, b) {
    VFunLit(_, _, clos) => clos,
    VNatLit(_) absurd,
    VBoolLit(_) absurd,
}

def Value(NatTy).asNatLit: Nat {
    VFunLit(_, _, _) absurd,
    VNatLit(n) => n,
    VBoolLit(_) absurd,
}

data Env(ctx: Ctx) {
    EEmpty: Env(Empty),
    EExtend(ctx: Ctx, a: Ty, env: Env(ctx), x: Value(a)): Env(Extend(ctx, a)),
}

def Index(ctx, a).evalIndex(ctx: Ctx, a: Ty, env: Env(ctx)): Value(a) {
    Stop(_, _) => env.match { 
        EEmpty absurd,
        EExtend(_, _, _, x) => x,
    },
    Pop(_, _, _, i) => env.match { 
        EEmpty absurd,
        EExtend(ctx, _, env, _) => i.evalIndex(ctx, a, env),
    },
}

def Prim(a).evalPrim(a: Ty): Value(a) {
    NatAdd => 
        VFunLit(NatTy, FunTy(NatTy, NatTy), \closAp(_, _, x) =>
            VFunLit(NatTy, NatTy, \closAp(_, _, y) => VNatLit(x.asNatLit.add(y.asNatLit)))),
    NatFact => VFunLit(NatTy, NatTy, \closAp(_, _, x) => VNatLit(x.asNatLit.fact)),
}

def Expr(ctx, a).evalExpr(ctx: Ctx, a: Ty, env: Env(ctx)): Value(a) {
    Let(ctx, a, b, defn, body) => 
        body.evalExpr(Extend(ctx, a), b, 
            EExtend(ctx, a, env, defn.evalExpr(ctx, a, env))),
    Var(ctx, a, i) => 
        i.evalIndex(ctx, a, env),
    FunLit(ctx, a, b, body) => 
        VFunLit(a, b, \closAp(_, _, x) => 
            body.evalExpr(Extend(ctx, a), b, EExtend(ctx, a, env, x))),
    FunApp(ctx, a, b, fn, arg) => 
        fn.evalExpr(ctx, FunTy(a, b), env)
            .asFunLit(a, b)
            .closAp(a, b, arg.evalExpr(ctx, a, env)),
    NatLit(ctx, n) => VNatLit(n),
    BoolLit(ctx, b) => VBoolLit(b),
    PrimExpr(_, _, prim) => prim.evalPrim(a),
}