StackEncoding.lean

/-
The persistent-stack encoding inside a scaffold.

A node's label is a tuple with one component per logical stack; component `j`
holds the top symbol of stack `j` at that node, or `none` to mean "stack `j`
is empty here".  Port `j` holds the predecessor for stack `j`.  Decoding
stack `j` from a node walks port `j`, reading component `j`, and stops at the
first `none`.

The `CleanEmpty` invariant below is the load-bearing discovery of this
formalization.  Nothing in the informal argument states it, but without it the
encoding is *wrong* for a pop of an empty stack: an empty stack's port may
still point at an older node whose `j`-component is labelled, and the
prescribed pop would then read that stale symbol and resurrect it.
`CleanEmpty` says this never happens, and `Compile` proves the compiler
maintains it.
-/

import FPRDTransfer.LMR
import FPRDTransfer.StackMachine

namespace FPRD

/-- The working alphabet: one component per logical stack. -/
abbrev Tuple (s : Nat) (A : Type) := Fin s → Option A

variable {s : Nat} {A : Type}

/-- Component `j` of the label at node `v`; `none` when `v` is unlabelled.

Written with `Option.bind` rather than a `match` so that it and the reader in
`Compile` reduce to syntactically identical terms. -/
def labelComp (S : Scaffold s (Tuple s A)) (v : Nat) (j : Fin s) : Option A :=
  (S.label v).bind fun g => g j

/-- `labelComp` extended to a possibly missing node. -/
def labelCompOpt (S : Scaffold s (Tuple s A)) (v : Option Nat) (j : Fin s) :
    Option A :=
  v.bind fun u => labelComp S u j

@[simp] theorem labelCompOpt_none (S : Scaffold s (Tuple s A)) (j : Fin s) :
    labelCompOpt S none j = none := rfl

@[simp] theorem labelCompOpt_some (S : Scaffold s (Tuple s A)) (u : Nat) (j : Fin s) :
    labelCompOpt S (some u) j = labelComp S u j := rfl

/-- Decode stack `j` from node `v`, top symbol first.

The `m < v` guard keeps this total on every scaffold, including one that used
`SELF`; a self-loop is not a predecessor, so decoding stops there.  For the
scaffolds built by the compiler the guard always holds, by `Backward`. -/
def decode (S : Scaffold s (Tuple s A)) (j : Fin s) (v : Nat) : List A :=
  match labelComp S v j with
  | none => []
  | some a =>
    match S.edge v j with
    | none => [a]
    | some m => if _h : m < v then a :: decode S j m else [a]
termination_by v
decreasing_by exact _h

/-- Decode from a possibly missing node. -/
def decodeOpt (S : Scaffold s (Tuple s A)) (j : Fin s) : Option Nat → List A
  | none => []
  | some m => decode S j m

@[simp] theorem decodeOpt_none (S : Scaffold s (Tuple s A)) (j : Fin s) :
    decodeOpt S j none = [] := rfl

@[simp] theorem decodeOpt_some (S : Scaffold s (Tuple s A)) (j : Fin s) (m : Nat) :
    decodeOpt S j (some m) = decode S j m := rfl

theorem decode_empty (S : Scaffold s (Tuple s A)) (j : Fin s) (v : Nat)
    (h : labelComp S v j = none) : decode S j v = [] := by
  rw [decode, h]

theorem decode_edge_none (S : Scaffold s (Tuple s A)) (j : Fin s) (v : Nat) (a : A)
    (h : labelComp S v j = some a) (he : S.edge v j = none) :
    decode S j v = [a] := by
  rw [decode, h, he]

theorem decode_edge_some (S : Scaffold s (Tuple s A)) (j : Fin s) (v m : Nat) (a : A)
    (h : labelComp S v j = some a) (he : S.edge v j = some m) (hlt : m < v) :
    decode S j v = a :: decode S j m := by
  rw [decode, h, he]
  simp [hlt]

/-- The head of a decoded stack is exactly the label component: the machine's
view of the stack tops is faithful. -/
theorem decode_head (S : Scaffold s (Tuple s A)) (j : Fin s) (v : Nat) :
    (decode S j v).head? = labelComp S v j := by
  cases h : labelComp S v j with
  | none => rw [decode_empty S j v h]; rfl
  | some a =>
    cases he : S.edge v j with
    | none => rw [decode_edge_none S j v a h he]; rfl
    | some m =>
      by_cases hlt : m < v
      · rw [decode_edge_some S j v m a h he hlt]; rfl
      · rw [decode, h, he]; simp [hlt]

/-- Under `Backward`, a nonempty decoded stack splits into its top symbol and
the decode of its port target. -/
theorem decode_cons (S : Scaffold s (Tuple s A)) (hb : S.Backward) (j : Fin s)
    (v : Nat) (a : A) (h : labelComp S v j = some a) :
    decode S j v = a :: decodeOpt S j (S.edge v j) := by
  cases he : S.edge v j with
  | none => rw [decode_edge_none S j v a h he]; simp
  | some m =>
    have hlt : m < v := hb v j m he
    rw [decode_edge_some S j v m a h he hlt]
    simp

/-! ### A structurally recursive twin, for kernel computation

`decode` is defined by well-founded recursion, which the kernel does not
reduce, so concrete checks cannot be closed by `decide` through it.
`decodeAux` is the same function with an explicit fuel parameter, hence
structurally recursive and reducible.  `decode_eq_decodeAux` licenses using it:
the examples state their claims about `decode` and compute with `decodeAux`. -/

def decodeAux (S : Scaffold s (Tuple s A)) (j : Fin s) : Nat → Nat → List A
  | 0, _ => []
  | fuel + 1, v =>
    match labelComp S v j with
    | none => []
    | some a =>
      match S.edge v j with
      | none => [a]
      | some m => if m < v then a :: decodeAux S j fuel m else [a]

theorem decodeAux_eq (S : Scaffold s (Tuple s A)) (j : Fin s) :
    ∀ (fuel v : Nat), v < fuel → decodeAux S j fuel v = decode S j v := by
  intro fuel
  induction fuel with
  | zero => intro v hv; exact absurd hv (by omega)
  | succ f ih =>
    intro v hv
    cases h : labelComp S v j with
    | none =>
      rw [decode_empty S j v h]
      simp only [decodeAux, h]
    | some a =>
      cases he : S.edge v j with
      | none =>
        rw [decode_edge_none S j v a h he]
        simp only [decodeAux, h, he]
      | some m =>
        by_cases hlt : m < v
        · rw [decode_edge_some S j v m a h he hlt]
          simp only [decodeAux, h, he, if_pos hlt]
          rw [ih m (by omega)]
        · rw [decode, h, he]
          simp only [decodeAux, h, he, hlt, if_false]
          simp

theorem decode_eq_decodeAux (S : Scaffold s (Tuple s A)) (j : Fin s) (v : Nat) :
    decode S j v = decodeAux S j (v + 1) v :=
  (decodeAux_eq S j (v + 1) v (Nat.lt_succ_self v)).symm

/-- Empty stacks have clean predecessor chains: if node `v` ends stack `j`,
then its port-`j` target also ends stack `j`. -/
def CleanEmpty (S : Scaffold s (Tuple s A)) : Prop :=
  ∀ v j, labelComp S v j = none → labelCompOpt S (S.edge v j) j = none

/-- The consequence that makes an empty pop correct. -/
theorem decodeOpt_of_empty (S : Scaffold s (Tuple s A)) (hc : CleanEmpty S)
    (v : Nat) (j : Fin s) (h : labelComp S v j = none) :
    decodeOpt S j (S.edge v j) = [] := by
  have hcv := hc v j h
  cases he : S.edge v j with
  | none => rfl
  | some m =>
    rw [he] at hcv
    simp only [labelCompOpt_some] at hcv
    simp only [decodeOpt_some]
    exact decode_empty S j m hcv

/-- Decoding only inspects nodes at most `n`, so two scaffolds agreeing up to
`n` decode alike from `n`.  Proved with an explicit bound so that no
strong-induction facility is needed. -/
theorem decode_congr_bounded (S S' : Scaffold s (Tuple s A)) (j : Fin s) :
    ∀ (bound n : Nat), n ≤ bound →
      (∀ u, u ≤ n → S'.label u = S.label u) →
      (∀ u, u ≤ n → S'.edge u j = S.edge u j) →
      decode S' j n = decode S j n := by
  intro bound
  induction bound with
  | zero =>
    intro n hn hlab hedge
    have hn0 : n = 0 := Nat.le_zero.mp hn
    subst hn0
    have hl : labelComp S' 0 j = labelComp S 0 j := by
      simp only [labelComp, hlab 0 (Nat.le_refl 0)]
    cases h : labelComp S 0 j with
    | none =>
      rw [decode_empty S' j 0 (by rw [hl, h]), decode_empty S j 0 h]
    | some a =>
      have he : S'.edge 0 j = S.edge 0 j := hedge 0 (Nat.le_refl 0)
      cases hE : S.edge 0 j with
      | none =>
        rw [decode_edge_none S' j 0 a (by rw [hl, h]) (by rw [he, hE]),
            decode_edge_none S j 0 a h hE]
      | some m =>
        rw [decode, hl, h, he, hE, decode, h, hE]
        simp
  | succ b ih =>
    intro n hn hlab hedge
    have hl : labelComp S' n j = labelComp S n j := by
      simp only [labelComp, hlab n (Nat.le_refl n)]
    cases h : labelComp S n j with
    | none =>
      rw [decode_empty S' j n (by rw [hl, h]), decode_empty S j n h]
    | some a =>
      have he : S'.edge n j = S.edge n j := hedge n (Nat.le_refl n)
      cases hE : S.edge n j with
      | none =>
        rw [decode_edge_none S' j n a (by rw [hl, h]) (by rw [he, hE]),
            decode_edge_none S j n a h hE]
      | some m =>
        by_cases hlt : m < n
        · have hmb : m ≤ b := by omega
          have hrec : decode S' j m = decode S j m :=
            ih m hmb (fun u hu => hlab u (by omega)) (fun u hu => hedge u (by omega))
          rw [decode_edge_some S' j n m a (by rw [hl, h]) (by rw [he, hE]) hlt,
              decode_edge_some S j n m a h hE hlt, hrec]
        · rw [decode, hl, h, he, hE, decode, h, hE]
          simp [hlt]

theorem decode_congr (S S' : Scaffold s (Tuple s A)) (j : Fin s) (n : Nat)
    (hlab : ∀ u, u ≤ n → S'.label u = S.label u)
    (hedge : ∀ u, u ≤ n → S'.edge u j = S.edge u j) :
    decode S' j n = decode S j n :=
  decode_congr_bounded S S' j n n (Nat.le_refl n) hlab hedge

/-- The base scaffold decodes to empty stacks and is clean. -/
theorem base_labelComp (j : Fin s) (v : Nat) :
    labelComp (Scaffold.base s (Tuple s A)) v j = none := rfl

theorem base_decode (j : Fin s) (v : Nat) :
    decode (Scaffold.base s (Tuple s A)) j v = [] :=
  decode_empty _ j v (base_labelComp j v)

theorem base_cleanEmpty : CleanEmpty (Scaffold.base s (Tuple s A)) := by
  intro v j _
  rfl

end FPRD