module cps where

-- Demonstration of CPS translation
-- uses de Bruijn representation of binding

open import Data.Bool
open import Data.Nat
open import Data.Unit
open import Data.Empty
open import Data.List
open import Data.Maybe
open import Data.Product hiding (map)
open import Relation.Binary.Core
open import Data.String hiding (_++_)


data Typ : Set where
  Nat  : Typ
  _~>_ : Typ -> Typ -> Typ  
  Cont : Typ -> Typ  -- type of continuation

i : Typ -> Set 
i Nat        = ℕ
i (t1 ~> t2) = i t1 -> i t2
i (Cont t)   = ⊥    -- arbitrary

Ctx = List Typ

data Var : Ctx -> Typ -> Set where
  VZ : { Γ : Ctx } { t : Typ } -> Var (t ∷ Γ) t
  VS : { Γ : Ctx } { t t' : Typ } -> Var Γ t -> Var ( t' ∷ Γ ) t

data Exp : Ctx -> Typ -> Set where 
  EVar  : { Γ : Ctx } { t : Typ } -> Var Γ t -> Exp Γ t
  ELit  : { Γ : Ctx } -> ℕ -> Exp Γ Nat
  ELam  : { Γ : Ctx } { t1 t2 : Typ } -> 
             (Exp ( t1 ∷ Γ ) t2) -> Exp Γ (t1 ~> t2)
  EApp  : { Γ : Ctx } {t1 t2 : Typ } -> 
             Exp Γ (t1 ~> t2) -> Exp Γ t1 -> Exp Γ t2
  EHalt : { Γ : Ctx } {t : Typ} -> Exp Γ Nat -> Exp Γ t

------------------------------------------------------------------------

-- An interpreter

data Env : Ctx -> Set where
  ENil  : Env []
  ECons : forall { g t } -> Env g -> i t -> Env (t ∷ g )

sLookup : ∀ {Γ t} -> Var Γ t -> Env Γ -> i t
sLookup VZ     (ECons g v) = v
sLookup (VS x) (ECons g v) = sLookup x g

-- How to complete interp ??

data HaltFree : forall {g t} -> Exp g t -> Set where
  HFN : forall {g i} -> HaltFree {g} (ELit i)
  HFL : forall {g t t1}{ e : Exp (t1 ∷ g) t} -> HaltFree e -> HaltFree (ELam e)
  HFV : forall {g t x} -> HaltFree {g}{t} (EVar x)
  HFA : forall {g t t1} {e1 : Exp g (t1 ~> t)} {e2 : Exp g t1} -> 
          HaltFree e1 -> HaltFree e2 -> HaltFree (EApp e1 e2)

interp : ∀ {G t} -> (e : Exp G t) -> (pf : HaltFree e) -> Env G -> i t
interp (ELit i)     HFN env = i
interp (ELam e)     (HFL pf) env = (\ y -> interp e pf (ECons env y))
interp (EVar x)     pf env = sLookup x env
interp (EApp e1 e2) (HFA p1 p2) env = (interp e1 p1 env) (interp e2 p2 env)
interp (EHalt e)    () env

-----------------------------------------------------------------------
-- interpreter examples
t1 : Exp [] Nat
t1 = EApp (EHalt (ELit 3)) (ELit 4)

-- v1 = interp t1 _ ENil

-----------------------------------------------------------------------
-- weakening...

weakvar : forall { g t } -> (t' : Typ) -> (g' : Ctx) -> 
             Var (g' ++ g) t -> Var (g' ++ (t' ∷ g)) t
weakvar t' [] x = VS x
weakvar t' (t1 ∷ g'') VZ = VZ
weakvar t' (t1 ∷ g'') (VS x) = VS (weakvar t' g'' x)
               
weak : forall {g' t g} -> 
          Exp (g' ++ g) t -> (t' : Typ) -> Exp (g' ++ (t' ∷ g)) t
weak {g'} (EVar x) t' = EVar (weakvar t' g' x)
weak {g'}{t1 ~> t2} (ELam u) t' = ELam (weak { (t1 ∷ g') } u t') 
weak {g'}{t2}{g}
   (EApp .{g' ++ g}{t1}.{t2} e1 e2) t' = 
    EApp (weak {g'}{t1 ~> t2}{g} e1 t') 
         (weak {g'}{t1}{g} e2 t')
weak (ELit i) t' = ELit i
weak {g'} (EHalt u) t' = EHalt (weak {g'} u t')

shift1 : forall { g t } -> ( t' : Typ) -> Exp g t -> Exp (t' ∷ g) t
shift1  t' e = weak { [] } e t'

-----------------------------------------
-- Result type 
-- not Void this time b/c we don't have Halt in Agda
RT : Typ
RT = Nat

-- CPS type conversion
CPS : Typ -> Typ
CPS Nat      = Nat
CPS (a ~> b) = CPS a ~> ((CPS b ~> RT) ~> RT)
CPS (Cont b) = CPS b ~> RT

{-
-- CPS conversion

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

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

cpsvar : forall { g t } -> Var g t -> Var (map CPS g) (CPS t)
cpsvar VZ = VZ
cpsvar (VS x) = (VS (cpsvar x)) 

mutual

  cps_prog : forall { g } -> Exp g Nat -> (Exp (map CPS g) RT) 
  cps_prog e = cps e (ELam (EVar VZ)) 

  cps : forall { g t } -> Exp g t -> (Exp (map CPS g) (CPS t ~> RT)) -> 
          Exp (map CPS g) RT 

  cps (EVar x) k = EApp k (EVar (cpsvar x))

  cps (ELit i) k = EApp k (ELit i)

  cps (EHalt e) k = cps_prog e

  cps { _ } { _ ~> t2 } (ELam e) k = 
     EApp k (ELam (ELam (cps (shift1 (Cont t2) e) (EVar VZ))))

  cps { g1 } { t2 } (EApp .{_}{t1}.{_} e1 e2) k = 
   let k2 : _ -- Exp ((CPS g, CPS (t1 -> t2)), CPS t1) (CPS (Cont t2)) 
       k2 =  shift1 _ (shift1 _ k) in
   let k1 : _ -- Exp (CPS g, CPS (t1 -> t2)) (CPS (Cont t1))
       k1 =  ELam (EApp (EApp (EVar (VS VZ)) (EVar VZ)) k2) in
   let e2' : _ -- Exp (g, t1 -> t2) t1
       e2' = shift1 (t1 ~> t2) e2 in 
   cps e1 (ELam (cps e2' k1))

cps-prog-haltfree : forall {g} -> (e : Exp g Nat) -> HaltFree (cps_prog e)
cps-prog-haltfree = {!!}