------------------------------------------------------------------------
-- An incomplete derivation of a compiler for a simple λ-calculus
------------------------------------------------------------------------

import Level
open import Relation.Binary.PropositionalEquality as P using (_≡_; _≗_)

-- For simplicity propositional equality of functions is assumed to be
-- extensional.

module CompositionBased.Lambda
         (ext : P.Extensionality Level.zero Level.zero) where

open import Data.List hiding (map)
open import Data.Product as Prod hiding (map)
open import Data.Unit
open import Function hiding (_∋_)

------------------------------------------------------------------------
-- Language

infixr 6 _⟶_
infix  4 _∋_ _⊢_

-- Types. Currently uninhabited.

data Type : Set where
  _⟶_ : (σ τ : Type) → Type

-- Contexts.

Ctxt : Set
Ctxt = List Type

-- Variables.

data _∋_ : Ctxt → Type → Set where
  zero : ∀ {Γ σ} → σ ∷ Γ ∋ σ
  suc  : ∀ {Γ σ τ} (x : Γ ∋ σ) → τ ∷ Γ ∋ σ

-- A simply typed λ-calculus

data _⊢_ (Γ : Ctxt) : Type → Set where
  var : ∀ {τ} (x : Γ ∋ τ) → Γ ⊢ τ
  ƛ   : ∀ {σ τ} (t : σ ∷ Γ ⊢ τ) → Γ ⊢ σ ⟶ τ
  _·_ : ∀ {σ τ} (t₁ : Γ ⊢ σ ⟶ τ) (t₂ : Γ ⊢ σ) → Γ ⊢ τ

------------------------------------------------------------------------
-- Interpreting contexts

-- Parametrised interpretation of contexts.

Ctxt-interpretation : (Type → Set) → (Ctxt → Set)
Ctxt-interpretation ⟦_⟧ []      = ⊤
Ctxt-interpretation ⟦_⟧ (σ ∷ Γ) =
  Σ[ x ∈ ⟦ σ ⟧ ] Ctxt-interpretation ⟦_⟧ Γ

-- Map.

map : {⟦_⟧₁ ⟦_⟧₂ : Type → Set} →
      (∀ {σ} → ⟦ σ ⟧₁ → ⟦ σ ⟧₂) →
      ∀ {Γ} → Ctxt-interpretation ⟦_⟧₁ Γ → Ctxt-interpretation ⟦_⟧₂ Γ
map f {[]}    = id
map f {σ ∷ Γ} = Prod.map f (map f)

-- Lookup.

lookup : ∀ {⟦_⟧ Γ σ} → Γ ∋ σ → Ctxt-interpretation ⟦_⟧ Γ → ⟦ σ ⟧
lookup zero    = proj₁
lookup (suc x) = lookup x ∘ proj₂

-- The lookup function is natural.

lookup-natural : ∀ {⟦_⟧₁ ⟦_⟧₂ : Type → Set} {Γ σ}
                 (x : Γ ∋ σ) (f : ∀ {σ} → ⟦ σ ⟧₁ → ⟦ σ ⟧₂)
                 (ρ : Ctxt-interpretation ⟦_⟧₁ Γ) →
                 lookup x (map f ρ) ≡ f (lookup x ρ)
lookup-natural zero    f = λ _ → P.refl
lookup-natural (suc x) f = lookup-natural x f ∘ proj₂

------------------------------------------------------------------------
-- Semantics

module Semantics where

  -- Interpretation of types.

  Value : Type → Set
  Value (σ ⟶ τ) = Value σ → Value τ

  -- Environments.

  Env = Ctxt-interpretation Value

  -- Semantic domains.

  Dom : Ctxt → Type → Set
  Dom Γ σ = Env Γ → Value σ

  -- The semantics of the language.

  ⟦_⟧ : ∀ {Γ σ} → Γ ⊢ σ → Dom Γ σ
  ⟦ var x   ⟧ ρ = lookup x ρ
  ⟦ ƛ t     ⟧ ρ = λ v → ⟦ t ⟧ (v , ρ)
  ⟦ t₁ · t₂ ⟧ ρ = (⟦ t₁ ⟧ ρ) (⟦ t₂ ⟧ ρ)

------------------------------------------------------------------------
-- Functions taking a well-typed stack and a well-typed environment to
-- a value

module StackEnv where

  open Semantics public using (Value; Env)

  -- Stacks.

  Stack : Ctxt → Set
  Stack = Env

  -- Semantic domains: stack → environment → result.

  Dom : Ctxt → Ctxt → Type → Set
  Dom Δ Γ σ = Stack Δ → Env Γ → Value σ

  -- Equality.

  infix 4 _≈_

  _≈_ : ∀ {Δ Γ σ} → Dom Δ Γ σ → Dom Δ Γ σ → Set
  f ≈ g = ∀ s ρ → f s ρ ≡ g s ρ

  -- Equational reasoning.

  infix  3 _∎
  infixr 2 _≈⟨_⟩_ _≈⟨⟩_

  _≈⟨_⟩_ : ∀ {Δ Γ σ} (f : Dom Δ Γ σ) {g h : Dom Δ Γ σ} →
           f ≈ g → g ≈ h → f ≈ h
  _ ≈⟨ f≈g ⟩ g≈h = λ s ρ → P.trans (f≈g s ρ) (g≈h s ρ)

  _≈⟨⟩_ : ∀ {Δ Γ σ} (f : Dom Δ Γ σ) {g : Dom Δ Γ σ} →
          f ≈ g → f ≈ g
  _ ≈⟨⟩ g≈h = g≈h

  _∎ : ∀ {Δ Γ σ} (f : Dom Δ Γ σ) → f ≈ f
  f ∎ = λ s ρ → P.refl

  sym : ∀ {Δ Γ σ} {f g : Dom Δ Γ σ} → f ≈ g → g ≈ f
  sym f≈g = λ s ρ → P.sym (f≈g s ρ)

  -- Helper functions for derivations.

  EqualTo : ∀ {Δ Γ σ} → Dom Δ Γ σ → Set
  EqualTo f = ∃ (λ g → f ≈ g)

  infix 0 ▷_

  ▷_ : ∀ {Δ Γ σ} {f g : Dom Δ Γ σ} → f ≈ g → EqualTo f
  ▷ f≈g = (_ , f≈g)

------------------------------------------------------------------------
-- Rewrite the evaluator in continuation-passing style

module CPS where

  private
    open module S = Semantics using () renaming (⟦_⟧ to ⟦_⟧S)
  open StackEnv

  -- A CPS-transformed variant of Dom.

  DomCPS : Ctxt → Type → Set
  DomCPS Γ σ = ∀ {Δ τ} → Dom (σ ∷ Δ) Γ τ → Dom Δ Γ τ

  -- A composition operator used to specify ⟦_⟧. The first argument
  -- gets access only to the environment, while the second one also
  -- gets to see the stack. TODO: Is this operator related to Kripke
  -- models? (Observe how the contexts change in the lambda and
  -- application cases below.)

  infixr 9 _⊙_ _⊙-cong_

  _⊙_ : ∀ {Γ σ} → S.Dom Γ σ → DomCPS Γ σ
  (f ⊙ κ) s ρ = κ (f ρ , s) ρ

  _⊙-cong_ : ∀ {Γ σ} {f₁ f₂ : S.Dom Γ σ} →
             (∀ ρ → f₁ ρ ≡ f₂ ρ) →
             ∀ {Δ τ} {κ₁ κ₂ : Dom (σ ∷ Δ) Γ τ} →
             κ₁ ≈ κ₂ →
             f₁ ⊙ κ₁ ≈ f₂ ⊙ κ₂
  _⊙-cong_ {f₁ = f₁} {f₂} f₁≈f₂ {κ₁ = κ₁} {κ₂} κ₁≈κ₂ =
    f₁ ⊙ κ₁  ≈⟨ (λ s ρ → P.cong (λ x → κ₁ (x , s) ρ) (f₁≈f₂ ρ)) ⟩
    f₂ ⊙ κ₁  ≈⟨ (λ s ρ → κ₁≈κ₂ (f₂ ρ , s) ρ) ⟩
    f₂ ⊙ κ₂  ∎

  -- The semantics.

  -- One could calculate/devise some other definition of the var′
  -- function, but allowing the unary representation of variables to
  -- affect the design of the machine/compiler seems like a bad idea.

  var′ : ∀ {Γ σ} → Γ ∋ σ → DomCPS Γ σ
  var′ x κ = lookup x ⊙ κ

  clos : ∀ {Γ σ τ} → Dom [] (σ ∷ Γ) τ → DomCPS Γ (σ ⟶ τ)
  clos f κ = (λ ρ v → f _ (v , ρ)) ⊙ κ

  clos-cong : ∀ {Γ σ τ} {f₁ f₂ : Dom [] (σ ∷ Γ) τ} → f₁ ≈ f₂ →
              ∀ {Δ χ} (κ : Dom (σ ⟶ τ ∷ Δ) Γ χ) →
              clos f₁ κ ≈ clos f₂ κ
  clos-cong f₁≈f₂ κ =
    (λ ρ → ext λ v → f₁≈f₂ _ (v , ρ)) ⊙-cong (κ ∎)

  top : ∀ {Γ σ} → Dom [ σ ] Γ σ
  top (v , s) ρ = v

  app : ∀ {Γ Δ σ τ χ} → Dom (τ ∷ Δ) Γ χ → Dom (σ ∷ σ ⟶ τ ∷ Δ) Γ χ
  app κ (x , f , s) = κ (f x , s)

  mutual

    ⟦_⟧ : ∀ {Γ σ} → Γ ⊢ σ → DomCPS Γ σ
    ⟦ t ⟧ κ = proj₁ (⟦ t ⟧′ κ)

    ⟦_⟧′ : ∀ {Γ σ} (t : Γ ⊢ σ) {Δ τ} (κ : Dom (σ ∷ Δ) Γ τ) →
           EqualTo (⟦ t ⟧S ⊙ κ)
    ⟦ var x ⟧′ κ = ▷
      ⟦ var x ⟧S ⊙ κ  ≈⟨⟩
      lookup x ⊙ κ    ≈⟨⟩
      var′ x κ        ∎

    ⟦ ƛ t ⟧′ κ = ▷
      ⟦ ƛ t ⟧S ⊙ κ                  ≈⟨⟩
      (λ ρ v → ⟦ t ⟧S (v , ρ)) ⊙ κ  ≈⟨⟩
      clos (⟦ t ⟧S ⊙ top) κ         ≈⟨ clos-cong (proj₂ $ ⟦ t ⟧′ top) κ ⟩
      clos (⟦ t ⟧ top) κ            ∎

    ⟦ t₁ · t₂ ⟧′ κ = ▷
      ⟦ t₁ · t₂ ⟧S ⊙ κ                   ≈⟨⟩
      (λ ρ → ⟦ t₁ ⟧S ρ (⟦ t₂ ⟧S ρ)) ⊙ κ  ≈⟨⟩
      ⟦ t₁ ⟧S ⊙ ⟦ t₂ ⟧S ⊙ app κ          ≈⟨ (λ _ → P.refl) ⊙-cong proj₂ (⟦ t₂ ⟧′ (app κ)) ⟩
      ⟦ t₁ ⟧S ⊙ ⟦ t₂ ⟧ (app κ)           ≈⟨ proj₂ $ ⟦ t₁ ⟧′ (⟦ t₂ ⟧ (app κ)) ⟩
      ⟦ t₁ ⟧ (⟦ t₂ ⟧ (app κ))            ∎

  correct : ∀ {Γ σ} (t : Γ ⊢ σ) ρ → ⟦ t ⟧S ρ ≡ ⟦ t ⟧ top _ ρ
  correct {Γ} t ρ = begin
    ⟦ t ⟧S ρ            ≡⟨ P.refl ⟩
    (⟦ t ⟧S ⊙ top) _ ρ  ≡⟨ proj₂ (⟦ t ⟧′ top) _ ρ ⟩
    ⟦ t ⟧ top _ ρ       □
    where open P.≡-Reasoning renaming (_∎ to _□)

  -- ⟦_⟧ preserves equality in its continuation argument.

  ⟦_⟧-cong : ∀ {Γ σ} (t : Γ ⊢ σ) {Δ τ} {κ₁ κ₂ : Dom (σ ∷ Δ) Γ τ} →
             κ₁ ≈ κ₂ → ⟦ t ⟧ κ₁ ≈ ⟦ t ⟧ κ₂
  ⟦_⟧-cong t {κ₁ = κ₁} {κ₂} κ₁≈κ₂ =
    ⟦ t ⟧ κ₁     ≈⟨ sym $ proj₂ $ ⟦ t ⟧′ κ₁ ⟩
    ⟦ t ⟧S ⊙ κ₁  ≈⟨ (λ _ → P.refl) ⊙-cong κ₁≈κ₂ ⟩
    ⟦ t ⟧S ⊙ κ₂  ≈⟨ proj₂ $ ⟦ t ⟧′ κ₂ ⟩
    ⟦ t ⟧ κ₂     ∎

------------------------------------------------------------------------
-- Defunctionalise

module Defunctionalised where

  open Semantics using () renaming (⟦_⟧ to ⟦_⟧S)
  open CPS using () renaming (⟦_⟧ to ⟦_⟧CPS; ⟦_⟧-cong to ⟦_⟧CPS-cong)
  open StackEnv

  -- I defunctionalise with respect to Dom.

  data Term : Ctxt → Ctxt → Type → Set where
    var     : ∀ {Γ Δ σ τ}
              (x : Γ ∋ σ) (t : Term (σ ∷ Δ) Γ τ) → Term Δ Γ τ
    closure : ∀ {Γ Δ σ τ χ}
              (t₁ : Term [] (σ ∷ Γ) τ) (t₂ : Term (σ ⟶ τ ∷ Δ) Γ χ) →
              Term Δ Γ χ
    top     : ∀ {Γ σ} → Term [ σ ] Γ σ
    app     : ∀ {Γ Δ σ τ χ}
              (t : Term (τ ∷ Δ) Γ χ) → Term (σ ∷ σ ⟶ τ ∷ Δ) Γ χ

  -- The semantics of Term.
  --
  -- Note that exec is /not/ a virtual machine: in the closure case
  -- there is a non-tail call.

  exec : ∀ {Γ Δ σ} → Term Δ Γ σ → Dom Δ Γ σ
  exec (var x t)       s           ρ = exec t (lookup x ρ , s) ρ
  exec (closure t₁ t₂) s           ρ = exec t₂ ((λ v → exec t₁ _ (v , ρ)) , s) ρ
  exec top             (v , s)     ρ = v
  exec (app t)         (x , f , s) ρ = exec t (f x , s) ρ

  -- Compiler.

  mutual

    comp : ∀ {Γ Δ σ τ} → Γ ⊢ σ → Term (σ ∷ Δ) Γ τ → Term Δ Γ τ
    comp t κ = proj₁ (comp′ t κ)

    comp′ : ∀ {Γ Δ σ τ} (t : Γ ⊢ σ) (κ : Term (σ ∷ Δ) Γ τ) →
            ∃ λ (t′ : Term Δ Γ τ) → ⟦ t ⟧CPS (exec κ) ≈ exec t′
    comp′ (var x) κ = var x κ , (
      ⟦ var x ⟧CPS (exec κ)  ≈⟨⟩
      CPS.var′ x (exec κ)    ≈⟨⟩
      exec (var x κ)         ∎)

    comp′ (ƛ t) κ = closure (comp t top) κ , (
      ⟦ ƛ t ⟧CPS (exec κ)                   ≈⟨⟩
      CPS.clos (⟦ t ⟧CPS CPS.top) (exec κ)  ≈⟨ CPS.clos-cong (proj₂ $ comp′ t top) (exec κ) ⟩
      exec (closure (comp t top) κ)         ∎)

    comp′ (t₁ · t₂) κ = comp t₁ (comp t₂ (app κ)) , (
      ⟦ t₁ · t₂ ⟧CPS (exec κ)                   ≈⟨⟩
      ⟦ t₁ ⟧CPS (⟦ t₂ ⟧CPS (CPS.app (exec κ)))  ≈⟨⟩
      ⟦ t₁ ⟧CPS (⟦ t₂ ⟧CPS (exec (app κ)))      ≈⟨ ⟦ t₁ ⟧CPS-cong (proj₂ $ comp′ t₂ (app κ)) ⟩
      ⟦ t₁ ⟧CPS (exec (comp t₂ (app κ)))        ≈⟨ proj₂ (comp′ t₁ (comp t₂ (app κ))) ⟩
      exec (comp t₁ (comp t₂ (app κ)))          ∎)

  correct : ∀ {Γ σ} (t : Γ ⊢ σ) ρ →
            ⟦ t ⟧S ρ ≡ exec (comp t top) _ ρ
  correct t ρ = begin
    ⟦ t ⟧S ρ               ≡⟨ CPS.correct t ρ ⟩
    ⟦ t ⟧CPS CPS.top _ ρ   ≡⟨ proj₂ (comp′ t top) _ ρ ⟩
    exec (comp t top) _ ρ  □
    where open P.≡-Reasoning renaming (_∎ to _□)