RootSharingMachine.lean

import FrontierDerivative.RootSharing

/-!
# Source-selecting stack machines

This module turns the coordinate-carrying one-step theorem into an exact
compiler. A transition reads every current top and lets each output slot
derive by push, keep, or pop from any source slot of the preceding round.

The source model is an interface for a stationary frontier allocator. This
module does not construct such an allocator from a bounded frontier.
-/

namespace FrontierDerivative.RootSharingMachine

open FPRD
open RootSharing

universe u

/-- A deterministic letter-synchronous fixed-slot machine with source-selecting
stack actions. -/
structure Machine (Sym : Type) (Q : Type) (A : Type u) (s : Nat) where
  step : Q → Sym → (Fin s → Option A) → Q × (Fin s → Derive s A)
  init : Q
  accept : Q → Bool

/-- Logical source configuration, with stacks held top-first. -/
structure Config (Q : Type) (A : Type u) (s : Nat) where
  ctrl : Q
  stacks : Fin s → List A

namespace Machine

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

def stepConfig (M : Machine Sym Q A s) (c : Config Q A s) (x : Sym) :
    Config Q A s :=
  let r := M.step c.ctrl x (fun j => (c.stacks j).head?)
  { ctrl := r.1, stacks := fun j => (r.2 j).apply c.stacks }

@[simp] theorem stepConfig_ctrl (M : Machine Sym Q A s) (c : Config Q A s)
    (x : Sym) :
    (M.stepConfig c x).ctrl =
      (M.step c.ctrl x (fun j => (c.stacks j).head?)).1 := rfl

@[simp] theorem stepConfig_stacks (M : Machine Sym Q A s) (c : Config Q A s)
    (x : Sym) (j : Fin s) :
    (M.stepConfig c x).stacks j =
      ((M.step c.ctrl x (fun j => (c.stacks j).head?)).2 j).apply c.stacks := rfl

def initConfig (M : Machine Sym Q A s) : Config Q A s where
  ctrl := M.init
  stacks := fun _ => []

def run (M : Machine Sym Q A s) (w : List Sym) : Config Q A s :=
  w.foldl M.stepConfig M.initConfig

def decides (M : Machine Sym Q A s) (w : List Sym) : Bool :=
  M.accept (M.run w).ctrl

@[simp] theorem run_nil (M : Machine Sym Q A s) : M.run [] = M.initConfig := rfl

theorem run_append (M : Machine Sym Q A s) (w : List Sym) (x : Sym) :
    M.run (w ++ [x]) = M.stepConfig (M.run w) x := by
  simp [run, List.foldl_append]

end Machine

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

/-! ## Depth-two observations -/

def readCells (n : Unfold s (RootLabel s A) 2) (j : Fin s) : Option (Cell s A) :=
  n.labelOf.bind fun g => g j

def readTops (n : Unfold s (RootLabel s A) 2) (j : Fin s) : Option A :=
  (readCells n j).map Cell.symbol

def readTailCell (n : Unfold s (RootLabel s A) 2) (i : Fin s) :
    Option (Cell s A) :=
  match readCells n i with
  | none => none
  | some c => (n.child i).labelOf.bind fun g => g c.next

def viewLabel (n : Unfold s (RootLabel s A) 2) (ops : Fin s → Derive s A) :
    RootLabel s A :=
  fun j =>
    match ops j with
    | .push i a => some ⟨a, i⟩
    | .keep i => readCells n i
    | .pop i => readTailCell n i

def viewTarget (n : Unfold s (RootLabel s A) 2) (ops : Fin s → Derive s A) :
    Fin s → Target s :=
  fun j =>
    match ops j with
    | .push _ _ => .path []
    | .keep i =>
      match readCells n i with
      | none => .missing
      | some _ => .path [i]
    | .pop i =>
      match readCells n i with
      | none => .missing
      | some c =>
        match readTailCell n i with
        | none => .missing
        | some _ => .path [i, c.next]

/-! ## Compiler -/

def compileDelta (M : Machine Sym Q A s) (q : Q) (x : Sym)
    (n : Unfold s (RootLabel s A) 2) :
    Q × RootLabel s A × (Fin s → Target s) :=
  let r := M.step q x (readTops n)
  (r.1, viewLabel n r.2, viewTarget n r.2)

def compile (M : Machine Sym Q A s) :
    Automaton Sym Q (RootLabel s A) s 2 where
  delta := compileDelta M
  init := M.init
  accept := M.accept

theorem viewTarget_distance (n : Unfold s (RootLabel s A) 2)
    (ops : Fin s → Derive s A) (j : Fin s) (p : List (Fin s))
    (h : viewTarget n ops j = .path p) : p.length ≤ 2 := by
  cases hop : ops j with
  | push i a =>
      simp [viewTarget, hop] at h
      subst p
      simp
  | keep i =>
      cases hc : readCells n i with
      | none => simp [viewTarget, hop, hc] at h
      | some c =>
          simp [viewTarget, hop, hc] at h
          subst p
          simp
  | pop i =>
      cases hc : readCells n i with
      | none => simp [viewTarget, hop, hc] at h
      | some c =>
          cases ht : readTailCell n i with
          | none => simp [viewTarget, hop, hc, ht] at h
          | some d =>
              simp [viewTarget, hop, hc, ht] at h
              subst p
              simp

theorem viewTarget_ne_self (n : Unfold s (RootLabel s A) 2)
    (ops : Fin s → Derive s A) (j : Fin s) : viewTarget n ops j ≠ .self := by
  cases hop : ops j with
  | push i a => simp [viewTarget, hop]
  | keep i =>
      cases hc : readCells n i <;> simp [viewTarget, hop, hc]
  | pop i =>
      cases hc : readCells n i with
      | none => simp [viewTarget, hop, hc]
      | some c =>
          cases ht : readTailCell n i <;> simp [viewTarget, hop, hc, ht]

theorem compile_distanceOk (M : Machine Sym Q A s) : (compile M).DistanceOk := by
  intro q x n j p h
  exact viewTarget_distance n (M.step q x (readTops n)).2 j p h

theorem compile_selfFree (M : Machine Sym Q A s) : (compile M).SelfFree := by
  intro q x n j
  exact viewTarget_ne_self n (M.step q x (readTops n)).2 j

/-! ## Observations agree with concrete scaffolds -/

theorem readCells_eq (S : Scaffold s (RootLabel s A)) (v : Nat) (j : Fin s) :
    readCells (S.viewAt 2 v) j = cellComp S v j := rfl

theorem readTops_eq (S : Scaffold s (RootLabel s A)) (v : Nat) (j : Fin s) :
    readTops (S.viewAt 2 v) j = (cellComp S v j).map Cell.symbol := rfl

theorem child_labelOf (S : Scaffold s (RootLabel 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 readTailCell_eq (S : Scaffold s (RootLabel s A)) (v : Nat) (i : Fin s) :
    readTailCell (S.viewAt 2 v) i =
      match cellComp S v i with
      | none => none
      | some c => cellCompOpt S (S.edge v i) c.next := by
  simp only [readTailCell, readCells_eq]
  cases hc : cellComp S v i with
  | none => rfl
  | some c =>
      rw [child_labelOf]
      cases he : S.edge v i with
      | none => rfl
      | some m => rfl

theorem viewLabel_eq (S : Scaffold s (RootLabel s A))
    (ops : Fin s → Derive s A) :
    viewLabel (S.viewAt 2 S.top) ops = RootSharing.newLabel S ops := by
  funext j
  cases hop : ops j with
  | push i a => simp [viewLabel, RootSharing.newLabel, hop]
  | keep i => simp [viewLabel, RootSharing.newLabel, hop, readCells_eq]
  | pop i =>
      cases hc : cellComp S S.top i with
      | none =>
          simp [viewLabel, RootSharing.newLabel, RootSharing.tailCell, hop, hc,
            readTailCell_eq]
      | some c =>
          simp [viewLabel, RootSharing.newLabel, RootSharing.tailCell, hop, hc,
            readTailCell_eq]

theorem viewTarget_eq (S : Scaffold s (RootLabel s A))
    (ops : Fin s → Derive s A) :
    viewTarget (S.viewAt 2 S.top) ops = RootSharing.newTarget S ops := by
  funext j
  cases hop : ops j with
  | push i a => simp [viewTarget, RootSharing.newTarget, hop]
  | keep i =>
      cases hc : cellComp S S.top i <;>
        simp [viewTarget, RootSharing.newTarget, hop, hc, readCells_eq]
  | pop i =>
      cases hc : cellComp S S.top i with
      | none =>
          simp [viewTarget, RootSharing.newTarget, hop, hc, readCells_eq]
      | some c =>
          cases ht : cellCompOpt S (S.edge S.top i) c.next <;>
            simp [viewTarget, RootSharing.newTarget, RootSharing.tailCell, hop, hc,
              ht, readCells_eq, readTailCell_eq]

theorem compileDelta_label (M : Machine Sym Q A s) (q : Q) (x : Sym)
    (S : Scaffold s (RootLabel s A)) :
    (compileDelta M q x (S.viewAt 2 S.top)).2.1 =
      RootSharing.newLabel S (M.step q x (readTops (S.viewAt 2 S.top))).2 := by
  exact viewLabel_eq S _

theorem compileDelta_target (M : Machine Sym Q A s) (q : Q) (x : Sym)
    (S : Scaffold s (RootLabel s A)) :
    (compileDelta M q x (S.viewAt 2 S.top)).2.2 =
      RootSharing.newTarget S (M.step q x (readTops (S.viewAt 2 S.top))).2 := by
  exact viewTarget_eq S _

/-! ## Exact simulation -/

structure Sim (c : Config Q A s) (q : Q) (S : Scaffold s (RootLabel s A)) : Prop where
  ctrl : c.ctrl = q
  stacks : ∀ j, RootSharing.decode S j S.top = c.stacks j
  backward : S.Backward
  clean : CleanMissing 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 : Machine Sym Q A s) :
    Sim M.initConfig (compile M).init (Scaffold.base s (RootLabel s A)) where
  ctrl := rfl
  stacks := fun j => by
    simp only [Machine.initConfig]
    exact decode_empty _ j 0 rfl
  backward := Scaffold.base_backward s (RootLabel s A)
  clean := by intro v j h; rfl
  aboveLabel := fun v hv => rfl
  aboveEdge := fun v hv i => rfl

theorem sim_tops {c : Config Q A s} {q : Q} {S : Scaffold s (RootLabel 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]
  cases hc : cellComp S S.top j with
  | none =>
      have hempty : c.stacks j = [] := by
        rw [← h.stacks j]
        exact RootSharing.decode_empty S j S.top hc
      simp [hempty]
  | some cell =>
      rw [← h.stacks j, RootSharing.decode_cons S h.backward j S.top cell hc]
      rfl

theorem compile_step_eq (M : Machine Sym Q A s) (q : Q)
    (S : Scaffold s (RootLabel s A)) (x : Sym) :
    (compile M).step (q, S) x =
      ((M.step q x (readTops (S.viewAt 2 S.top))).1,
        S.append
          (RootSharing.newLabel S (M.step q x (readTops (S.viewAt 2 S.top))).2)
          (RootSharing.newTarget S (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 : Machine Sym Q A s) {c : Config Q A s}
    {p : Q × Scaffold s (RootLabel 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
  change Sim c q S at h
  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 RootSharing.decode _ j (S.top + 1) =
      ((M.step c.ctrl x (fun j => (c.stacks j).head?)).2 j).apply c.stacks
    rw [hstep]
    exact RootSharing.decode_new S h.backward c.stacks h.stacks _ j
  · exact RootSharing.append_backward S h.backward _
  · exact RootSharing.cleanMissing_append S h.clean _
  · intro v hv
    change S.top + 1 < v at hv
    rw [S.append_label_old _ _ (by omega : v ≠ S.top + 1)]
    exact h.aboveLabel v (by omega)
  · intro v hv i
    change S.top + 1 < v at hv
    rw [S.append_edge_old _ _ (by omega : v ≠ S.top + 1) i]
    exact h.aboveEdge v (by omega) i

theorem sim_run (M : Machine Sym Q A s) :
    ∀ (w : List Sym) (c : Config Q A s) (p : Q × Scaffold s (RootLabel 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)

theorem compile_correct (M : Machine 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 (RootLabel s A)) (sim_init M)

theorem compile_ctrl (M : Machine 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 : Machine Sym Q A s) (w : List Sym) (j : Fin s) :
    RootSharing.decode ((compile M).run w).2 j ((compile M).run w).2.top =
      (M.run w).stacks j :=
  (compile_correct M w).stacks j

theorem compile_decides (M : Machine 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

/-! ## Shared-source calibration -/

def splitMachine : Machine Bool Bool Nat 2 where
  init := false
  accept := fun q => q
  step := fun _ x _ =>
    (x, fun j => if j = 0 then .push 0 7 else .push 0 8)

theorem splitMachine_source_run :
    (splitMachine.run [true]).stacks 0 = [7] ∧
      (splitMachine.run [true]).stacks 1 = [8] := by
  native_decide

theorem splitMachine_compiled_run :
    RootSharing.decode ((compile splitMachine).run [true]).2 0
        ((compile splitMachine).run [true]).2.top = [7] ∧
      RootSharing.decode ((compile splitMachine).run [true]).2 1
        ((compile splitMachine).run [true]).2.top = [8] := by
  constructor
  · rw [compile_stacks, splitMachine_source_run.1]
  · rw [compile_stacks, splitMachine_source_run.2]

theorem splitMachine_accepts_true : (compile splitMachine).decides [true] = true := by
  rw [compile_decides]
  native_decide

theorem splitMachine_rejects_empty : (compile splitMachine).decides [] = false := by
  rfl

#check compile_distanceOk
#check compile_selfFree
#check sim_step
#check sim_run
#check compile_decides
#check splitMachine_compiled_run

end FrontierDerivative.RootSharingMachine