Compile.lean

/-
The compiler from a letter-synchronous multi-stack machine to an exact LMR
scaffolding automaton, and its correctness for every input word.

The compiled transition uses only `Unfold.labelOf` at depth 0 and depth 1 of
the tree unfolding, and names its targets by the port words `[]`, `[j]`, and
`[j, j]`.  It never inspects node identity, never compares two port words for
aliasing, never uses `SELF`, and appends exactly one node per input symbol.
-/

import FPRDTransfer.StackEncoding

namespace FPRD

variable {Sym Q A : Type} {s : Nat}

/-- Stack tops, read from the root label of the unfolding. -/
def readTops (n : Unfold s (Tuple s A) 2) (j : Fin s) : Option A :=
  n.labelOf.bind fun g => g j

/-- Second elements, read from the depth-1 label along port `j`.  A missing
edge and an unlabelled node both yield `none`. -/
def readSeconds (n : Unfold s (Tuple s A) 2) (j : Fin s) : Option A :=
  (n.child j).labelOf.bind fun g => g j

/-- The compiled transition function. -/
def compileDelta (M : StackMachine Sym Q A s) (q : Q) (x : Sym)
    (n : Unfold s (Tuple s A) 2) : Q × Tuple s A × (Fin s → Target s) :=
  let r := M.step q x (readTops n)
  (r.1,
    (fun j =>
      match r.2 j with
      | .push a => some a
      | .keep => readTops n j
      | .pop => readSeconds n j),
    fun j =>
      match r.2 j with
      | .push _ => .path []
      | .keep => .path [j]
      | .pop => .path [j, j])

/-- The compiled scaffolding automaton: degree `s`, distance `2`. -/
def compile (M : StackMachine Sym Q A s) : Automaton Sym Q (Tuple s A) s 2 where
  delta := compileDelta M
  init := M.init
  accept := M.accept

theorem compile_distanceOk (M : StackMachine Sym Q A s) : (compile M).DistanceOk := by
  intro q x n j p h
  simp only [compile, compileDelta] at h
  cases hop : (M.step q x (readTops n)).2 j with
  | push a =>
    rw [hop] at h; injection h with h'; subst h'; simp
  | keep =>
    rw [hop] at h; injection h with h'; subst h'; simp
  | pop =>
    rw [hop] at h; injection h with h'; subst h'; simp

theorem compile_selfFree (M : StackMachine Sym Q A s) : (compile M).SelfFree := by
  intro q x n j
  simp only [compile, compileDelta]
  cases (M.step q x (readTops n)).2 j <;> simp

/-! ### Reading the unfolding agrees with the scaffold -/

theorem readTops_eq (S : Scaffold s (Tuple s A)) (v : Nat) (j : Fin s) :
    readTops (S.viewAt 2 v) j = labelComp S v j := rfl

theorem child_labelOf (S : Scaffold s (Tuple s A)) (v : Nat) (j : Fin s) :
    ((S.viewAt 2 v).child j).labelOf = (S.edge v j).bind fun m => S.label m := by
  rw [Scaffold.viewAt_child]
  cases S.edge v j with
  | none => rfl
  | some m => exact S.viewAt_labelOf 1 m

theorem readSeconds_eq (S : Scaffold s (Tuple s A)) (v : Nat) (j : Fin s) :
    readSeconds (S.viewAt 2 v) j = labelCompOpt S (S.edge v j) j := by
  rw [readSeconds, child_labelOf]
  cases S.edge v j with
  | none => rfl
  | some m => rfl

/-! ### Resolving the three port words -/

theorem resolve_nil (S : Scaffold s (Tuple s A)) (t new : Nat) :
    S.resolve t new (.path []) = some t := rfl

theorem resolve_one (S : Scaffold s (Tuple s A)) (t new : Nat) (j : Fin s) :
    S.resolve t new (.path [j]) = S.edge t j :=
  S.follow_single t j

theorem resolve_two (S : Scaffold s (Tuple s A)) (t new : Nat) (j : Fin s) :
    S.resolve t new (.path [j, j]) = (S.edge t j).bind (fun m => S.edge m j) := by
  show S.follow t [j, j] = _
  cases he : S.edge t j with
  | none => rw [S.follow_cons_none t j [j] he]; rfl
  | some m => rw [S.follow_cons_some t j [j] m he, S.follow_single m j]; rfl

/-! ### Decoding is unchanged below the new node -/

theorem decode_append_below (S : Scaffold s (Tuple s A)) (g : Tuple s A)
    (tgt : Fin s → Target s) (j : Fin s) (u : Nat) (hu : u ≤ S.top) :
    decode (S.append g tgt) j u = decode S j u := by
  apply decode_congr
  · intro w hw
    exact S.append_label_old g tgt (by omega)
  · intro w hw
    exact S.append_edge_old g tgt (by omega) j

theorem decodeOpt_append_below (S : Scaffold s (Tuple s A)) (g : Tuple s A)
    (tgt : Fin s → Target s) (j : Fin s) (o : Option Nat)
    (ho : ∀ m, o = some m → m ≤ S.top) :
    decodeOpt (S.append g tgt) j o = decodeOpt S j o := by
  cases ho' : o with
  | none => rfl
  | some m =>
    simp only [decodeOpt_some]
    exact decode_append_below S g tgt j m (ho m ho')

/-! ### The simulation invariant -/

/-- The invariant relating a machine configuration to a compiled scaffold.

`clean` is the invariant the informal argument omits; it is what makes a pop of
an empty stack decode to the empty stack rather than resurrecting a stale
symbol reachable through a leftover port. -/
structure Sim (c : Config Q A s) (q : Q) (S : Scaffold s (Tuple s A)) : Prop where
  ctrl : c.ctrl = q
  stacks : ∀ j, decode S j S.top = c.stacks j
  backward : S.Backward
  clean : CleanEmpty S
  aboveLabel : ∀ v, S.top < v → S.label v = none
  aboveEdge : ∀ v, S.top < v → ∀ i, S.edge v i = none

theorem sim_init (M : StackMachine Sym Q A s) :
    Sim M.initConfig (compile M).init (Scaffold.base s (Tuple s A)) where
  ctrl := rfl
  stacks := fun j => by
    simp only [StackMachine.initConfig]
    exact base_decode j _
  backward := Scaffold.base_backward s (Tuple s A)
  clean := base_cleanEmpty
  aboveLabel := fun _ _ => rfl
  aboveEdge := fun _ _ _ => rfl

/-- The tops the compiled transition reads are exactly the machine's tops. -/
theorem sim_tops {c : Config Q A s} {q : Q} {S : Scaffold s (Tuple s A)}
    (h : Sim c q S) : readTops (S.viewAt 2 S.top) = fun j => (c.stacks j).head? := by
  funext j
  rw [readTops_eq, ← decode_head, h.stacks j]

/-! ### The new node, in terms of the scaffold rather than the unfolding

`newLabel` and `newTarget` restate the compiled transition's output using
`labelComp` and `Scaffold.edge`.  `compileDelta_label` and
`compileDelta_target` show they agree with what the transition actually emits
from the unfolding, so the rest of the proof never mentions `Unfold` again. -/

/-- The label the compiler puts on the new node. -/
def newLabel (S : Scaffold s (Tuple s A)) (ops : Fin s → StackOp A) : Tuple s A :=
  fun j =>
    match ops j with
    | .push a => some a
    | .keep => labelComp S S.top j
    | .pop => labelCompOpt S (S.edge S.top j) j

/-- The port words the compiler names for the new node. -/
def newTarget (ops : Fin s → StackOp A) : Fin s → Target s :=
  fun j =>
    match ops j with
    | .push _ => .path []
    | .keep => .path [j]
    | .pop => .path [j, j]

theorem newTarget_ne_self (ops : Fin s → StackOp A) (j : Fin s) :
    newTarget ops j ≠ .self := by
  simp only [newTarget]
  cases ops j <;> simp

theorem compileDelta_label (M : StackMachine Sym Q A s) (q : Q) (x : Sym)
    (S : Scaffold s (Tuple s A)) :
    (compileDelta M q x (S.viewAt 2 S.top)).2.1 =
      newLabel S (M.step q x (readTops (S.viewAt 2 S.top))).2 := by
  funext j
  simp only [compileDelta, newLabel]
  cases (M.step q x (readTops (S.viewAt 2 S.top))).2 j with
  | push a => rfl
  | keep => exact readTops_eq S S.top j
  | pop => exact readSeconds_eq S S.top j

theorem compileDelta_target (M : StackMachine Sym Q A s) (q : Q) (x : Sym)
    (S : Scaffold s (Tuple s A)) :
    (compileDelta M q x (S.viewAt 2 S.top)).2.2 =
      newTarget (M.step q x (readTops (S.viewAt 2 S.top))).2 := by
  funext j
  simp only [compileDelta, newTarget]

/-! ### The new node decodes to the machine's action -/

/-- The heart of the bridge: after appending the compiler's node, stack `j`
decodes to exactly the result of stack `j`'s action.

The `pop`-of-an-empty-stack case is where `CleanEmpty` is consumed. -/
theorem decode_new (S : Scaffold s (Tuple s A)) (hb : S.Backward) (hc : CleanEmpty S)
    (stacks : Fin s → List A) (hst : ∀ j, decode S j S.top = stacks j)
    (ops : Fin s → StackOp A) (j : Fin s) :
    decode (S.append (newLabel S ops) (newTarget ops)) j (S.top + 1)
      = (ops j).apply (stacks j) := by
  have hb' : (S.append (newLabel S ops) (newTarget ops)).Backward :=
    S.append_backward hb _ _ (newTarget_ne_self ops)
  have hlab : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
      = newLabel S ops j := by
    simp only [labelComp, Scaffold.append_label_new]
    rfl
  have hedge : (S.append (newLabel S ops) (newTarget ops)).edge (S.top + 1) j
      = S.resolve S.top (S.top + 1) (newTarget ops j) := by
    rw [Scaffold.append_edge_new]
  have hbelow : ∀ o : Option Nat, (∀ m, o = some m → m ≤ S.top) →
      decodeOpt (S.append (newLabel S ops) (newTarget ops)) j o = decodeOpt S j o :=
    fun o ho => decodeOpt_append_below S _ _ j o ho
  -- edges out of the old top land strictly below it
  have htop_lt : ∀ m, S.edge S.top j = some m → m < S.top := fun m hm => hb _ _ _ hm
  cases hop : ops j with
  | push a =>
    have hl : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
        = some a := by rw [hlab]; simp only [newLabel, hop]
    have ht : (S.append (newLabel S ops) (newTarget ops)).edge (S.top + 1) j
        = some S.top := by rw [hedge]; simp only [newTarget, hop]; rfl
    rw [decode_cons _ hb' j (S.top + 1) a hl, ht, decodeOpt_some,
      decode_append_below S _ _ j S.top (Nat.le_refl _), hst j]
    rfl
  | keep =>
    have ht : (S.append (newLabel S ops) (newTarget ops)).edge (S.top + 1) j
        = S.edge S.top j := by
      rw [hedge]; simp only [newTarget, hop]; exact resolve_one S S.top (S.top + 1) j
    cases hc0 : labelComp S S.top j with
    | none =>
      have hl : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
          = none := by rw [hlab]; simp only [newLabel, hop]; exact hc0
      have hzero : stacks j = [] := by rw [← hst j]; exact decode_empty S j S.top hc0
      rw [decode_empty _ j _ hl, hzero]
      rfl
    | some a =>
      have hl : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
          = some a := by rw [hlab]; simp only [newLabel, hop]; exact hc0
      rw [decode_cons _ hb' j (S.top + 1) a hl, ht,
        hbelow (S.edge S.top j) (fun m hm => Nat.le_of_lt (htop_lt m hm)),
        ← decode_cons S hb j S.top a hc0, hst j]
      rfl
  | pop =>
    have ht : (S.append (newLabel S ops) (newTarget ops)).edge (S.top + 1) j
        = (S.edge S.top j).bind (fun m => S.edge m j) := by
      rw [hedge]; simp only [newTarget, hop]; exact resolve_two S S.top (S.top + 1) j
    cases hc0 : labelComp S S.top j with
    | none =>
      -- empty stack: CleanEmpty forces the new label component to be none
      have hclean : labelCompOpt S (S.edge S.top j) j = none := hc S.top j hc0
      have hl : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
          = none := by rw [hlab]; simp only [newLabel, hop]; exact hclean
      have hempty : stacks j = [] := by rw [← hst j]; exact decode_empty S j S.top hc0
      rw [decode_empty _ j _ hl, hempty]
      rfl
    | some a =>
      have hsplit : stacks j = a :: decodeOpt S j (S.edge S.top j) := by
        rw [← hst j]; exact decode_cons S hb j S.top a hc0
      cases he : S.edge S.top j with
      | none =>
        have hl : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
            = none := by rw [hlab]; simp only [newLabel, hop, he]; rfl
        rw [decode_empty _ j _ hl, hsplit, he]
        rfl
      | some m =>
        have hmlt : m < S.top := htop_lt m he
        cases hcm : labelComp S m j with
        | none =>
          have hl : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
              = none := by rw [hlab]; simp only [newLabel, hop, he]; exact hcm
          rw [decode_empty _ j _ hl, hsplit, he, decodeOpt_some,
            decode_empty S j m hcm]
          rfl
        | some b =>
          have hl : labelComp (S.append (newLabel S ops) (newTarget ops)) (S.top + 1) j
              = some b := by rw [hlab]; simp only [newLabel, hop, he]; exact hcm
          have ht2 : (S.append (newLabel S ops) (newTarget ops)).edge (S.top + 1) j
              = S.edge m j := by rw [ht, he]; rfl
          have hbound : ∀ m', S.edge m j = some m' → m' ≤ S.top := by
            intro m' hm'
            have : m' < m := hb _ _ _ hm'
            omega
          rw [decode_cons _ hb' j (S.top + 1) b hl, ht2, hbelow (S.edge m j) hbound,
            hsplit, he, decodeOpt_some, decode_cons S hb j m b hcm]
          rfl


/-! ### Auxiliary rewriting of labels below the new node -/

theorem labelComp_append_old (S : Scaffold s (Tuple s A)) (g : Tuple s A)
    (tgt : Fin s → Target s) {v : Nat} (hv : v ≠ S.top + 1) (j : Fin s) :
    labelComp (S.append g tgt) v j = labelComp S v j := by
  simp only [labelComp, S.append_label_old g tgt hv]

theorem labelCompOpt_append_below (S : Scaffold s (Tuple s A)) (g : Tuple s A)
    (tgt : Fin s → Target s) (j : Fin s) (o : Option Nat)
    (ho : ∀ m, o = some m → m ≤ S.top) :
    labelCompOpt (S.append g tgt) o j = labelCompOpt S o j := by
  cases ho' : o with
  | none => rfl
  | some m =>
    simp only [labelCompOpt_some]
    exact labelComp_append_old S g tgt (by have := ho m ho'; omega) j

theorem newTarget_isPath (ops : Fin s → StackOp A) (j : Fin s) :
    ∃ p, newTarget ops j = Target.path p := by
  simp only [newTarget]
  cases ops j
  · exact ⟨[], rfl⟩
  · exact ⟨[j], rfl⟩
  · exact ⟨[j, j], rfl⟩

/-- Every target the compiler names lands at or below the old top.  This is
`Scaffold.follow_le` applied to the three port words. -/
theorem resolve_newTarget_le (S : Scaffold s (Tuple s A)) (hb : S.Backward)
    (ops : Fin s → StackOp A) (j : Fin s) (m : Nat)
    (h : S.resolve S.top (S.top + 1) (newTarget ops j) = some m) : m ≤ S.top := by
  obtain ⟨p, hp⟩ := newTarget_isPath ops j
  rw [hp] at h
  exact S.follow_le hb p S.top m h

/-! ### `CleanEmpty` is preserved -/

theorem cleanEmpty_append (S : Scaffold s (Tuple s A)) (hb : S.Backward)
    (hc : CleanEmpty S) (hae : ∀ v, S.top < v → ∀ i, S.edge v i = none)
    (ops : Fin s → StackOp A) :
    CleanEmpty (S.append (newLabel S ops) (newTarget ops)) := by
  intro v j hv
  by_cases hvt : v = S.top + 1
  · -- the new node
    subst hvt
    have hedge : (S.append (newLabel S ops) (newTarget ops)).edge (S.top + 1) j
        = S.resolve S.top (S.top + 1) (newTarget ops j) := by rw [Scaffold.append_edge_new]
    have hlow : labelCompOpt (S.append (newLabel S ops) (newTarget ops))
        (S.resolve S.top (S.top + 1) (newTarget ops j)) j
        = labelCompOpt S (S.resolve S.top (S.top + 1) (newTarget ops j)) j :=
      labelCompOpt_append_below S _ _ j _ (resolve_newTarget_le S hb ops j)
    rw [hedge, hlow]
    have hnl : newLabel S ops j = none := by
      rw [← hv]
      simp only [labelComp, Scaffold.append_label_new]
      rfl
    simp only [newLabel] at hnl
    cases hop : ops j with
    | push a =>
      rw [hop] at hnl
      exact absurd hnl (by simp)
    | keep =>
      rw [hop] at hnl
      have := hc S.top j hnl
      simp only [newTarget, hop]
      rw [resolve_one S S.top (S.top + 1) j]
      exact this
    | pop =>
      rw [hop] at hnl
      simp only [newTarget, hop]
      rw [resolve_two S S.top (S.top + 1) j]
      cases he : S.edge S.top j with
      | none => rfl
      | some m =>
        rw [he] at hnl
        simp only [labelCompOpt_some] at hnl
        exact hc m j hnl
  · -- an old node, or a node above the new top
    rw [S.append_edge_old (newLabel S ops) (newTarget ops) hvt j]
    rw [labelComp_append_old S _ _ hvt j] at hv
    by_cases hle : v ≤ S.top
    · have hb' : ∀ m, S.edge v j = some m → m ≤ S.top := by
        intro m hm
        have := hb v j m hm
        omega
      rw [labelCompOpt_append_below S _ _ j _ hb']
      exact hc v j hv
    · have : S.edge v j = none := hae v (by omega) j
      rw [this]
      rfl

/-! ### One step of the simulation -/

theorem compile_step_eq (M : StackMachine Sym Q A s) (q : Q)
    (S : Scaffold s (Tuple s A)) (x : Sym) :
    (compile M).step (q, S) x =
      ((M.step q x (readTops (S.viewAt 2 S.top))).1,
        S.append (newLabel S (M.step q x (readTops (S.viewAt 2 S.top))).2)
          (newTarget (M.step q x (readTops (S.viewAt 2 S.top))).2)) := by
  show ((compileDelta M q x (S.viewAt 2 S.top)).1,
      S.append (compileDelta M q x (S.viewAt 2 S.top)).2.1
        (compileDelta M q x (S.viewAt 2 S.top)).2.2) = _
  rw [compileDelta_label, compileDelta_target]
  rfl

theorem sim_step (M : StackMachine Sym Q A s) {c : Config Q A s}
    {p : Q × Scaffold s (Tuple s A)} (h : Sim c p.1 p.2) (x : Sym) :
    Sim (M.stepConfig c x) ((compile M).step p x).1 ((compile M).step p x).2 := by
  obtain ⟨q, S⟩ := p
  have htops : readTops (S.viewAt 2 S.top) = fun j => (c.stacks j).head? := sim_tops h
  have hstep : M.step c.ctrl x (fun j => (c.stacks j).head?)
      = M.step q x (readTops (S.viewAt 2 S.top)) := by
    rw [h.ctrl, htops]
  rw [compile_step_eq M q S x]
  refine
    { ctrl := ?_, stacks := ?_, backward := ?_, clean := ?_
      aboveLabel := ?_, aboveEdge := ?_ }
  · show (M.step c.ctrl x (fun j => (c.stacks j).head?)).1 = _
    rw [hstep]
  · intro j
    show decode _ j (S.top + 1) = ((M.step c.ctrl x (fun j => (c.stacks j).head?)).2 j).apply _
    rw [hstep]
    exact decode_new S h.backward h.clean c.stacks h.stacks _ j
  · exact S.append_backward h.backward _ _ (newTarget_ne_self _)
  · exact cleanEmpty_append S h.backward h.clean h.aboveEdge _
  · intro v hvv
    have hv' : S.top + 1 < v := hvv
    rw [S.append_label_old _ _ (by omega : v ≠ S.top + 1)]
    exact h.aboveLabel v (by show S.top < v; omega)
  · intro v hvv i
    have hv' : S.top + 1 < v := hvv
    rw [S.append_edge_old _ _ (by omega : v ≠ S.top + 1) i]
    exact h.aboveEdge v (by show S.top < v; omega) i

/-! ### The theorem: simulation holds after every prefix -/

theorem sim_run (M : StackMachine Sym Q A s) :
    ∀ (w : List Sym) (c : Config Q A s) (p : Q × Scaffold s (Tuple s A)),
      Sim c p.1 p.2 →
      Sim (w.foldl M.stepConfig c) (w.foldl (compile M).step p).1
        (w.foldl (compile M).step p).2 := by
  intro w
  induction w with
  | nil => intro c p h; exact h
  | cons x w ih =>
    intro c p h
    simp only [List.foldl_cons]
    exact ih (M.stepConfig c x) ((compile M).step p x) (sim_step M h x)

/-- **Stack compiler correctness.**  For every input word, the compiled LMR
scaffolding automaton and the source machine have the same control state, and
every logical stack decoded from the scaffold's top equals the source stack. -/
theorem compile_correct (M : StackMachine Sym Q A s) (w : List Sym) :
    Sim (M.run w) ((compile M).run w).1 ((compile M).run w).2 :=
  sim_run M w M.initConfig (M.init, Scaffold.base s (Tuple s A)) (sim_init M)

theorem compile_ctrl (M : StackMachine Sym Q A s) (w : List Sym) :
    (M.run w).ctrl = ((compile M).run w).1 :=
  (compile_correct M w).ctrl

theorem compile_stacks (M : StackMachine Sym Q A s) (w : List Sym) (j : Fin s) :
    decode ((compile M).run w).2 j ((compile M).run w).2.top = (M.run w).stacks j :=
  (compile_correct M w).stacks j

/-- The compiled automaton decides exactly the source machine's language. -/
theorem compile_decides (M : StackMachine Sym Q A s) (w : List Sym) :
    (compile M).decides w = M.decides w := by
  show (compile M).accept ((compile M).run w).1 = M.accept (M.run w).ctrl
  rw [compile_ctrl M w]
  rfl

end FPRD