{-# OPTIONS --safe #-}

module Problem6 where

open import Data.Nat
open import Data.Nat.Properties using (1+n≢0; +-assoc; +-identityʳ)
open import Data.Empty using (⊥-elim)
open import Data.Fin using (Fin; suc; raise; inject+)
open import Data.Fin.Patterns
open import Data.Fin.Properties using (toℕ-injective; toℕ-inject+)
open import Data.List
open import Data.List.Properties using (++-assoc; ++-identityʳ)
open import Data.Product using (_×_; _,_)
open import Relation.Binary.PropositionalEquality hiding ([_])

open ≡-Reasoning

data Expr : Set where
  Const : ℕ → Expr
  _⊞_ : Expr → Expr → Expr
  _⊠_ : Expr → Expr → Expr

eval : Expr → ℕ
eval (Const x) = x
eval (p ⊞ q) = eval p + eval q
eval (p ⊠ q) = eval p * eval q

data Instr : Set where
  Push : ℕ → Instr
  Mult Add : Instr

instr-semantics : Instr → List ℕ → List ℕ
instr-semantics (Push x) s       = x ∷ s
instr-semantics Mult (y ∷ x ∷ s) = x * y ∷ s
instr-semantics Add  (y ∷ x ∷ s) = x + y ∷ s
instr-semantics _ _ = []

execute : List Instr → List ℕ → List ℕ
execute []       s = s
execute (i ∷ is) s = execute is (instr-semantics i s)

compile : Expr → List Instr
compile (Const x) = [ Push x ]
compile (e₁ ⊞ e₂) = compile e₁ ++ compile e₂ ++ [ Add  ]
compile (e₁ ⊠ e₂) = compile e₁ ++ compile e₂ ++ [ Mult ]

-- Use difference lists because they are nicer to work with.

compile′ : Expr → List Instr → List Instr
compile′ (Const x) is = Push x ∷ is
compile′ (e₁ ⊞ e₂) is = compile′ e₁ (compile′ e₂ (Add ∷ is))
compile′ (e₁ ⊠ e₂) is = compile′ e₁ (compile′ e₂ (Mult ∷ is))

compile≗compile′ : ∀ e is → compile e ++ is ≡ compile′ e is
compile≗compile′ (Const x) is = refl
compile≗compile′ (e₁ ⊞ e₂) is = begin
  (compile e₁ ++ compile e₂ ++ [ Add ]) ++ is ≡⟨ ++-assoc (compile e₁) (compile e₂ ++ [ Add ]) is ⟩
  compile e₁ ++ (compile e₂ ++ [ Add ]) ++ is ≡⟨ cong (compile e₁ ++_) (++-assoc (compile e₂) [ Add ] is) ⟩
  compile e₁ ++ compile e₂ ++ Add ∷ is        ≡⟨ cong (compile e₁ ++_) (compile≗compile′ e₂ (Add ∷ is)) ⟩
  compile e₁ ++ compile′ e₂ (Add ∷ is)        ≡⟨ compile≗compile′ e₁ (compile′ e₂ (Add ∷ is)) ⟩
  compile′ e₁ (compile′ e₂ (Add ∷ is))        ∎
compile≗compile′ (e₁ ⊠ e₂) is = begin
  (compile e₁ ++ compile e₂ ++ [ Mult ]) ++ is ≡⟨ ++-assoc (compile e₁) (compile e₂ ++ [ Mult ]) is ⟩
  compile e₁ ++ (compile e₂ ++ [ Mult ]) ++ is ≡⟨ cong (compile e₁ ++_) (++-assoc (compile e₂) [ Mult ] is) ⟩
  compile e₁ ++ compile e₂ ++ Mult ∷ is        ≡⟨ cong (compile e₁ ++_) (compile≗compile′ e₂ (Mult ∷ is)) ⟩
  compile e₁ ++ compile′ e₂ (Mult ∷ is)        ≡⟨ compile≗compile′ e₁ (compile′ e₂ (Mult ∷ is)) ⟩
  compile′ e₁ (compile′ e₂ (Mult ∷ is))        ∎

-- Generalise expressions to open terms

data Expr′ (n : ℕ) : Set where
  Var   : Fin n → Expr′ n
  Const : ℕ → Expr′ n
  _⊞_   : Expr′ n → Expr′ n → Expr′ n
  _⊠_   : Expr′ n → Expr′ n → Expr′ n

forget : Expr → Expr′ 0
forget (Const x) = Const x
forget (e₁ ⊞ e₂) = forget e₁ ⊞ forget e₂
forget (e₁ ⊠ e₂) = forget e₁ ⊠ forget e₂

eval′ : ∀ {n} → Expr′ n → (Fin n → ℕ) → ℕ
eval′ (Var x)   γ = γ x
eval′ (Const x) γ = x
eval′ (e₁ ⊞ e₂) γ = eval′ e₁ γ + eval′ e₂ γ
eval′ (e₁ ⊠ e₂) γ = eval′ e₁ γ * eval′ e₂ γ

eval′-correct : (e : Expr) → eval′ (forget e) (λ ()) ≡ eval e
eval′-correct (Const x) = refl
eval′-correct (e₁ ⊞ e₂) = cong₂ _+_ (eval′-correct e₁) (eval′-correct e₂)
eval′-correct (e₁ ⊠ e₂) = cong₂ _*_ (eval′-correct e₁) (eval′-correct e₂)

-- We can substitute free variables.

wkn : ∀ {k} n → Expr′ k → Expr′ (k + n)
wkn n (Var x)   = Var (inject+ n x)
wkn n (Const x) = Const x
wkn n (e₁ ⊞ e₂) = wkn n e₁ ⊞ wkn n e₂
wkn n (e₁ ⊠ e₂) = wkn n e₁ ⊠ wkn n e₂

sub : ∀ {k n} → Expr′ k → Expr′ (suc n) → Expr′ (k + n)
sub         e′ (Var 0F)      = wkn _ e′
sub {k = k} e′ (Var (suc i)) = Var (raise k i)
sub         e′ (Const x)     = Const x
sub         e′ (e₁ ⊞ e₂)     = sub e′ e₁ ⊞ sub e′ e₂
sub         e′ (e₁ ⊠ e₂)     = sub e′ e₁ ⊠ sub e′ e₂

-- These facts generalise to `n` variables, but the types are nasty

wkn0 : (e : Expr′ 0) → wkn 0 e ≡ e
wkn0 (Const x) = refl
wkn0 (e₁ ⊞ e₂) = cong₂ _⊞_ (wkn0 e₁) (wkn0 e₂)
wkn0 (e₁ ⊠ e₂) = cong₂ _⊠_ (wkn0 e₁) (wkn0 e₂)

wkn-assoc : ∀ k n (e : Expr′ 0) → wkn (k + n) e ≡ wkn n (wkn k e)
wkn-assoc k n (Const x) = refl
wkn-assoc k n (e₁ ⊞ e₂) = cong₂ _⊞_ (wkn-assoc k n e₁) (wkn-assoc k n e₂)
wkn-assoc k n (e₁ ⊠ e₂) = cong₂ _⊠_ (wkn-assoc k n e₁) (wkn-assoc k n e₂)

sub-wkn : (e′ : Expr′ 0) (e : Expr′ 0) → sub e′ (wkn 1 e) ≡ e
sub-wkn e′ (Const x) = refl
sub-wkn e′ (e₁ ⊞ e₂) = cong₂ _⊞_ (sub-wkn e′ e₁) (sub-wkn e′ e₂)
sub-wkn e′ (e₁ ⊠ e₂) = cong₂ _⊠_ (sub-wkn e′ e₁) (sub-wkn e′ e₂)

wkn-sub : ∀ {k} n (e′ : Expr′ 0) (e : Expr′ (suc k)) → sub e′ (wkn n e) ≡ wkn n (sub e′ e)
wkn-sub n e′ (Const x) = refl
wkn-sub n e′ (e₁ ⊞ e₂) = cong₂ _⊞_ (wkn-sub n e′ e₁) (wkn-sub n e′ e₂)
wkn-sub n e′ (e₁ ⊠ e₂) = cong₂ _⊠_ (wkn-sub n e′ e₁) (wkn-sub n e′ e₂)
wkn-sub n e′ (Var 0F)      = wkn-assoc _ n e′
wkn-sub n e′ (Var (suc i)) = refl

sub-assoc :
  ∀ {m n} (e₁ : Expr′ 0) (e₂ : Expr′ (suc m)) (e₃ : Expr′ (suc n)) →
  sub e₁ (sub e₂ e₃) ≡ sub (sub e₁ e₂) e₃
sub-assoc e₁ e₂ (Const x) = refl
sub-assoc e₁ e₂ (e₃ ⊞ e₄) = cong₂ _⊞_ (sub-assoc e₁ e₂ e₃) (sub-assoc e₁ e₂ e₄)
sub-assoc e₁ e₂ (e₃ ⊠ e₄) = cong₂ _⊠_ (sub-assoc e₁ e₂ e₃) (sub-assoc e₁ e₂ e₄)
sub-assoc e₁ e₂ (Var 0F)      = wkn-sub _ e₁ e₂
sub-assoc e₁ e₂ (Var (suc i)) = refl

-- I want to work only on well-formed instruction sequences. A sequence "eats"
-- `n` arguments if correct execution requires `n` elements on the stack.
--
-- The arity of an instruction gives how many arguments it consumes.

arity : Instr → ℕ
arity (Push x) = 0
arity Add      = 2
arity Mult     = 2

data WellFormed : List Instr → ℕ → Set where
  []   : WellFormed [] 1
  _∷_  : ∀ {is n} → (i : Instr) → WellFormed is (suc n) → WellFormed (i ∷ is) (arity i + n)

-- Compilation only produces well-formed sequences that produce one element.

compile′-wf : ∀ {n is} (e : Expr) → WellFormed is (suc n) → WellFormed (compile′ e is) n
compile′-wf (Const x) wf = Push x ∷ wf
compile′-wf (e₁ ⊞ e₂) wf = compile′-wf e₁ (compile′-wf e₂ (Add ∷ wf))
compile′-wf (e₁ ⊠ e₂) wf = compile′-wf e₁ (compile′-wf e₂ (Mult ∷ wf))

-- We can form open expressions from instruction sequences.

decompile₁ : (i : Instr) → Expr′ (arity i)
decompile₁ (Push x) = Const x
decompile₁ Add      = Var 1F ⊞ Var 0F
decompile₁ Mult     = Var 1F ⊠ Var 0F

decompile : ∀ {n is} → WellFormed is n → Expr′ n
decompile []       = Var 0F
decompile (i ∷ wf) = sub (decompile₁ i) (decompile wf)

-- Provide semantics for function environments

instr-semantics′ : ∀ {n} (i : Instr) → (Fin (arity i + n) → ℕ) → Fin (suc n) → ℕ
instr-semantics′ i        f (suc x) = f (raise (arity i) x)
instr-semantics′ (Push x) f 0F      = x
instr-semantics′ Add      f 0F      = f 1F + f 0F
instr-semantics′ Mult     f 0F      = f 1F * f 0F

instr-semantics-cong :
  ∀ (i : Instr) {n xs} {f : Fin (arity i + n) → ℕ} →
  xs ≡ tabulate f → instr-semantics i xs ≡ tabulate (instr-semantics′ i f)
instr-semantics-cong (Push x) refl = refl
instr-semantics-cong Add      refl = refl
instr-semantics-cong Mult     refl = refl

-- Prove decompilation is correct

eval-sub :
  ∀ (i : Instr) {n} (f : Fin (arity i + n) → ℕ) (e : Expr′ (suc n)) →
  eval′ e (instr-semantics′ i f) ≡ eval′ (sub (decompile₁ i) e) f
eval-sub i f (Const x)     = refl
eval-sub i f (e₁ ⊞ e₂)     = cong₂ _+_ (eval-sub i f e₁) (eval-sub i f e₂)
eval-sub i f (e₁ ⊠ e₂)     = cong₂ _*_ (eval-sub i f e₁) (eval-sub i f e₂)
eval-sub i f (Var (suc x)) = refl
eval-sub (Push x) f (Var 0F) = refl
eval-sub Add      f (Var 0F) = refl
eval-sub Mult     f (Var 0F) = refl

eval-decompile :
  ∀ {n is} (wf : WellFormed is n) (xs : List ℕ) (f : Fin n → ℕ) →
  xs ≡ tabulate f → execute is xs ≡ [ eval′ (decompile wf) f ]
eval-decompile [] xs f xs≡f = xs≡f
eval-decompile {is = .i ∷ is} (i ∷ wf) xs f xs≡f = begin
  execute is (instr-semantics i xs)               ≡⟨ eval-decompile wf (instr-semantics i xs) (instr-semantics′ i f) (instr-semantics-cong i xs≡f) ⟩
  [ eval′ (decompile wf) (instr-semantics′ i f) ] ≡⟨ cong [_] (eval-sub i f (decompile wf)) ⟩
  [ eval′ (sub (decompile₁ i) (decompile wf)) f ] ∎

-- Prove decompilation is an inverse

decompile-compile′ : ∀ {is n} (e : Expr) (wf : WellFormed is (suc n)) → decompile (compile′-wf e wf) ≡ sub (forget e) (decompile wf)
decompile-compile′ (Const x) wf = refl
decompile-compile′ (e₁ ⊞ e₂) wf = begin
  decompile (compile′-wf e₁ (compile′-wf e₂ (Add ∷ wf)))                       ≡⟨ decompile-compile′ e₁ (compile′-wf e₂ (Add ∷ wf)) ⟩
  sub (forget e₁) (decompile (compile′-wf e₂ (Add ∷ wf)))                      ≡⟨ cong (sub (forget e₁)) (decompile-compile′ e₂ (Add ∷ wf)) ⟩
  sub (forget e₁) (sub (forget e₂) (sub (Var 1F ⊞ Var 0F) (decompile wf)))     ≡⟨ cong (sub (forget e₁)) (sub-assoc (forget e₂) (Var 1F ⊞ Var 0F) (decompile wf)) ⟩
  sub (forget e₁) (sub (Var 0F ⊞ wkn 1 (forget e₂)) (decompile wf))            ≡⟨ sub-assoc (forget e₁) (Var 0F ⊞ wkn 1 (forget e₂)) (decompile wf) ⟩
  sub (wkn 0 (forget e₁) ⊞ sub (forget e₁) (wkn 1 (forget e₂))) (decompile wf) ≡⟨ cong₂ (λ -₁ -₂ → sub (-₁ ⊞ -₂) (decompile wf)) (wkn0 (forget e₁)) (sub-wkn (forget e₁) (forget e₂)) ⟩
  sub (forget e₁ ⊞ forget e₂) (decompile wf)                                   ∎
decompile-compile′ (e₁ ⊠ e₂) wf = begin
  decompile (compile′-wf e₁ (compile′-wf e₂ (Mult ∷ wf)))                      ≡⟨ decompile-compile′ e₁ (compile′-wf e₂ (Mult ∷ wf)) ⟩
  sub (forget e₁) (decompile (compile′-wf e₂ (Mult ∷ wf)))                     ≡⟨ cong (sub (forget e₁)) (decompile-compile′ e₂ (Mult ∷ wf)) ⟩
  sub (forget e₁) (sub (forget e₂) (sub (Var 1F ⊠ Var 0F) (decompile wf)))     ≡⟨ cong (sub (forget e₁)) (sub-assoc (forget e₂) (Var 1F ⊠ Var 0F) (decompile wf)) ⟩
  sub (forget e₁) (sub (Var 0F ⊠ wkn 1 (forget e₂)) (decompile wf))            ≡⟨ sub-assoc (forget e₁) (Var 0F ⊠ wkn 1 (forget e₂)) (decompile wf) ⟩
  sub (wkn 0 (forget e₁) ⊠ sub (forget e₁) (wkn 1 (forget e₂))) (decompile wf) ≡⟨ cong₂ (λ -₁ -₂ → sub (-₁ ⊠ -₂) (decompile wf)) (wkn0 (forget e₁)) (sub-wkn (forget e₁) (forget e₂)) ⟩
  sub (forget e₁ ⊠ forget e₂) (decompile wf)                                   ∎

-- Prove evaluation correct.

compile-correct : ∀ e → execute (compile e) [] ≡ [ eval e ]
compile-correct e = begin
  execute (compile e) []                          ≡˘⟨ cong (λ - → execute - []) (++-identityʳ (compile e)) ⟩
  execute (compile e ++ []) []                    ≡⟨ cong (λ - → execute - []) (compile≗compile′ e []) ⟩
  execute (compile′ e []) []                      ≡⟨ eval-decompile (compile′-wf e []) [] (λ ()) refl ⟩
  [ eval′ (decompile (compile′-wf e [])) (λ ()) ] ≡⟨ cong (λ - → [ eval′ - (λ ()) ]) (decompile-compile′ e []) ⟩
  [ eval′ (wkn 0 (forget e)) (λ ()) ]             ≡⟨ cong (λ - → [ eval′ - (λ ()) ]) (wkn0 (forget e)) ⟩
  [ eval′ (forget e) (λ ()) ]                     ≡⟨ cong [_] (eval′-correct e) ⟩
  [ eval e ]                                      ∎