{-# OPTIONS --guardedness #-}
module Lang where

open import Data.Integer using (ℤ; _+_; +_)
open import Data.Nat using (ℕ; suc; _<_)
open import Data.Nat.Properties using (m≤n⇒m≤o+n; ≤-trans; +-monoʳ-≤; ≤-refl; n≮n)
open import Data.Bool using (Bool; true; false)
open import Data.List using (List; _∷_; []; length)
open import Data.Unit using (⊤; tt)
open import Data.Empty using (⊥; ⊥-elim)
open import Data.Product using (Σ; _×_; _,_)

data Value : Set where
    int : ℤ → Value
    bool : Bool → Value

data Instruction : Set where
    pushint : ℤ → Instruction
    pushbool : Bool → Instruction
    add : Instruction

    loop : ℕ → Instruction
    break : Instruction
    end : Instruction

Program : Set
Program = List Instruction

Env : Set
Env = List (List Instruction)

variable
    i : Instruction
    is is' is'' : Program
    S S' S'' : Env
    s s' s'' : List Value
    n : ℕ
    b : Bool
    z z₁ z₂ : ℤ

data _⇒ₓ_ : Program → Program → Set where
    pushint-skip : is ⇒ₓ is' → (pushint z ∷ is) ⇒ₓ is'
    pushbool-skip : is ⇒ₓ is' → (pushbool b ∷ is) ⇒ₓ is'
    add-skip :     is ⇒ₓ is' → (add ∷ is) ⇒ₓ is'
    break-skip :   is ⇒ₓ is' → (break ∷ is) ⇒ₓ is'
    loop-skip :    is ⇒ₓ is' → is' ⇒ₓ is'' → (loop n ∷ is) ⇒ₓ is''
    end-skip :     (end ∷ is) ⇒ₓ is

-- Skipping always drops instructions, because it at least consumes 'end'
skip-shrink : is ⇒ₓ is' → length is' < length is
skip-shrink (pushint-skip rec) = m≤n⇒m≤o+n 1 (skip-shrink rec)
skip-shrink (pushbool-skip rec) = m≤n⇒m≤o+n 1 (skip-shrink rec)
skip-shrink (add-skip rec) = m≤n⇒m≤o+n 1 (skip-shrink rec)
skip-shrink (break-skip rec) = m≤n⇒m≤o+n 1 (skip-shrink rec)
skip-shrink (loop-skip is⇒ₓis' is'⇒ₓis'') =
    m≤n⇒m≤o+n 1 (≤-trans (m≤n⇒m≤o+n 1 (skip-shrink is'⇒ₓis''))
                         (skip-shrink is⇒ₓis'))
skip-shrink end-skip = ≤-refl

-- Since skipping consumes end, skipping can't return the input program
skip-not-noop : is ⇒ₓ is → ⊥
skip-not-noop {is} is⇒ₓis = n≮n (length is) (skip-shrink is⇒ₓis)

Triple : Set
Triple = List Value × Program × Env

-- Small-step evaluation relation, paremeterized by stack, program, and continuation list
-- Each 'loop' introduces a continuation (what to do when the loop is over).
data _⇒_ : Triple → Triple → Set where
    pushint-eval : ( s , (pushint z ∷ is) , S ) ⇒ ( ((int z) ∷ s) , is , S )
    pushbool-eval : ( s , (pushbool b ∷ is) , S ) ⇒ ( ((bool b) ∷ s) , is , S )
    add-eval : ( (int z₁ ∷ int z₂ ∷ s) , (add ∷ is), S ) ⇒ ( ((int (z₁ + z₂)) ∷ s) , is , S )

    loop-eval⁰ : is ⇒ₓ is' →
                 ( s , (loop 0 ∷ is) , S ) ⇒ ( s , is' , S )
    loop-evalⁿ : ( s , (loop (suc n) ∷ is) ,  S ) ⇒ ( s , is , ((loop n ∷ is) ∷ S) )
    end-eval : ( s , (end ∷ is) , (is' ∷ S) ) ⇒ ( s , is' , S )
    break-evalᶠ : ( (bool false ∷ s) , (break ∷ is) , S ) ⇒ ( s , is , S )
    break-evalᵗ : is ⇒ₓ is' →
                  ( (bool true ∷ s) , (break ∷ is) , (is'' ∷ S) ) ⇒ ( s , is' , S )

-- Transitive evaluation relation is trivial
infixr 30 _thenᵉ_
data _⇒*_ : Triple → Triple → Set where
    doneᵉ : ( s , [] , S  ) ⇒* ( s , [] , S )
    _thenᵉ_ : ( s , is , S ) ⇒ ( s' , is' , S' ) →
             ( s' , is' , S' ) ⇒* ( s'' , is'' , S'' ) →
             ( s , is , S ) ⇒* ( s'' , is'' , S'' )

-- Some examples of programs, though (2) and (3) don't maintain well-typed stacks
basic₁ : ( [] , (pushint (+ 1) ∷ pushint (+ 2) ∷ pushint (+ 3) ∷ add ∷ add ∷ []) , [] ) ⇒* ( (int (+ 6) ∷ []) , [] , [] )
basic₁ = pushint-eval thenᵉ pushint-eval thenᵉ pushint-eval thenᵉ add-eval thenᵉ add-eval thenᵉ doneᵉ

basic₂ : ( [] , (loop 1 ∷ pushbool true ∷ break ∷ pushint (+ 1) ∷ end ∷ []) , [] ) ⇒* ( [] , [] , [] )
basic₂ = loop-evalⁿ thenᵉ pushbool-eval thenᵉ break-evalᵗ (pushint-skip end-skip) thenᵉ doneᵉ

basic₃ : ( [] , (loop 1 ∷ pushbool false ∷ break ∷ pushint (+ 1) ∷ end ∷ []) , [] ) ⇒* ( (int (+ 1) ∷ []) , [] , [] )
basic₃ = loop-evalⁿ thenᵉ pushbool-eval thenᵉ break-evalᶠ thenᵉ pushint-eval thenᵉ end-eval thenᵉ (loop-eval⁰ (pushbool-skip (break-skip (pushint-skip end-skip)))) thenᵉ doneᵉ

data ValueType : Set where
    tint tbool : ValueType

-- A stack type describes the current stack and the continuation types
infix 35 _⟨_⟩
record StackType : Set where
    constructor _⟨_⟩
    field
        current : List ValueType
        frames : List (List ValueType)

variable
    st st' st'' : List ValueType
    F F' F'' : List (List ValueType)
    σ σ' σ'' : StackType

-- Typing relation for a single instruction.
-- Importantly, loops are supposed to preserve the type of the stack once
-- they exit (because if you loop 0, the stack doesn't change). This means
-- that 'break' must only be valid when the current stack type matches
-- the after-loop continuation type (and thus the pre-loop type). The same
-- is true for 'end', since the next iteration of the loop expects the same
-- stack type as the first.
data _⊢_↦_ : StackType → Instruction → StackType → Set where
    pushint-t : (st ⟨ F ⟩) ⊢ (pushint z) ↦ ((tint ∷ st) ⟨ F ⟩)
    pushbool-t : (st ⟨ F ⟩) ⊢ (pushbool b) ↦ ((tbool ∷ st) ⟨ F ⟩)
    add-t : ((tint ∷ tint ∷ st) ⟨ F ⟩) ⊢ add ↦ ((tint ∷ st) ⟨ F ⟩)

    loop-t : (st ⟨ F ⟩) ⊢ (loop n) ↦ (st ⟨ st ∷ F ⟩)
    end-t : (st ⟨ st ∷ F ⟩) ⊢ end ↦ (st ⟨ F ⟩)
    break-t : ((tbool ∷ st) ⟨ st ∷ F ⟩) ⊢ break ↦ (st ⟨ st ∷ F ⟩)

-- Transitive typing relation is also trivial
infixr 30 _thenᵗ_
data _⊢_↦*_ : StackType → Program → StackType → Set where
    doneᵗ : (st ⟨ F ⟩) ⊢ [] ↦* (st ⟨ F ⟩)
    _thenᵗ_ : (st ⟨ F ⟩) ⊢ i ↦ (st' ⟨ F' ⟩) →
              (st' ⟨ F' ⟩) ⊢ is ↦* (st'' ⟨ F'' ⟩) →
              (st ⟨ F ⟩) ⊢ (i ∷ is) ↦* (st'' ⟨ F'' ⟩)

-- A "well-typed" program which terminates without leftover breaks/ends etc.
_⊢_↦|_ : StackType → Program → List ValueType → Set
_⊢_↦|_ σ is st = σ ⊢ is ↦* st ⟨ [] ⟩

-- Only the first program is well-typed, the others violate the loop-stack
-- consistency principle described above.
basict₁ : [] ⟨ [] ⟩ ⊢ (pushint (+ 1) ∷ pushint (+ 2) ∷ pushint (+ 3) ∷ add ∷ add ∷ []) ↦* (tint ∷ []) ⟨ [] ⟩
basict₁ = pushint-t thenᵗ pushint-t thenᵗ pushint-t thenᵗ add-t thenᵗ add-t thenᵗ doneᵗ

-- doesn't typecheck: violates type discipline because on break, stack has different type than on end.
--
-- basict₂ : ([] ⟨ []  ⟩) ⊢ (loop 1 ∷ pushbool true ∷ break ∷ pushint (+ 1) ∷ end ∷ []) ↦* ([] ⟨ [] ⟩)
-- basict₂ = loop-t thenᵗ pushbool-t thenᵗ break-t thenᵗ pushint-t thenᵗ end-t thenᵗ doneᵗ

-- Fast forwarding to a 'break' will peel off one continuation (since that
-- continuation encodes the post-loop type).
lemma₁ : ∀ {lst} → st ⟨ lst ∷ F ⟩ ⊢ is ↦* σ → is ⇒ₓ is' → lst ⟨ F ⟩ ⊢ is' ↦* σ
lemma₁ doneᵗ ()
lemma₁ (pushint-t thenᵗ stⁱ⊢is↦st') (pushint-skip isⁱ⇒ₓis') = lemma₁ stⁱ⊢is↦st' isⁱ⇒ₓis'
lemma₁ (pushbool-t thenᵗ stⁱ⊢is↦st') (pushbool-skip isⁱ⇒ₓis') = lemma₁ stⁱ⊢is↦st' isⁱ⇒ₓis'
lemma₁ (add-t thenᵗ stⁱ⊢is↦st') (add-skip isⁱ⇒ₓis') = lemma₁ stⁱ⊢is↦st' isⁱ⇒ₓis'
lemma₁ (loop-t thenᵗ stⁱ⊢is↦st') (loop-skip isⁱ⇒ₓisⁱ' isⁱ'⇒ₓis') = lemma₁ (lemma₁ stⁱ⊢is↦st' isⁱ⇒ₓisⁱ') isⁱ'⇒ₓis'
lemma₁ (end-t thenᵗ stⁱ⊢is↦st') end-skip = stⁱ⊢is↦st'
lemma₁ (break-t thenᵗ stⁱ⊢is↦st') (break-skip isⁱ⇒ₓis') = lemma₁ stⁱ⊢is↦st' isⁱ⇒ₓis'

-- A well-typed, terminating program in which there is a current continuation (st')
-- can be fast-forwarded, since a continuation means a loop was encountered,
-- and well-typedness implies loops are all closed.
{-# TERMINATING #-} -- each returned instruction list is smaller (skip-shrink), so it's okay to keep using them.
lemma₂ : st ⟨ st' ∷ F ⟩ ⊢ is ↦| st'' → Σ Program (λ is' → is ⇒ₓ is' × (st' ⟨ F ⟩ ⊢ is' ↦| st''))
lemma₂ (pushint-t thenᵗ ist⊢is↦|st'') =
    let (is' , is''⇒ₓis' , is'↦|st'') = lemma₂ ist⊢is↦|st'' in
    (_ , pushint-skip is''⇒ₓis' , is'↦|st'')
lemma₂ (pushbool-t thenᵗ ist⊢is↦|st'') =
    let (is' , is''⇒ₓis' , is'↦|st'') = lemma₂ ist⊢is↦|st'' in
    (_ , pushbool-skip is''⇒ₓis' , is'↦|st'')
lemma₂ (add-t thenᵗ ist⊢is↦|st'') =
    let (is' , is''⇒ₓis' , is'↦|st'') = lemma₂ ist⊢is↦|st'' in
    (_ , add-skip is''⇒ₓis' , is'↦|st'')
lemma₂ (break-t thenᵗ ist⊢is↦|st'') =
    let (is' , is''⇒ₓis' , is'↦|st'') = lemma₂ ist⊢is↦|st'' in
    (_ , break-skip is''⇒ₓis' , is'↦|st'')
lemma₂ (loop-t thenᵗ ist⊢is↦|st'') =
    let (is'₁ , is''⇒ₓis' , is'₁↦|st'') = lemma₂ ist⊢is↦|st'' in
    let (is'₂ , is'₁⇒ₓis'₂ , is'₂↦|st'') = lemma₂ is'₁↦|st'' in
    (_ , loop-skip is''⇒ₓis' is'₁⇒ₓis'₂  , is'₂↦|st'')
lemma₂ (end-t thenᵗ rest) = (_ , end-skip , rest)

-- what it means for list of value types to describe a list of values
data ⟦_⟧ˢ : List ValueType → List Value → Set where
    ⟦⟧ˢ-empty : ⟦ [] ⟧ˢ []
    ⟦⟧ˢ-int : ⟦ st ⟧ˢ s → ⟦ tint ∷ st ⟧ˢ (int z ∷ s)
    ⟦⟧ˢ-bool : ⟦ st ⟧ˢ s → ⟦ tbool ∷ st ⟧ˢ (bool b ∷ s)

-- Parameterize by final stack type. The reason to do this is that we want
-- all continuations to end in the same type, so that control flow doesn't change
-- what final type the program is evaluated to. This is an unfortunate consequence
-- of having to include copies of the rest of the program as part of execution.
module _ (stˡ : List ValueType) where
    data ⟦_⟧ᶠ : List (List ValueType) → Program → Env → Set where
        ⟦⟧ᶠ-empty : ⟦ [] ⟧ᶠ is []
        ⟦⟧ᶠ-step : (∀ {is'} → is ⇒ₓ is' → ⟦ F ⟧ᶠ is' S) → -- Loop invariant one: breaking out of the loop is well-typed
                   ⟦ F ⟧ᶠ is' S →                         -- Loop invariant two: continuing once reaching loop end is well-typed
                   st ⟨ F ⟩ ⊢ is' ↦| stˡ  →               -- The continuation is well-typed
                   ⟦ st ∷ F ⟧ᶠ is (is' ∷ S)

    -- Some helper properties. A non-branching instruction preserves the validity
    -- of the current stack type (map), the number of iterations doesn't
    -- affect the validity of a stcack type given a loop (dec), and seeking
    -- past matched loop-end tokens doesn't affect the validity of a stack type (skip).
    private
        map : (s , (i ∷ is) , S) ⇒ (s' , is , S) →
              ⟦ F ⟧ᶠ (i ∷ is) S →
              ⟦ F ⟧ᶠ is S
        map {F = []} _ (⟦⟧ᶠ-empty) = ⟦⟧ᶠ-empty
        map (pushint-eval) (⟦⟧ᶠ-step is⇒ₓok ⟦F'⟧ is'↦|) = ⟦⟧ᶠ-step (λ is⇒ₓis' → is⇒ₓok (pushint-skip is⇒ₓis')) ⟦F'⟧ is'↦|
        map (pushbool-eval) (⟦⟧ᶠ-step is⇒ₓok ⟦F'⟧ is'↦|) = ⟦⟧ᶠ-step (λ is⇒ₓis' → is⇒ₓok (pushbool-skip is⇒ₓis')) ⟦F'⟧ is'↦|
        map (add-eval) (⟦⟧ᶠ-step is⇒ₓok ⟦F'⟧ is'↦|) = ⟦⟧ᶠ-step (λ is⇒ₓis' → is⇒ₓok (add-skip is⇒ₓis')) ⟦F'⟧ is'↦|
        map (break-evalᶠ) (⟦⟧ᶠ-step is⇒ₓok ⟦F'⟧ is'↦|) = ⟦⟧ᶠ-step (λ is⇒ₓis' → is⇒ₓok (break-skip is⇒ₓis')) ⟦F'⟧ is'↦|
        map (loop-eval⁰ is⇒ₓis) _ = ⊥-elim (skip-not-noop is⇒ₓis)

        dec : ⟦ F ⟧ᶠ (loop (suc n) ∷ is) S →
              ⟦ F ⟧ᶠ (loop n ∷ is) S
        dec ⟦⟧ᶠ-empty = ⟦⟧ᶠ-empty
        dec (⟦⟧ᶠ-step is⇒ₓok ⟦F'⟧ is'↦|) = ⟦⟧ᶠ-step (λ { (loop-skip is⇒ₓis' is'⇒ₓis'') → is⇒ₓok (loop-skip is⇒ₓis' is'⇒ₓis'') }) ⟦F'⟧ is'↦|

        skip : ⟦ F ⟧ᶠ (loop n ∷ is) S →
               is ⇒ₓ is' →
               ⟦ F ⟧ᶠ is' S
        skip ⟦⟧ᶠ-empty _ = ⟦⟧ᶠ-empty
        skip (⟦⟧ᶠ-step f t rec) is⇒ₓis' = ⟦⟧ᶠ-step (λ is'⇒ₓis'' → f (loop-skip is⇒ₓis' is'⇒ₓis'')) t rec

    ⟦_⟧ : StackType → List Value → Program → List Program → Set
    ⟦_⟧ (st ⟨ F ⟩) s is S = ⟦ st ⟧ˢ s × ⟦ F ⟧ᶠ is S × (st ⟨ F ⟩ ⊢ is ↦* stˡ ⟨ [] ⟩)

    preservation : (s , is , S) ⇒ (s' , is' , S') →
                   ⟦ σ ⟧ s is S →
                   Σ StackType (λ σ'' → ⟦ σ'' ⟧ s' is' S')
    preservation e@pushint-eval (⟦st⟧ , ⟦F⟧ , pushint-t thenᵗ is↦*) = (_ , (⟦⟧ˢ-int ⟦st⟧ , map e ⟦F⟧ , is↦*))
    preservation e@pushbool-eval (⟦st⟧ , ⟦F⟧ , pushbool-t thenᵗ is↦*) = (_ , (⟦⟧ˢ-bool ⟦st⟧ , map e ⟦F⟧ , is↦*))
    preservation e@add-eval ((⟦⟧ˢ-int (⟦⟧ˢ-int ⟦st⟧)) , ⟦F⟧ , add-t thenᵗ is↦*) = (_ , (⟦⟧ˢ-int ⟦st⟧ , map e ⟦F⟧ , is↦*))
    preservation (loop-eval⁰ is⇒ₓis') (⟦st⟧ , ⟦F⟧ , loop-t thenᵗ is↦*) =  (_ , (⟦st⟧ , skip ⟦F⟧ is⇒ₓis' , lemma₁ is↦* is⇒ₓis'))
    preservation loop-evalⁿ (⟦st⟧ , ⟦F⟧ , loop-t thenᵗ is↦*) = (_ , (⟦st⟧ , ⟦⟧ᶠ-step (λ is⇒ₓis'' → skip ⟦F⟧ is⇒ₓis'') (dec ⟦F⟧) (loop-t thenᵗ is↦*) , is↦*))
    preservation end-eval (⟦st⟧ , (⟦⟧ᶠ-step _ ⟦F⟧ t) , end-t thenᵗ is↦*) = (_ , (⟦st⟧ , ⟦F⟧ , t))
    preservation  e@break-evalᶠ (⟦⟧ˢ-bool ⟦st⟧ , ⟦F⟧ , break-t thenᵗ is↦*) = (_ , (⟦st⟧ , map e ⟦F⟧ , is↦*))
    preservation (break-evalᵗ is⇒ₓis') (⟦⟧ˢ-bool ⟦st⟧ , (⟦⟧ᶠ-step is⇒ₓis'⇒ok ⟦F⟧ t) , break-t thenᵗ is↦*) = (_ , (⟦st⟧ , is⇒ₓis'⇒ok (break-skip is⇒ₓis') , lemma₁ is↦* is⇒ₓis'))

    progress : ⟦ σ ⟧ s (i ∷ is) S → Σ (List Value × Program × Env) ((s , (i ∷ is) , S) ⇒_)
    progress (⟦st⟧ , ⟦F⟧ , pushint-t thenᵗ is↦*) = ((_ , _ , _) , pushint-eval)
    progress (⟦st⟧ , ⟦F⟧ , pushbool-t thenᵗ is↦*) = ((_ , _ , _) , pushbool-eval)
    progress (⟦⟧ˢ-int (⟦⟧ˢ-int ⟦st⟧) , ⟦F⟧ , add-t thenᵗ is↦*) = ((_ , _ , _) , add-eval)
    progress {i = loop 0} (⟦st⟧ , ⟦F⟧ , (loop-t thenᵗ is↦*)) = let (is' , is⇒ₓis' , _) = lemma₂ is↦* in ((_ , _ , _) , loop-eval⁰ is⇒ₓis')
    progress {i = loop (suc n)} (⟦st⟧ , ⟦F⟧ , _) = ((_ , _ , _) , loop-evalⁿ)
    progress {s = bool true ∷ s} {i = break} {S = _ ∷ S} (⟦st⟧ , ⟦F⟧ , (break-t thenᵗ is↦*)) = let (is' , is⇒ₓis' , _) = lemma₂ is↦* in ((s , is' , S) , break-evalᵗ is⇒ₓis')
    progress {s = bool false ∷ s} {i = break} (⟦st⟧ , ⟦F⟧ , (break-t thenᵗ is↦*)) = ((_ , _ , _) , break-evalᶠ)
    progress {S = is' ∷ S} (⟦st⟧ , ⟦F⟧ , (end-t thenᵗ is↦*)) = ((_ , is' , S) , end-eval)

{-# TERMINATING #-} -- evaluation always terminates, but proving it is tedious
static-dynamic : ∀ stˡ → (⟦_⟧ stˡ σ s is S) → Σ (List Value) (λ sˡ → (s , is , S) ⇒* (sˡ , [] , []) × ⟦ stˡ ⟧ˢ sˡ)
static-dynamic {s = s} stˡ (⟦st⟧ , ⟦⟧ᶠ-empty , doneᵗ) = (s , doneᵉ , ⟦st⟧)
static-dynamic {is = is@(i ∷ is')} stˡ ⟦σ⟧ =
    let ((s' , is' , S') , s,is,S⇒s',is',S') = progress stˡ ⟦σ⟧ in
    let (σ' , ⟦σ'⟧) = preservation stˡ s,is,S⇒s',is',S' ⟦σ⟧ in
    let (sˡ , (s',is'S'⇒*sˡ,[],[] , ⟦stˡ⟧)) = static-dynamic stˡ ⟦σ'⟧ in
    (sˡ , (s,is,S⇒s',is',S' thenᵉ s',is'S'⇒*sˡ,[],[] , ⟦stˡ⟧))

open import Data.Nat.Show using () renaming (show to showℕ)
open import Data.Integer.Show using () renaming (show to showℤ)
open import Data.Bool.Show using () renaming (show to showBool)
open import Data.String using (String; _++_)
open import Data.Product using (proj₁)

stackToString : List Value → String
stackToString [] = "[]"
stackToString (int i ∷ s) = showℤ i ++ " :: " ++ stackToString s
stackToString (bool b ∷ s) = showBool b ++ " :: " ++ stackToString s

open import IO
open import Level using (0ℓ)

main = run {0ℓ} (putStrLn (stackToString (proj₁ (static-dynamic _ (⟦⟧ˢ-empty , ⟦⟧ᶠ-empty , basict₁)))))