Evaluate LambdaMu terms using CPS transform

LambdaMu.idr

import Data.List.Elem 

-- | Lets start by defining some boilerplate data-types.

-- | Contexts/Variable Environments 
data All : {a : Type} -> (p : a -> Type) -> List a -> Type where 
  Nil : All p [] 
  (::) : {x : a} -> (px : p x) -> (pxs : All p xs) -> All p (x :: xs)

lookup : All p ctx -> Elem x ctx -> p x
lookup (value :: rest) Here  = value 
lookup (value :: rest) (There later) = lookup rest later 

-- | Co-Contexts. A classical term outputs a disjunction of values. 
data Any : (p : k -> Type) -> List k -> Type where 
  This : {p : k -> Type} -> p x -> Any p (x :: xs)
  That : Any p xs -> Any p (x :: xs)

-- Monoidal functor on Co-contexts
merge : {p : k -> Type} -> {x, x', x'' : k} -> 
  (p x -> p x' -> p x'') -> Any p (x::xs) -> Any p (x'::xs) -> Any p (x''::xs)
merge op (This t1) (This t2) = This (op t1 t2)
merge op (This _) (That t2) = That t2
merge op (That t1) _ = That t1 

-- Inject a value into a given location
fromElem : {p : k -> Type} -> p a -> Elem a d -> Any p d
fromElem v Here = This v 
fromElem v (There t) = That (fromElem v t)

-- Extract a from singleton [a]
fromAny : {p : k -> Type} -> Any p [a] -> p a
fromAny (This a) = a 
fromAny (That p) impossible

-- | Continuations
Cont : Type -> Type -> Type 
Cont r a = (a -> r) -> r

KlCont : Type -> Type -> Type -> Type 
KlCont r a b = a -> Cont r b 

toCPS : a -> (a -> r) -> r
toCPS x f = f x 

-- Functor over Cont
mapC : (a -> b) -> Cont r a -> Cont r b 
mapC f c = (\br => c (br . f)) 

-- Monoidal functor over Cont 
ap : (a -> b -> c) -> ((a -> r) -> r) -> ((b -> r) -> r) -> (c -> r) -> r
ap op ar br k = ar (\x => br (\y => k (op x y)))

-- CallCC
cc : ((forall b. a -> Cont r b) -> Cont r a) -> Cont r a
cc f = \k => (f $ \a, br => k a) k

-- Practice run : Evaluating STLC terms into Kleisli(Cont)
namespace STLC
  data Ty = N | Imp Ty Ty 

  data Term : List Ty -> Ty -> Type where
    Var : {g : List Ty} -> {a : Ty} -> Elem a g -> Term g a
    Lam : {g : List Ty} -> {a, b : Ty} -> Term (a::g) b -> Term g (Imp a b)
    App : {g : List Ty} -> {a, b : Ty} -> Term g (Imp a b) -> Term g a -> Term g b

    -- Some basic primitives
    KInt : Int -> Term g N 
    Plus : Term g N -> Term g N -> Term g N 

  EvTy : Type -> Ty -> Type
  EvTy r N = Int 
  EvTy r (Imp t1 t2) = EvTy r t1 -> Cont r (EvTy r t2) 

  -- Evaluate Terms of STLC into the Cont Monad, ie (eval: Term g t -> All [g] -> Cont r [t])
  eval : {r : Type} -> Term g t -> KlCont r (All (EvTy r) g) (EvTy r t)
  eval (Var v) env k = k (lookup env v) 
  eval (Lam t) env k = let ev = eval {r} t in k (\x => ev (x::env)) 
  eval (App t1 t2) env k = let 
    ev1 = eval t1 env 
    ev2 = eval t2 env in 
      ev1 (\x => ev2 (\y => x y k))  
  eval (KInt n) env k = k n
  eval (Plus t1 t2) env k = ap ((+)) (eval t1 env) (eval t2 env) k 

  -- Example term
  ex1 : Term [] N 
  ex1 = App (Lam (Plus (Var Here) (KInt 6))) (KInt 5)

  public export
  test : {r : Type} -> Cont r Int
  test = eval ex1 [] toCPS -- == 11
  
-- Evaluating LambdaMu into Kleisli(Cont)
namespace LambdaMu
  -- We add Bottom to types
  data Ty = N | Imp Ty Ty | Bot  

  --          Input Context  Output type   Co-Context
  data Term : List Ty ->     Ty ->         List Ty ->  Type where
    
    -- STLC Connectives
    Var   : {a : Ty} -> Elem a g -> Term g a d
    Lam   : {a, b :Ty} -> {g : List Ty} -> Term (a::g) b d -> Term g (Imp a b) d
    App   : {a : Ty} -> Term g (Imp a b) d -> Term g a d -> Term g b d
    Let : {s, t : Ty} -> Term g s d -> Term (s::g) t d -> Term g t d  

    -- Classical Connectives
    Mu    : Term g Bot (a::d) -> Term g a d              -- activate / catch / bottom elimination / proof by contradiction (a != Bot)
    Named : {a : Ty} -> Elem a d -> Term g a d -> Term g Bot d -- passivate / throw / bottom introduction / non-contradiction
    
    -- Basic Primitives
    KInt : Int -> Term g N d 
    Plus : Term g N d -> Term g N d -> Term g N d 

  EvTy : Type -> Ty -> Type
  EvTy r N = Int 
  EvTy r Bot = Void 
  EvTy r (Imp t1 t2) = EvTy r t1 -> (Cont r (EvTy r t2)) 

  -- Each term now consumes a context g of inputs, and produces a co-context of outputs 
  -- eval : Term g t d -> All [g] -> Cont r (Any (t::d))
  eval : {t : Ty} -> {r : Type} -> Term g t d -> KlCont r (All (EvTy r) g) (Any (EvTy r) (t::d))
  eval (Var v) env k = k (This (lookup env v))

  -- The two classical connectives can be figured out by just following the types
  eval (Mu t) env k = (eval t env) (\f => case f of 
    This n impossible
    That t => k t)
  eval (Named v t) env k = (eval t env) (\f => case f of 
    This n => k $ That (fromElem n v)
    That t => k $ That t)
  
  -- Let Binding
  eval (Let this inThis) env k = let 
    ev1 = eval this env in 
      ev1 (\f => case f of 
        This t => (eval inThis (t::env)) k
        That th => k (That th)) 

  -- Lambda application 
  eval (App {a,b} t1 t2) env k = let 
    ev1 = eval t1 env 
    ev2 = eval t2 env in 
      ev1 (\x => ev2 (\y => case (x, y) of 
        (This t1, This t2) => ?x1 t1 t2 -- x1: (EvTy r a -> (EvTy r b -> r) -> r)) -> EvTy r a -> r
        (This t1, That t2) => (k (That t2))
        (That t1, _) => k (That t1)))
  
  -- Lambda abstraction requires backtracking. Not 100% sure this works
  eval (Lam t) env k = 
    k (This (\a, b => (eval {r} t) (a::env) $ \c => case c of 
        This t => b t
        That th => k (That th)))

  eval (KInt n) env k = k (This n) 
  eval (Plus t1 t2) env k = ap (merge (+)) (eval t1 env) (eval t2 env) k   

  ex0 : Term [] N d
  ex0 = Plus (KInt 6) (KInt 7)

  ex1 : Term [] N d
  ex1 = App (Lam (Plus (Var Here) (KInt 6))) (KInt 5)

  public export
  test1 : {r : Type} -> Cont r Int
  test1 = mapC fromAny (eval ex1 []) toCPS -- == 11