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

-- Formalisation of "Compiling Exceptions Correctly" and "Calculating
-- an Exceptional Machine" by Graham Hutton and Joel Wright.

-- This module formalises "Hutton's razor", i.e. a simple language
-- with numbers and addition and nothing else.

open import Relation.Binary.PropositionalEquality

module HuttonsRazor
         -- Assume extensionality; it appears to be necessary for a
         -- proof below.
         (ext : forall {a b : Set} {f g : a -> b} ->
                  (forall x -> f x ≡ g x) -> f ≡ g)
       where

open import Data.Star
open import Data.Star.Nat
open import Data.Star.Decoration
open import Data.Star.Properties
open import Data.Star.Vec
open ≡-Reasoning
open import Derivation

------------------------------------------------------------------------
-- Expressions

data Expr : Set where
  val : ℕ -> Expr
  _⊕_ : Expr -> Expr -> Expr

⟦_⟧ : Expr -> ℕ
⟦ val n   ⟧ = n
⟦ e₁ ⊕ e₂ ⟧ = ⟦ e₁ ⟧ + ⟦ e₂ ⟧

------------------------------------------------------------------------
-- Direct implementation and proof of correctness

module Direct where

  ----------------------------------------------------------------------
  -- Stack-machine code

  Stack : ℕ -> Set
  Stack = Vec ℕ

  data Op : ℕ -> ℕ -> Set where
    push : forall {n} -> ℕ -> Op n (1# + n)
    add  : forall {n} -> Op (2# + n) (1# + n)

  Code : ℕ -> ℕ -> Set
  Code = Star Op

  ⟪_⟫ : forall {m n} -> Code m n -> Stack m -> Stack n
  ⟪ ε            ⟫ s               = s
  ⟪ push n ◅ ops ⟫ s               = ⟪ ops ⟫ (↦ n ◅ s)
  ⟪ add ◅ ops    ⟫ (↦ m ◅ ↦ n ◅ s) = ⟪ ops ⟫ (↦ (n + m) ◅ s)

  ----------------------------------------------------------------------
  -- Compiler

  comp : forall {n} -> (e : Expr) -> Code n (1# + n)
  comp (val n)   = push n ◅ ε
  comp (e₁ ⊕ e₂) = comp e₁ ◅◅ comp e₂ ◅◅ add ◅ ε

  ----------------------------------------------------------------------
  -- Correctness proof

  distrib : forall {i j k}
                   (s : Stack i) (c₁ : Code i j) (c₂ : Code j k) ->
            ⟪ c₁ ◅◅ c₂ ⟫ s ≡ ⟪ c₂ ⟫ (⟪ c₁ ⟫ s)
  distrib s               ε             c₂ = refl
  distrib s               (push n ◅ c₁) c₂ = distrib (↦ n ◅ s) c₁ c₂
  distrib (↦ m ◅ ↦ n ◅ s) (add ◅ c₁)    c₂ =
    distrib (↦ (n + m) ◅ s) c₁ c₂

  thm1 : forall {n} (s : Stack n) e -> ⟪ comp e ⟫ s ≡ ↦ ⟦ e ⟧ ◅ s
  thm1 s (val n)   = refl
  thm1 s (e₁ ⊕ e₂) = begin
    ⟪ comp (e₁ ⊕ e₂) ⟫ s
      ≡⟨ distrib s (comp e₁) (comp e₂ ◅◅ add ◅ ε) ⟩
    ⟪ comp e₂ ◅◅ add ◅ ε ⟫ (⟪ comp e₁ ⟫ s)
      ≡⟨ cong ⟪ comp e₂ ◅◅ add ◅ ε ⟫ (thm1 s e₁) ⟩
    ⟪ comp e₂ ◅◅ add ◅ ε ⟫ (↦ ⟦ e₁ ⟧ ◅ s)
      ≡⟨ distrib (↦ ⟦ e₁ ⟧ ◅ s) (comp e₂) (add ◅ ε) ⟩
    ⟪ add ◅ ε ⟫ (⟪ comp e₂ ⟫ (↦ ⟦ e₁ ⟧ ◅ s))
      ≡⟨ cong ⟪ add ◅ ε ⟫ (thm1 (↦ ⟦ e₁ ⟧ ◅ s) e₂) ⟩
    ↦ ⟦ e₁ ⊕ e₂ ⟧ ◅ s
      ∎

  -- Graham and Joel claim that the following proof is shorter, but
  -- they don't mention their use of associativity of _◅◅_. (Note that
  -- a continuation-passing formulation of the compiler would obviate
  -- the need to prove associativity.)

  thm2 : forall {i j} (s : Stack i) e (ops : Code (1# + i) j) ->
         ⟪ comp e ◅◅ ops ⟫ s ≡ ⟪ ops ⟫ (↦ ⟦ e ⟧ ◅ s)
  thm2 s (val n)   ops = refl
  thm2 s (e₁ ⊕ e₂) ops = begin
    ⟪ comp (e₁ ⊕ e₂) ◅◅ ops ⟫ s
      ≡⟨ refl ⟩
    ⟪ (comp e₁ ◅◅ comp e₂ ◅◅ add ◅ ε) ◅◅ ops ⟫ s
      ≡⟨ cong (\ops' -> ⟪ ops' ⟫ s)
              (◅◅-assoc (comp e₁) (comp e₂ ◅◅ add ◅ ε) ops) ⟩
    ⟪ comp e₁ ◅◅ ((comp e₂ ◅◅ add ◅ ε) ◅◅ ops) ⟫ s
      ≡⟨ cong (\s' -> ⟪ comp e₁ ◅◅ s' ⟫ s)
              (◅◅-assoc (comp e₂) (add ◅ ε) ops) ⟩
    ⟪ comp e₁ ◅◅ (comp e₂ ◅◅ add ◅ ops) ⟫ s
      ≡⟨ thm2 s e₁ (comp e₂ ◅◅ add ◅ ops) ⟩
    ⟪ comp e₂ ◅◅ add ◅ ops ⟫ (↦ ⟦ e₁ ⟧ ◅ s)
      ≡⟨ thm2 (↦ ⟦ e₁ ⟧ ◅ s) e₂ (add ◅ ops) ⟩
    ⟪ ops ⟫ (↦ ⟦ e₁ ⊕ e₂ ⟧ ◅ s)
      ∎

------------------------------------------------------------------------
-- Calculation of an abstract machine

module Calculational where

  -- ⑴ Add continuations.

  module Step₁ where

    Cont : Set
    Cont = ℕ -> ℕ

    mutual

      -- Let us derive a function eval which satisfies
      --   eval e c ≡ c ⟦ e ⟧.
      -- The derivation is done in evalP, whose type guarantees
      -- correctness.

      eval : Expr -> Cont -> ℕ
      eval e c = witness (evalP e c)

      evalP : (e : Expr) (c : Cont) -> EqualTo (c ⟦ e ⟧)
      evalP (val n)   c = ▷ begin c n ∎
      evalP (e₁ ⊕ e₂) c = ▷ begin
        c (⟦ e₁ ⟧ + ⟦ e₂ ⟧)
          ≡⟨ refl ⟩
        (\(m : _) -> c (m + ⟦ e₂ ⟧)) ⟦ e₁ ⟧
          ≡⟨ proof (evalP e₁ (\m -> c (m + ⟦ e₂ ⟧))) ⟩
        eval e₁ (\m -> c (m + ⟦ e₂ ⟧))
          ≡⟨ refl ⟩
        eval e₁ (\m -> (\(n : _) -> c (m + n)) ⟦ e₂ ⟧)
          ≡⟨ cong (eval e₁)
                  (ext (\m -> proof (evalP e₂ (\n -> c (m + n))))) ⟩
        eval e₁ (\m -> eval e₂ (\n -> c (m + n)))
          ∎

    -- The code above contains some continuations. Let us name them.
    -- The identity continuation, used to connect eval to ⟦_⟧, is also
    -- named.

    c₁ : Cont
    c₁ = \n -> n

    c₃ : ℕ -> Cont -> Cont
    c₃ m c = \n -> c (m + n)

    c₂ : Expr -> Cont -> Cont
    c₂ e c = \m -> eval e (c₃ m c)

    -- The top-level evaluation function can be recovered by using the
    -- identity continuation.

    ⟦_⟧' : Expr -> ℕ
    ⟦ e ⟧' = eval e c₁

    correct : forall e -> ⟦ e ⟧ ≡ ⟦ e ⟧'
    correct e = proof (evalP e c₁)

  -- ⑵ Defunctionalise, i.e. get rid of the higher-order
  --   continuations.

  module Step₂ where

    data Cont : Set where
      c₁ : Cont
      c₂ : Expr -> Cont -> Cont
      c₃ : ℕ -> Cont -> Cont

    -- Define the semantics of the continuations in terms of the
    -- higher-order ones.

    ⟦_⟧c : Cont -> Step₁.Cont
    ⟦ c₁     ⟧c = Step₁.c₁
    ⟦ c₂ e c ⟧c = Step₁.c₂ e ⟦ c ⟧c
    ⟦ c₃ m c ⟧c = Step₁.c₃ m ⟦ c ⟧c

    mutual

      -- Define a variant of Step₁.eval which uses first-order
      -- continuations.

      eval : Expr -> Cont -> ℕ
      eval e c = witness (evalP e c)

      evalP : (e : Expr) (c : Cont) ->
              EqualTo (Step₁.eval e ⟦ c ⟧c)
      evalP (val n)   c = ▷ begin ⟦ c ⟧c n ∎
      evalP (e₁ ⊕ e₂) c = ▷ begin
        Step₁.eval e₁ (\m -> Step₁.eval e₂ (\n -> ⟦ c ⟧c (m + n)))
          ≡⟨ refl ⟩
        Step₁.eval e₁ (Step₁.c₂ e₂ ⟦ c ⟧c)
          ≡⟨ proof (evalP e₁ (c₂ e₂ c)) ⟩
        eval e₁ (c₂ e₂ c)
          ∎

    mutual

      -- Note that the code in this block does not satisfy the
      -- termination checker.

      -- Also make sure that the semantics of continuations does not
      -- refer to Step₁.eval (via Step₁.c₂).

      ⟦_⟧c' : Cont -> Step₁.Cont
      ⟦ c ⟧c' n = witness (cSem c n)

      eval' : Expr -> Cont -> ℕ
      eval' e c = witness (evalP' e c)

      cSem : (c : Cont) (n : ℕ) -> EqualTo (⟦ c ⟧c n)
      cSem c₁       n = ▷ begin n ∎
      cSem (c₂ e c) n = ▷ begin
        Step₁.eval e (Step₁.c₃ n ⟦ c ⟧c)
          ≡⟨ refl ⟩
        Step₁.eval e ⟦ c₃ n c ⟧c
          ≡⟨ proof (evalP e (c₃ n c)) ⟩
        eval e (c₃ n c)
          ≡⟨ proof (evalP' e (c₃ n c)) ⟩
        eval' e (c₃ n c)
          ∎
      cSem (c₃ m c) n = ▷ begin
        ⟦ c ⟧c (m + n)
          ≡⟨ proof (cSem c (m + n)) ⟩
        ⟦ c ⟧c' (m + n)
          ∎

      -- Note that eval above uses ⟦_⟧c, so it has to be updated.

      evalP' : (e : Expr) (c : Cont) ->
               EqualTo (eval e c)
      evalP' (val n) c = ▷ begin
        ⟦ c ⟧c n
          ≡⟨ proof (cSem c n) ⟩
        ⟦ c ⟧c' n
          ∎
      evalP' (e₁ ⊕ e₂) c = ▷ begin
        eval e₁ (c₂ e₂ c)
          ≡⟨ proof (evalP' e₁ (c₂ e₂ c)) ⟩
        eval' e₁ (c₂ e₂ c)
          ∎

    -- The top-level evaluation function can be recovered.

    ⟦_⟧' : Expr -> ℕ
    ⟦ e ⟧' = eval' e c₁

    correct : forall e -> ⟦ e ⟧ ≡ ⟦ e ⟧'
    correct e = begin
      ⟦ e ⟧
        ≡⟨ Step₁.correct e ⟩
      Step₁.⟦_⟧' e
        ≡⟨ proof (evalP e c₁) ⟩
      eval e c₁
        ≡⟨ proof (evalP' e c₁) ⟩
      ⟦ e ⟧'
        ∎

  -- ⑶ Refactor.

  module Step₃ where

    -- Note that the Hutton/Wright paper contains a slight mistake
    -- here: They define Step₂.eval by using Step₂.⟦_⟧c, but in this
    -- step they assume (with no comment) that Step₂.⟦_⟧c and
    -- Step₂.⟦_⟧c' are the same thing (they overload the name "apply",
    -- so this is hard to spot in the paper).
    --
    -- I tried to fix their mistake, but this makes the termination
    -- checker complain.

    open Step₂ public
      using (Cont; correct)
      renaming ( c₁ to STOP; c₂ to EVAL; c₃ to ADD
               ; eval' to eval; ⟦_⟧c' to exec; ⟦_⟧' to run
               )