-- A typesystem and CPS interpreter for algebraic effects and handlers
-- (deep and shallow)

module algebraic2 where

open import Data.Nat using (ℕ; zero; suc; _+_)
open import Data.Bool using (true; false; if_then_else_) renaming (Bool to 𝔹)
open import Data.String using (String)
open import Data.Unit using (⊤; tt)
open import Data.Empty using (⊥)
open import Data.Product using (_×_; _,_)
open import Relation.Binary.PropositionalEquality
open Relation.Binary.PropositionalEquality.≡-Reasoning

-- Types
data Ty : Set
data Tr : Set
data Mc : Set
data H  : Set

data Ty where
  Nat  : Ty
  Bool : Ty
  Str  : Ty
  _⇒_⟨_,_⟩_⟨_,_⟩_ : Ty → Ty → Tr → Mc → Ty → Tr → Mc → Ty → Ty

data Tr where
  ● : Tr
  _⇨⟨_,_⟩_ : Ty → Tr → Mc → Ty → Tr

infix  _⇨⟨_,_⟩_

data Mc where
  [] : Mc
  _⇨⟨_,_⟩_×_×_::_ : Ty → Tr → Mc → Ty → Tr → H → Mc → Mc

infix  _⇨⟨_,_⟩_×_×_::_

data H where
  ● : H
  _⇒[_⇒⟨_,_⟩_]⇒_⇒_⟨_,_⟩_⟨_,_⟩_ :
    Ty → Ty → Tr → Mc → Ty → Tr → Ty → Tr → Mc → Ty → Tr → Mc → Ty → H

infix  _⇒_⟨_,_⟩_⟨_,_⟩_

-- cons and append of trail
compatible : Tr → Tr → Tr → Set
compatible ● μ₂ μ₃ = μ₂ ≡ μ₃
compatible (τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₁') ● μ₃ = (τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₁') ≡ μ₃
compatible (τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₁') (τ₂ ⇨⟨ μ₂ , σ₂ ⟩ τ₂') ● = ⊥
compatible (τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₁') (τ₂ ⇨⟨ μ₂ , σ₂ ⟩ τ₂')
           (τ₃ ⇨⟨ μ₃ , σ₃ ⟩ τ₃') =
             (τ₁ ≡ τ₃) × (compatible (τ₂ ⇨⟨ μ₂ , σ₂ ⟩ τ₂') μ₃ μ₁) ×
             (σ₁ ≡ σ₃) × (τ₁' ≡ τ₃')

-- initial continuation
id-cont-type : Ty → Tr → Mc → Ty → Set
id-cont-type τ ● [] τ' = τ ≡ τ'
id-cont-type τ ● (τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₂ × μ₂ × η :: σ₂) τ' =
               (τ ≡ τ₁) × (μ₁ ≡ μ₂) × (σ₁ ≡ σ₂) × (τ' ≡ τ₂)
id-cont-type τ (τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₁') σ τ' =
               (τ ≡ τ₁) × (μ₁ ≡ ●) × (σ ≡ σ₁) × (τ' ≡ τ₁')

-- the following variables automatically become implicit arguments
variable
  τ τ' τ₀ τ₀' τ₁ τ₁' τ₂ τ₂' τ₃ τ₄ τ₅ τv τvc α β γ γ' δ : Ty
  μ μ' μ₀ μ₀' μ₁ μ₁' μ₂ μ₂' μ₃ μ₄ μα μβ μγ μid : Tr
  σ σ' σ₀ σ₀' σ₁ σ₁' σ₂ σ₂' σ₃ σ₄ σα σβ σγ σid : Mc
  η : H

ηd : ℕ → (τ' τ₀ τ₀' τ₁ τ₁' τ₂ τ₂' δ : Ty) → (μ₀ μ₀' μ₁ μ₂ μ₂' : Tr) →
      (σ₀ σ₀' σ₂ σ₂' : Mc) → H
ηd zero τ' τ₀ τ₀' τ₁ τ₁' τ₂ τ₂' δ μ₀ μ₀' μ₁ μ₂ μ₂' σ₀ σ₀' σ₂ σ₂' = ●
ηd (suc n) τ' τ₀ τ₀' τ₁ τ₁' τ₂ τ₂' δ μ₀ μ₀' μ₁ μ₂ μ₂' σ₀ σ₀' σ₂ σ₂' =
  τ' ⇒[ τ₁ ⇒⟨ μ₁ , τ₂ ⇨⟨ μ₂ , σ₂ ⟩ τ₂' × μ₂' ×
               ηd n τ' τ₀ τ₀' τ₁ τ₁' τ₂ τ₂' δ μ₀ μ₀' μ₁ μ₂ μ₂' σ₀ σ₀'
                  σ₂ σ₂' :: σ₂' ⟩
         τ₁' ]⇒ μ₁ ⇒ τ₀ ⟨ μ₀ , σ₀ ⟩ τ₀' ⟨ μ₀' , σ₀' ⟩ δ


-- Terms
-- e : Exp var τ μα σα α μβ σβ β  means that e
--  * has type τ
--  * requires
--      - a context that yields a computation of type α
--        when given a trail of type μα and a metacontext of type σα
--      - a trail of type μβ
--      - a metacontext of type σβ
--  * eventually returns a value of type β
data Exp (var : Ty → Set) : Ty → Tr → Mc → Ty → Tr → Mc → Ty → Set where
  Var      : var τ → Exp var τ μα σα α μα σα α
  Num      : ℕ → Exp var Nat μα σα α μα σα α
  Bol      : 𝔹 → Exp var Bool μα σα α μα σα α
  Lam      :   (var τ₂ → Exp var τ₁ μβ σβ β μγ σγ γ)
             → Exp var (τ₂ ⇒ τ₁ ⟨ μβ , σβ ⟩ β ⟨ μγ , σγ ⟩ γ)
                   μα σα α μα σα α
  App      :   Exp var (τ₂ ⇒ τ₁ ⟨ μα , σα ⟩ α ⟨ μ₂ , σ₂ ⟩ β)
                   μ₁ σ₁ γ μβ σβ δ
             → Exp var τ₂ μ₂ σ₂ β μ₁ σ₁ γ
             → Exp var τ₁ μα σα α μβ σβ δ
  Plus     :   Exp var Nat μγ σγ γ μβ σβ β
             → Exp var Nat μα σα α μγ σγ γ
             → Exp var Nat μα σα α μβ σβ β
  TryWithD :   (n : ℕ)
             → id-cont-type γ μid σid γ'
             → Exp var γ μid σid γ' ●
                   (τ ⇨⟨ μα , σα ⟩ α × μβ ×
                    (τv ⇒[ τ₁ ⇒⟨ μ₁ , τ₂ ⇨⟨ μ₂ , σ₂ ⟩ τ₂' × μ₂' ×
                                   ηd n τv τ₀ τ₀' τ₁ τ₁' τ₂ τ₂' δ μ₀ μ₀'
                                      μ₁ μ₂ μ₂' σ₀ σ₀' σ₂ σ₂' :: σ₂' ⟩
                          τ₁' ]⇒ μ₁ ⇒ τ₀ ⟨ μ₀ , σ₀ ⟩ τ₀' ⟨ μ₀' , σ₀' ⟩ δ)
                    :: σβ)
                   β
             → (var τv → var (τ₁ ⇒ τ₂ ⟨ μ₂ , σ₂ ⟩ τ₂' ⟨ μ₂' , σ₂' ⟩ τ₁') →
                Exp var τ₀ μ₀ σ₀ τ₀' μ₀' σ₀' δ)
             → Exp var τ μα σα α μβ σβ β
  TryWithS :   id-cont-type γ μid σid γ'
             → compatible (τ₂ ⇨⟨ μ₂ , σ₂ ⟩ τ₂') μ₂' μ₃
             → compatible μ₁' μ₃ μ₁
             → Exp var γ μid σid γ' ●
                   (τ ⇨⟨ μα , σα ⟩ α × μβ ×
                    (τv ⇒[ τ₁ ⇒⟨ μ₁ , σ₁ ⟩ τ₁' ]⇒ μ₁' ⇒
                     τ₀ ⟨ μ₀ , σ₀ ⟩ τ₀' ⟨ μ₀' , σ₀' ⟩ δ)
                    :: σβ)
                   β
             → (var τv → var (τ₁ ⇒ τ₂ ⟨ μ₂ , σ₂ ⟩ τ₂' ⟨ μ₂' , σ₁ ⟩ τ₁') →
                Exp var τ₀ μ₀ σ₀ τ₀' μ₀' σ₀' δ)
             → Exp var τ μα σα α μβ σβ β
  Perform  :   Exp var τv μγ
                   (τ₀ ⇨⟨ μ₀ , σ₀ ⟩ τ₀' × μ₀' ×
                    (τv ⇒[ τ ⇒⟨ μα , σα ⟩ α ]⇒ μγ ⇒
                     τ₀ ⟨ μ₀ , σ₀ ⟩ τ₀' ⟨ μ₀' , σ₀' ⟩ γ)
                    :: σ₀')
                   γ μβ σβ β
             → Exp var τ μα σα α μβ σβ β

-- Interpretation of types
〚_〛τ : Ty → Set
〚_〛μ : Tr → Set
〚_〛σ : Mc → Set
〚_〛η : H  → Set

--- Ty
〚 Nat 〛τ = ℕ
〚 Bool 〛τ = 𝔹
〚 Str 〛τ = String
〚 τ₁ ⇒ τ₂ ⟨ μα , σα ⟩ α ⟨ μβ , σβ ⟩ β 〛τ =
  〚 τ₁ 〛τ → (〚 τ₂ 〛τ → 〚 μα 〛μ → 〚 σα 〛σ → 〚 α 〛τ) →
  〚 μβ 〛μ → 〚 σβ 〛σ → 〚 β 〛τ

--- Tr
〚 ● 〛μ = ⊤
〚 τ₁ ⇨⟨ μ , σ ⟩ τ₂ 〛μ = 〚 τ₁ 〛τ → 〚 μ 〛μ → 〚 σ 〛σ → 〚 τ₂ 〛τ

--- Mc
〚 [] 〛σ = ⊤
〚 τ ⇨⟨ μ , σ ⟩ τ' × μ' × η :: σ' 〛σ =
  ((〚 τ 〛τ → 〚 μ 〛μ → 〚 σ 〛σ → 〚 τ' 〛τ) × 〚 μ' 〛μ × 〚 η 〛η)
  × 〚 σ' 〛σ

--- H
〚 ● 〛η = ⊤
〚 τv ⇒[ τ₁ ⇒⟨ μ₁ , σ₁ ⟩ τ₁' ]⇒ μ₄
   ⇒ τ ⟨ μα , σα ⟩ α ⟨ μβ , σβ ⟩ β 〛η =
  〚 τv 〛τ → (〚 τ₁ 〛τ → 〚 μ₁ 〛μ → 〚 σ₁ 〛σ → 〚 τ₁' 〛τ) →
  〚 μ₄ 〛μ → (〚 τ 〛τ → 〚 μα 〛μ → 〚 σα 〛σ → 〚 α 〛τ) →
  〚 μβ 〛μ → 〚 σβ 〛σ → 〚 β 〛τ

-- compose trail
compose-trail : compatible μ₁ μ₂ μ₃ → 〚 μ₁ 〛μ → 〚 μ₂ 〛μ → 〚 μ₃ 〛μ
compose-trail {●} refl tt t₂ = t₂
compose-trail {τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₁'} {●} refl t₁ tt = t₁
compose-trail {τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₁'} {τ₂ ⇨⟨ μ₂ , σ₂ ⟩ τ₂'}
              {τ₃ ⇨⟨ μ₃ , σ₃ ⟩ τ₃'} (refl , c , refl , refl) t₁ t₂ =
                λ v t m → t₁ v (compose-trail c t₂ t) m

-- initial continaution
idc : id-cont-type τ μ σ τ' → 〚 τ 〛τ → 〚 μ 〛μ → 〚 σ 〛σ → 〚 τ' 〛τ
idc {μ = ●} {[]} refl v tt tt = v
idc {μ = ●} {τ₁ ⇨⟨ μ₁ , σ₁ ⟩ τ₂ × .μ₁ × η :: .σ₁}
    (refl , refl , refl , refl) v tt ((c , t , h) , m) = c v t m
idc {μ = τ ⇨⟨ .● , σ ⟩ τ'} (refl , refl , refl , refl) v c m = c v tt m

-- CPS interpreter
g : {var : Ty → Set} → Exp 〚_〛τ τ μα σα α μβ σβ β →
    (〚 τ 〛τ → 〚 μα 〛μ → 〚 σα 〛σ → 〚 α 〛τ) →
    〚 μβ 〛μ → 〚 σβ 〛σ → 〚 β 〛τ
g (Var x) c t m = c x t m
g (Num n) c t m = c n t m
g (Bol b) c t m = c b t m
g (Lam f) c t m = c (λ v c₁ t₁ m₁ → g {var = 〚_〛τ} (f v) c₁ t₁ m₁) t m
g (App e₁ e₂) c t m =
  g {var = 〚_〛τ} e₁
    (λ v₁ t₁ m₁ → g {var = 〚_〛τ} e₂ (λ v₂ t₂ m₂ → v₁ v₂ c t₂ m₂) t₁ m₁) t m
g (Plus e₁ e₂) c t m =
  g {var = 〚_〛τ} e₁
    (λ n₁ t₁ m₁ →
      g {var = 〚_〛τ} e₂ (λ n₂ t₂ m₂ → c (n₁ + n₂) t₂ m₂) t₁ m₁) t m
g (TryWithD n id-c-t e₁ e₂) c t m =
  g {var = 〚_〛τ} e₁ (idc id-c-t) tt
    ((c , t , h n)
    , m)
  where h : (n : ℕ) → 〚 _ ⇒[ _ ⇒⟨ _ ,
                 _ ⇨⟨ _ , _ ⟩ _ × _ ×
                 ηd n _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ :: _ ⟩ _
                 ]⇒ _ ⇒ _ ⟨ _ , _ ⟩ _ ⟨ _ , _ ⟩ _ 〛η
        h zero v c₁ t₁ c₀ t₀ m₀ =
          g {var = 〚_〛τ}
            (e₂ v (λ v₂ c₂ t₂ m₂ → c₁ v₂ t₁ ((c₂ , t₂ , tt) , m₂)))
            c₀ t₀ m₀
        h (suc n) v c₁ t₁ c₀ t₀ m₀ =
          g {var = 〚_〛τ}
            (e₂ v (λ v₂ c₂ t₂ m₂ → c₁ v₂ t₁ ((c₂ , t₂ , h n) , m₂)))
            c₀ t₀ m₀
g (TryWithS id-c-t ct₁ ct₂ e₁ e₂) c t m =
  g {var = 〚_〛τ} e₁ (idc id-c-t) tt
    ((c , t ,
      λ { v c₁ t₁ c₀ t₀ m₀ →
        g {var = 〚_〛τ}
          (e₂ v (λ v₂ c₂ t₂ m₂ →
                   c₁ v₂ (compose-trail ct₂ t₁ (compose-trail ct₁ c₂ t₂)) m₂))
          c₀ t₀ m₀})
    , m)
g (Perform e) c t m =
  g {var = 〚_〛τ} e
    (λ { v t' ((c₀ , t₀ , h₀) , m₀) →
       h₀ v c t' c₀ t₀ m₀})
    t m

-- Top-level evaluation
f : Exp 〚_〛τ τ ● [] τ ● [] τ → 〚 τ 〛τ
f e = g {var = 〚_〛τ} e (λ v _ _ → v) tt tt

-- try perform 12 + perform 13 with x, k -> k (x + 1)
term1 : {τ : Ty} → Exp 〚_〛τ Nat ● [] τ ● [] τ
term1 = TryWithD (suc zero) -- cannot be zero
                 (refl , refl , refl , refl)
                 (Plus (Perform (Num )) (Perform (Num )))
                 λ x k → App (Var k) (Plus (Var x) (Num ))