Staged.agda

module Staged where

open import Data.Nat using (ℕ; zero; suc)
open import Data.Empty using (⊥)
open import Relation.Binary.PropositionalEquality
open ≡-Reasoning

-- indices, we could open import Data.Fin renaming too.
-- probably should (so we get all the lemmas).
-- but for now it's defined inline.

data Idx : ℕ → Set where
  iz : ∀ {n} → Idx (suc n)
  is : ∀ {n} → Idx n → Idx (suc n)


-- Syntax
---------------------------------------------------------------

-- simple types

data Ty : Set where
  fun  : Ty → Ty → Ty
  code : Ty → Ty
  int  : Ty

pattern _⇒_ α β = fun α β
infixr 5 _⇒_

data Term (n : ℕ) : Set where
  var    : Idx n → Term n                                -- x
  app    : Term n → Term n → Term n                      -- f t
  lam    : Ty → Term (suc n) → Term n                    -- λ (x : α) → t
  zero   : Term n                                        -- Z
  succ   : Term n → Term n                               -- S n
  elim   : Term n → Ty → Term n → Term (suc n) → Term n  -- elim α z (λ m → s) n
  quot   : Term n → Term n                               -- ⟦ t ⟧
  splice : Term n → Term n                               -- $t

pattern var₀ = var iz

-- Environments, essentially Data.Vec,
-- but we want own syntax and snoc behavior (do we?)
-- well that will make it look more like in the literature.
data Env (A : Set) : ℕ → Set where
  ∙   : Env A zero
  _,_ : ∀ {n} → Env A n → A → Env A (suc n)

infixl 5 _,_

mapEnv : ∀ {A B} → (A → B) → ∀ {n} → Env A n → Env B n
mapEnv f ∙       = ∙
mapEnv f (Γ , x) = mapEnv f Γ , f x

mapEnv-comp : ∀ {A B C} → (f : B → C) (g : A → B) → ∀ {n} → (Γ : Env A n)
  → mapEnv f (mapEnv g Γ) ≡ mapEnv (λ x → f (g x)) Γ
mapEnv-comp f g ∙       = refl
mapEnv-comp f g (Γ , x) = cong (_, f (g x)) (mapEnv-comp f g Γ)

-- this is a bit tricky to avoid needing funExt
mapEnv-id : ∀ {A} → (f : A → A) → (∀ x → f x ≡ x) → ∀ {n} → (Γ : Env A n)
  → mapEnv f Γ ≡ Γ
mapEnv-id f f-id ∙       = refl
mapEnv-id f f-id (Γ , x) = cong₂ _,_ (mapEnv-id f f-id Γ) (f-id x)

-- Typing relation
---------------------------------------------------------------

infix 0 _∋_⦂_ _⊢_⦂_

-- variable with a type in context.
-- we /could/ (define and) use lookup,
-- but IIRC it's handier to have inductive type
-- as we can always match on it.
--
-- alternatively varₜ would need to be defined like
-- varₜ : ∀ {x α} → lookup Γ x ≡ α → Γ ⊢ var x ⦂ α
--
-- either one works.
data _∋_⦂_ : ∀ {n} → Env Ty n → Idx n → Ty → Set where
  izₜ : ∀ {α n}     {Γ : Env Ty n} → Γ , α ∋ iz ⦂ α
  isₜ : ∀ {α n β x} {Γ : Env Ty n} → Γ ∋ x ⦂ α → Γ , β ∋ is x ⦂ α 

-- Typing relation.
-- The best (ab)use of Agda mixfix support.
data _⊢_⦂_ {n} (Γ : Env Ty n) : Term n → Ty → Set where
  varₜ : ∀ {x α}
    → Γ ∋ x ⦂ α
    ---------------
    → Γ ⊢ var x ⦂ α

  appₜ : ∀ {f t α β}
    → Γ ⊢ f ⦂ fun α β
    → Γ ⊢ t ⦂ α
    -----------------
    → Γ ⊢ app f t ⦂ β

  lamₜ : ∀ {t α β}
    → Γ , α ⊢ t ⦂ β
    -----------------------
    → Γ ⊢ lam α t ⦂ fun α β

  quotₜ : ∀ {t α}
    → Γ ⊢ t ⦂ α
    ---------------------
    → Γ ⊢ quot t ⦂ code α

  spliceₜ : ∀ {t α}
    → Γ ⊢ t ⦂ code α
    ---------------------
    → Γ ⊢ splice t ⦂ α

  -- TODO: add rules for zero, succ and elim


-- Examples
---------------------------------------------------------------

-- ⟦ λ x → x ⟧
timely : Term zero
timely = quot (lam int var₀)

-- λ y → $y
hasty : Term zero
hasty = lam (code int) (splice var₀)

-- λ z → ⟦ z ⟧
tardy : Term zero
tardy = lam int (quot var₀)

-- These are all /well-typed/.
--
-- And in fact, C-a solves these, as the typing-relation
-- is syntax directed.

timelyₜ : ∙ ⊢ timely ⦂ code (fun int int)
timelyₜ = quotₜ (lamₜ (varₜ izₜ))

hastyₜ : ∙ ⊢ hasty ⦂ fun (code int) int
hastyₜ = lamₜ (spliceₜ (varₜ izₜ))

tardyₜ : ∙ ⊢ tardy ⦂ fun int (code int)
tardyₜ = lamₜ (quotₜ (varₜ izₜ))

-- but only timely is well-staged!
-- what means to be well-staged?
-- we have to define that.

-- Staging relation
---------------------------------------------------------------

-- staged levels are integers.
import Data.Integer as ℤ
import Data.Integer.Properties as ℤP
open ℤ using (ℤ; 0ℤ)

-- We could use some fancy syntax here like
-- Γ ⊢ x staged, i.e. _⊢_staged.
-- but maybe better not to for now.
--
-- ₓ subscript is for "variables" (x)
data well-stagedₓ : ∀ {n} → Env ℤ n → Idx n → Set where
  izₛ : ∀ {n}     {Γ : Env ℤ n} → well-stagedₓ (Γ , 0ℤ) iz
  isₛ : ∀ {n x l} {Γ : Env ℤ n} → well-stagedₓ Γ x → well-stagedₓ (Γ , l) (is x) 

data well-staged {n} (Γ : Env ℤ n) : Term n → Set where
  varₛ : ∀ {x} → well-stagedₓ Γ x → well-staged Γ (var x)

  appₛ : ∀ {f t}
    → well-staged Γ f
    → well-staged Γ t
    -----------------
    → well-staged Γ (app f t)

  lamₛ : ∀ {t α}
    → well-staged (Γ , 0ℤ) t
    --------------------------------
    → well-staged Γ (lam α t)

  quotₛ : ∀ {t}
    → well-staged (mapEnv ℤ.suc Γ) t
    --------------------------------
    → well-staged Γ (quot t)

  spliceₜ : ∀ {t}
    → well-staged (mapEnv ℤ.pred Γ) t
    ---------------------------------
    → well-staged Γ (splice t)

-- Usually the staging relation is defined as "at staged n".
-- We define here "at level 0" and instead adjust the context
-- with mapEnv in quotₛ and spliceₜ
--
-- I think it's cleaner, and I have a reasoning for that
-- to be explained later.

-- In this development typing and staging relation are separate.
-- In practice it's probably better to combine them,
-- after all they are both syntactically directed,
-- so are similar
--
-- And it's hard to think why you would want be able
-- to consider say just well typed terms, which are possibly
-- ill-staged. Except for example purposes as in here.

-- Staging examples
---------------------------------------------------------------

-- only timely is well-staged!
timelyₛ : well-staged ∙ timely
timelyₛ = quotₛ (lamₛ (varₛ izₛ))

-- (her we see the benefit of defining well-stagedₓ inductively).
¬hastyₛ : well-staged ∙ hasty → ⊥
¬hastyₛ (lamₛ (spliceₜ (varₛ ())))

¬tardyₛ : well-staged ∙ tardy → ⊥
¬tardyₛ (lamₛ (quotₛ (varₛ ())))
 

-- There are also examples which is well staged but not well typed
--
-- λ (x : α) → x x
ohnoes : Ty → Term zero
ohnoes α = lam α (app var₀ var₀)

-- ohnoes is well staged
ohnoesₛ : (α : Ty) → well-staged ∙ (ohnoes α)
ohnoesₛ _ = lamₛ (appₛ (varₛ izₛ) (varₛ izₛ))

-- ... but not well-typed
-- though we'd need to proof uniqueness of typing to show that in general.
-- but for specific instantiation it's simple.
¬ohnoesₜ-int : ∙ ⊢ ohnoes int ⦂ int ⇒ int → ⊥
¬ohnoesₜ-int (lamₜ (appₜ (varₜ ()) _))


-- Computation
---------------------------------------------------------------

-- If we define small-step relation for this calculus
-- the quot and splice have simple rules
-- quot (splice t) ↝ t  i.e. ⟦ $t ⟧ ↝ t
-- splice (quot t) ↝ t  i.e. $⟦ t ⟧ ↝ t

-- these preserve both typing and staging, e.g. the first one

preservation-qsₜ : ∀ {n Γ α} {t : Term n}
  → Γ ⊢ quot (splice t) ⦂ code α
  ------------------------------
  → Γ ⊢ t ⦂ code α
preservation-qsₜ (quotₜ (spliceₜ tₜ)) = tₜ

preservation-qsₛ : ∀ {n Γ} {t : Term n}
  → well-staged Γ (quot (splice t))
  ---------------------------------
  → well-staged Γ t
preservation-qsₛ (quotₛ (spliceₜ tₛ)) = subst (λ Γ → well-staged Γ _) (lemma _) tₛ
  where
    lemma : ∀ {n} (Γ : Env ℤ n) →  mapEnv ℤ.pred (mapEnv ℤ.suc Γ) ≡ Γ
    lemma Γ =
      mapEnv ℤ.pred (mapEnv ℤ.suc Γ)     ≡⟨ mapEnv-comp _ _ Γ ⟩
      mapEnv (λ x → ℤ.pred (ℤ.suc x)) Γ  ≡⟨ mapEnv-id _ ℤP.pred-suc Γ ⟩ 
      Γ                                  ∎

-- Most of small-step relation are boring congruence rules like
--
-- t ↝ t′
-- ------------------
-- lam α t ↝ lam α t′
--
-- And in fact, for quot and splice these can look like that too!
--
-- However, to actually do staging (and not just add extra syntax)
-- the rules need to be changed.
--
-- We'd need to index the small-step relation by level,
-- make quot and splice congruence rules increase/decrease it
-- and allow-reduction only at appropriate levels.
--
-- That is possible, but not very fun.
--
-- Instead we could keep the small-step relation unindexed
-- and change /the syntax/.
--
-- That's idea of splice environments in matthew's thesis
-- and staging with class paper
-- https://www.andres-loeh.de/StagingWithClass/
--
-- Instead of quot and splice, we'd have splice environments
--
-- spliceEnv : Splices n m → Term (n + m) → Term n
--   where Splices n m = Env (Term n) m
--
-- then the "boring" small-step congruence rule
--
-- env ↝ env′
-- ----------------------------------
-- spliceEnv env t ↝ spliceEnv env′ t
--
-- will be enough.
-- We can reduce the splices, but the code part (t) will stay as is.
--
-- It would be fun to formalise such staging setup in Agda in this style,
-- but it will take longer than an afternoon.