------------------------------------------------------------------------
-- Based on "From Interpreter to Compiler: A Representational
-- Derivation" by Mitchell Wand
------------------------------------------------------------------------

{-# OPTIONS --no-termination-check #-}

open import Data.Star.Nat

module Wand2
         -- Semantic domains of base types.
         (Base : ℕ -> Set)
         -- The result type for the continuations.
         (Result : Set)
         where

open import Data.Star
open import Data.Star.List
open import Data.Star.Decoration
import Data.Star.Environment as Env
open import Data.Unit
open import Relation.Binary
open import Function hiding (_∋_)
open import Relation.Binary.PropositionalEquality

------------------------------------------------------------------------
-- Type system

-- The language used is the simply typed lambda calculus with
-- multi-argument functions.

infixl 70 _·_
infix  50 ƛ_
infixr 30 _⟶_
infix   4 _⊢_

data Ty : Set where
  base : ℕ -> Ty
  _⟶_  : List Ty -> Ty -> Ty

open Env Ty

mutual

  data _⊢_ : Ctxt -> Ty -> Set where
    var  : forall {Γ τ}    -> Γ ∋ τ -> Γ ⊢ τ
    ƛ_   : forall {Γ Γ⁺ τ} -> Γ ▻▻ Γ⁺ ⊢ τ -> Γ ⊢ Γ⁺ ⟶ τ
    _·_  : forall {Γ Γ⁺ τ} -> Γ ⊢ Γ⁺ ⟶ τ -> Args Γ Γ⁺ -> Γ ⊢ τ

  -- An argument sequence is defined as a context extension decorated
  -- with terms.

  Args : Ctxt -> Ctxt -> Set
  Args Γ Γ⁺ = All (\σ -> Γ ⊢ σ) Γ⁺

------------------------------------------------------------------------
-- An example, just to see how things work

module Example (x : ε ⊢ base 0#) (y : ε ⊢ base 1#) where

  Types : Ctxt
  Types = ε ▻ base 0# ▻ base 1#

  fun : ε ⊢ Types ⟶ base 0#
  fun = ƛ var (vs vz)

  args : Args ε Types
  args = ε {I = NonEmpty _} ▻ ↦ x ▻ ↦ y

  application : ε ⊢ base 0#
  application = fun · args

------------------------------------------------------------------------
-- The semantics which Wand starts off with

-- Continuations.

Cont : Set -> Set
Cont A = (A -> Result) -> Result

-- Semantic domains. Note that the code in Wand's paper is untyped.
-- Before I understood what the semantic domains should be I had
-- trouble understanding his code. Wand presumably had something
-- resembling these types in his head; writing them out would have
-- made his paper more accessible (at least to people like me who
-- usually write interpreters in a different way).

⟦_⟧⋆ : Ty -> Set
⟦ base n ⟧⋆ = Base n
⟦ Γ⁺ ⟶ τ ⟧⋆ = Env ⟦_⟧⋆ Γ⁺ -> Cont ⟦ τ ⟧⋆

-- The function above is terminating, since ⟦_⟧⋆ is only applied to
-- the elements of Γ⁺. The termination checker cannot see this,
-- though.

-- The interpreter.

-- The interpreter's result type.

Dom : Ctxt -> Ty -> Set
Dom Γ σ = Env ⟦_⟧⋆ Γ -> Cont ⟦ σ ⟧⋆

-- These two functions are also terminating. Due to the use of reverse
-- this is not apparent to the termination checker, though.

mutual

  ⟦_⟧ : forall {Γ σ} -> Γ ⊢ σ -> Dom Γ σ
  ⟦ var x  ⟧ ρ k = k (lookup x ρ)
  ⟦ ƛ t    ⟧ ρ k = k (\vs k' -> ⟦ t ⟧ (ρ ▻▻▻ vs) k')
  ⟦ t · ts ⟧ ρ k = ⟦ t • ε • reverse id ts ⟧ ρ k

  -- Note that this function exhibits a clowns/jokers structure, by
  -- the way (a zipper where the things on the two sides have
  -- different types).

  ⟦_•_•_⟧ : forall {Γ Γ⁺ Δ⁺ τ} ->
            Γ ⊢ Γ⁺ ⟶ τ -> Env ⟦_⟧⋆ Δ⁺ ->
            Star (flip (DecoratedWith (\σ -> Γ ⊢ σ)))
                 (nonEmpty Δ⁺) (nonEmpty Γ⁺) ->
            Dom Γ τ
  ⟦ t • ρ' • ε        ⟧ ρ k = ⟦ t ⟧ ρ (\v -> v ρ' k)
  ⟦ t • ρ' • ↦ a ◅ as ⟧ ρ k = ⟦ a ⟧ ρ (\v -> ⟦ t • ρ' ▻ ↦ v • as ⟧ ρ k)

------------------------------------------------------------------------
-- Back to the example

module Example₂ (x : ε ⊢ base 0#) (y : ε ⊢ base 1#) where

  open Example x y

  sem : Cont (Base 0#)
  sem = ⟦ application ⟧ ε

------------------------------------------------------------------------
-- First step of the derivation: name subexpressions

module Naming where

  *fetch : forall {Γ τ} -> Γ ∋ τ -> Dom Γ τ
  *fetch x = \ρ k -> k (lookup x ρ)

  *close : forall {Γ Γ⁺ τ} -> Dom (Γ ▻▻ Γ⁺) τ -> Dom Γ (Γ⁺ ⟶ τ)
  *close v = \ρ k -> k (\vs k' -> v (ρ ▻▻▻ vs) k')

  Dom' : Ctxt -> Ctxt -> Ty -> Set
  Dom' Δ⁺ Γ τ = Env ⟦_⟧⋆ Δ⁺ -> Dom Γ τ

  *clear-ev : forall {Γ τ} -> Dom' ε Γ τ -> Dom Γ τ
  *clear-ev f = \ρ k -> f ε ρ k

  *jsr : forall {Γ Δ⁺ σ τ} ->
         Dom Γ σ -> Dom' (Δ⁺ ▻ σ) Γ τ -> Dom' Δ⁺ Γ τ
  *jsr a as ρ' = \ρ k -> a ρ (\v -> as (ρ' ▻ ↦ v) ρ k)

  *ev-fn : forall {Γ Γ⁺ τ} -> Dom Γ (Γ⁺ ⟶ τ) -> Dom' Γ⁺ Γ τ
  *ev-fn f ρ' = \ρ k -> f ρ (\v -> v ρ' k)

  -- Note that *jsr and *ev-fn have similar structure. Try to merge
  -- them.

  *ap : forall {Γ Γ⁺ τ} -> Dom' (Γ⁺ ▻ Γ⁺ ⟶ τ) Γ τ
  *ap (↦ v ◅ ρ') ρ k = v ρ' k

  -- Note that *ev-fn f = *jsr f *ap:

  prop : forall {Γ Γ⁺ τ} (f : Dom Γ (Γ⁺ ⟶ τ)) (ρ' : Env ⟦_⟧⋆ Γ⁺) ->
         *ev-fn {τ = τ} f ρ' ≡ *jsr {τ = τ} f (*ap {τ = τ}) ρ'
  prop f ρ' = refl

------------------------------------------------------------------------
-- Second step: defunctionalise

module Compiler where

  mutual

    data Code : Ctxt -> Ty -> Set where
      *fetch    : forall {Γ τ} -> Γ ∋ τ -> Code Γ τ
      *close    : forall {Γ Γ⁺ τ} ->
                  Code (Γ ▻▻ Γ⁺) τ -> Code Γ (Γ⁺ ⟶ τ)
      *clear-ev : forall {Γ τ} -> Code' ε Γ τ -> Code Γ τ

    data Code' : Ctxt -> Ctxt -> Ty -> Set where
      *jsr : forall {Γ Δ⁺ σ τ} ->
             Code Γ σ -> Code' (Δ⁺ ▻ σ) Γ τ -> Code' Δ⁺ Γ τ
      *ap  : forall {Γ Γ⁺ τ} -> Code' (Γ⁺ ▻ Γ⁺ ⟶ τ) Γ τ

  -- Compiler. It is terminating, but the use of reverse means that
  -- the termination checker cannot see this.

  mutual

    comp : forall {Γ σ} -> Γ ⊢ σ -> Code Γ σ
    comp (var x)   = *fetch x
    comp (ƛ t)     = *close (comp t)
    comp (t₁ · t₂) = *clear-ev (comp' t₁ (reverse id t₂))

    comp' : forall {Γ Γ⁺ Δ⁺ τ} ->
            Γ ⊢ Γ⁺ ⟶ τ ->
            Star (flip (DecoratedWith (\σ -> Γ ⊢ σ)))
                 (nonEmpty Δ⁺) (nonEmpty Γ⁺) ->
            Code' Δ⁺ Γ τ
    comp' t₁ ε          = *jsr (comp t₁) *ap
    comp' t₁ (↦ a ◅ as) = *jsr (comp a) (comp' t₁ as)

------------------------------------------------------------------------
-- Back to the example again

module Example₃ (x : ε ⊢ base 0#) (y : ε ⊢ base 1#) where

  open Example x y
  open Compiler

  code : Code ε (base 0#)
  code = comp application