{-# LANGUAGE GADTs, ExistentialQuantification, TypeFamilies, EmptyDataDecls, ScopedTypeVariables, RankNTypes, TypeOperators #-}

module CPS where

-- Demonstration of type families with CPS translation

data Void -- necessary for Halt

data Var g t where
   Z :: Var (g, t) t
   S :: Var g t -> Var (g,t') t

data Exp g t where 
    Var  :: Var g t -> Exp g t
    Lit  :: Int ->  Exp g Int
    Lam  :: (Exp (g, t1) t2) -> Exp g (t1 -> t2)
    App  :: Exp g (t1 -> t2) -> Exp g t1 -> Exp g t2
    Halt :: Exp g t -> Exp g Void

data Env g where
    ENil  :: Env ()
    ECons :: Env g -> t -> Env (g, t)

instance Show (Exp g n) where
  show (Var v)  = "Var"
  show (Lit i)   = "(Lit " ++ show i ++ ")"
  show (Lam t)   = "(Lam " ++ show t ++ ")"
  show (App t u) = "(App " ++ show t ++ " " ++ show u ++ ")"
  show (Halt t)  = "(Halt " ++ show t ++ ")"

sLookup :: Var env t -> Env env -> t
sLookup Z     (ECons g v) = v
sLookup (S x) (ECons g v) = sLookup x g

interp :: Env env -> Exp env t ->  t
interp e (Var x)     = sLookup x e
interp e (Lit i)     = i
interp e (Lam t)     = \x -> interp (ECons e x) t
interp e (App t1 t2) = (interp e t1) (interp e t2)
interp e (Halt t)    = error "Machine halted"

--weakening...
shift1 :: Exp g t -> Exp (g, t') t
shift1 = undefined

-- type of continuation
data Cont t

-- CPS coercion
type family CPS t
type instance CPS Int = Int
type instance CPS (a -> b) = CPS a -> (CPS (Cont b)) -> Void
type instance CPS (a,b) = (CPS a, CPS b)
type instance CPS () = ()
type instance CPS (Cont b) = CPS b -> Void

-- standard CPS conversion

-- [[  x  ]]  c = App c x
-- [[ \x.t ]] c = App c (Lam x. Lam c'.[[t]]c')
-- [[ t1 t2]] c = [[t1]] (Lam v1. [[t2]] (Lam v2. (App (App v1 v2) c)))
-- [[ halt t]] c = P [[ t ]] 

-- P[[ t ]] = [[ t ]] (Lam x.halt x)

initk :: Exp g (t -> Void)
initk = (Lam (Halt (Var Z)))

cps_prog :: Exp g t -> Exp (CPS g) Void
cps_prog t = cps t initk

cpsvar :: Var g t -> Var (CPS g) (CPS t)
cpsvar Z = Z 
cpsvar (S x) = S (cpsvar x)

cps :: Exp g t -> (Exp (CPS g) (CPS t -> Void)) -> Exp (CPS g) Void 
cps (Var x) c = App c (Var (cpsvar x))
cps (Lit i) c = App c (Lit i)