RootSharing.lean

import FPRDTransfer.Compile

/-!
# Coordinate-carrying persistent root sharing

The fixed-port stack decoder cannot clone one historical stack into another
coordinate. Here each cell carries the coordinate to use *after* its symbol.
The decoder follows the current coordinate's port and then switches to the
carried coordinate. This permits several output slots to derive, with
different operations, from one old stack root.

This module proves the representation theorem only. It does not infer a
stationary slot allocator from a bound on frontier cardinality.
-/

namespace FrontierDerivative.RootSharing

open FPRD

universe u

/-- A symbol and the coordinate at which decoding continues after its edge. -/
structure Cell (s : Nat) (A : Type u) where
  symbol : A
  next : Fin s
deriving DecidableEq, Repr

abbrev RootLabel (s : Nat) (A : Type u) := Fin s → Option (Cell s A)

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

def cellComp (S : Scaffold s (RootLabel s A)) (v : Nat) (j : Fin s) :
    Option (Cell s A) :=
  (S.label v).bind fun g => g j

def cellCompOpt (S : Scaffold s (RootLabel s A)) (v : Option Nat) (j : Fin s) :
    Option (Cell s A) :=
  v.bind fun m => cellComp S m j

/-- Read component `j`, follow port `j`, then use the cell's next coordinate. -/
def decode (S : Scaffold s (RootLabel s A)) (j : Fin s) (v : Nat) : List A :=
  match cellComp S v j with
  | none => []
  | some c =>
    match S.edge v j with
    | none => [c.symbol]
    | some m => if _h : m < v then c.symbol :: decode S c.next m else [c.symbol]
termination_by v
decreasing_by exact _h

def decodeOpt (S : Scaffold s (RootLabel s A)) (j : Fin s) :
    Option Nat → List A
  | none => []
  | some m => decode S j m

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

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

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

theorem decode_cons (S : Scaffold s (RootLabel s A)) (hb : S.Backward)
    (j : Fin s) (v : Nat) (c : Cell s A) (h : cellComp S v j = some c) :
    decode S j v = c.symbol :: decodeOpt S c.next (S.edge v j) := by
  rw [decode, h]
  cases he : S.edge v j with
  | none => simp
  | some m =>
      have hlt : m < v := hb v j m he
      simp [hlt]

/-! A reducible twin supports exact kernel-checked calibrations. -/

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

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

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

/-- An empty cell owns no edge, ruling out stale tails. -/
def CleanMissing (S : Scaffold s (RootLabel s A)) : Prop :=
  ∀ v j, cellComp S v j = none → S.edge v j = none

/-! ## Congruence below a freshly appended node -/

theorem decodeAux_congr (S T : Scaffold s (RootLabel s A)) :
    ∀ (fuel : Nat) (j : Fin s) (v : Nat),
    (∀ u, u ≤ v → S.label u = T.label u) →
    (∀ u, u ≤ v → ∀ k, S.edge u k = T.edge u k) →
    decodeAux S fuel j v = decodeAux T fuel j v := by
  intro fuel
  induction fuel with
  | zero => intro j v hlab hedge; rfl
  | succ f ih =>
      intro j v hlab hedge
      have hcell : cellComp S v j = cellComp T v j := by
        simp only [cellComp, hlab v (Nat.le_refl v)]
      simp only [decodeAux, hcell]
      cases hc : cellComp T v j with
      | none => rfl
      | some c =>
          have he : S.edge v j = T.edge v j := hedge v (Nat.le_refl v) j
          rw [he]
          cases htarget : T.edge v j with
          | none => rfl
          | some m =>
              by_cases hlt : m < v
              · simp only [if_pos hlt, List.cons.injEq, true_and]
                exact ih c.next m
                  (fun u hu => hlab u (Nat.le_trans hu (Nat.le_of_lt hlt)))
                  (fun u hu => hedge u (Nat.le_trans hu (Nat.le_of_lt hlt)))
              · simp [hlt]

theorem decode_congr (S T : Scaffold s (RootLabel s A)) (v : Nat)
    (hlab : ∀ u, u ≤ v → S.label u = T.label u)
    (hedge : ∀ u, u ≤ v → ∀ j, S.edge u j = T.edge u j) :
    ∀ j, decode S j v = decode T j v := by
  intro j
  rw [decode_eq_decodeAux, decode_eq_decodeAux]
  exact decodeAux_congr S T (v + 1) j v hlab hedge

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

theorem decodeOpt_append_below (S : Scaffold s (RootLabel s A))
    (g : RootLabel 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 h : o with
  | none => rfl
  | some m =>
      simp only [decodeOpt_some]
      exact decode_append_below S g tgt j m (ho m h)

/-! ## Simultaneous source-selecting actions -/

inductive Derive (s : Nat) (A : Type u) where
  | push (source : Fin s) (symbol : A)
  | keep (source : Fin s)
  | pop (source : Fin s)
deriving Repr

namespace Derive

def source : Derive s A → Fin s
  | .push i _ | .keep i | .pop i => i

def apply (op : Derive s A) (stacks : Fin s → List A) : List A :=
  match op with
  | .push i a => a :: stacks i
  | .keep i => stacks i
  | .pop i => (stacks i).tail

end Derive

def tailCell (S : Scaffold s (RootLabel s A)) (i : Fin s) : Option (Cell s A) :=
  match cellComp S S.top i with
  | none => none
  | some c => cellCompOpt S (S.edge S.top i) c.next

def newLabel (S : Scaffold s (RootLabel s A)) (ops : Fin s → Derive s A) :
    RootLabel s A :=
  fun j =>
    match ops j with
    | .push i a => some ⟨a, i⟩
    | .keep i => cellComp S S.top i
    | .pop i => tailCell S i

/-- Each output coordinate owns its edge; pop uses the source's carried next
coordinate as the second port in a distance-two path. -/
def newTarget (S : Scaffold s (RootLabel s A)) (ops : Fin s → Derive s A) :
    Fin s → Target s :=
  fun j =>
    match ops j with
    | .push _ _ => .path []
    | .keep i =>
      match cellComp S S.top i with
      | none => .missing
      | some _ => .path [i]
    | .pop i =>
      match cellComp S S.top i with
      | none => .missing
      | some c =>
        match tailCell S i with
        | none => .missing
        | some _ => .path [i, c.next]

theorem newTarget_ne_self (S : Scaffold s (RootLabel s A))
    (ops : Fin s → Derive s A) (j : Fin s) : newTarget S ops j ≠ .self := by
  cases hop : ops j with
  | push i a => simp [newTarget, hop]
  | keep i =>
      cases hc : cellComp S S.top i <;> simp [newTarget, hop, hc]
  | pop i =>
      cases hc : cellComp S S.top i with
      | none => simp [newTarget, hop, hc]
      | some c =>
          cases ht : tailCell S i <;> simp [newTarget, hop, hc, ht]

theorem newTarget_distance (S : Scaffold s (RootLabel s A))
    (ops : Fin s → Derive s A) (j : Fin s) (p : List (Fin s))
    (h : newTarget S ops j = .path p) : p.length ≤ 2 := by
  cases hop : ops j with
  | push i a =>
      simp [newTarget, hop] at h
      subst p
      simp
  | keep i =>
      cases hc : cellComp S S.top i with
      | none => simp [newTarget, hop, hc] at h
      | some c =>
          simp [newTarget, hop, hc] at h
          subst p
          simp
  | pop i =>
      cases hc : cellComp S S.top i with
      | none => simp [newTarget, hop, hc] at h
      | some c =>
          cases ht : tailCell S i with
          | none => simp [newTarget, hop, hc, ht] at h
          | some d =>
              simp [newTarget, hop, hc, ht] at h
              subst p
              simp

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

theorem append_backward (S : Scaffold s (RootLabel s A)) (hb : S.Backward)
    (ops : Fin s → Derive s A) :
    (S.append (newLabel S ops) (newTarget S ops)).Backward :=
  S.append_backward hb _ _ (newTarget_ne_self S ops)

/-! ## Exact one-step root sharing -/

theorem decode_new (S : Scaffold s (RootLabel s A)) (hb : S.Backward)
    (stacks : Fin s → List A) (hst : ∀ i, decode S i S.top = stacks i)
    (ops : Fin s → Derive s A) (j : Fin s) :
    decode (S.append (newLabel S ops) (newTarget S ops)) j (S.top + 1) =
      (ops j).apply stacks := by
  let S' := S.append (newLabel S ops) (newTarget S ops)
  have hb' : S'.Backward := append_backward S hb ops
  have hlabel : cellComp S' (S.top + 1) j = newLabel S ops j := by
    simp only [S', cellComp, Scaffold.append_label_new, Option.bind_some]
  have hedge : S'.edge (S.top + 1) j =
      S.resolve S.top (S.top + 1) (newTarget S ops j) := by
    simp only [S', Scaffold.append_edge_new]
  cases hop : ops j with
  | push i a =>
      have hl : cellComp S' (S.top + 1) j = some ⟨a, i⟩ := by
        rw [hlabel]; simp [newLabel, hop]
      rw [decode_cons S' hb' j (S.top + 1) ⟨a, i⟩ hl, hedge]
      simp only [newTarget, hop, Scaffold.resolve, Scaffold.follow_nil,
        decodeOpt_some, Derive.apply]
      rw [decode_append_below S _ _ i S.top (Nat.le_refl _), hst i]
  | keep i =>
      cases hc : cellComp S S.top i with
      | none =>
          have hl : cellComp S' (S.top + 1) j = none := by
            rw [hlabel]; simp [newLabel, hop, hc]
          have hempty : stacks i = [] := by
            rw [← hst i]; exact decode_empty S i S.top hc
          rw [decode_empty S' j (S.top + 1) hl, Derive.apply, hempty]
      | some c =>
          have hl : cellComp S' (S.top + 1) j = some c := by
            rw [hlabel]; simp [newLabel, hop, hc]
          have ht : S'.edge (S.top + 1) j = S.edge S.top i := by
            rw [hedge]; simp [newTarget, hop, hc, Scaffold.resolve, Scaffold.follow_single]
          have hbound : ∀ m, S.edge S.top i = some m → m ≤ S.top := by
            intro m hm; exact Nat.le_of_lt (hb S.top i m hm)
          rw [decode_cons S' hb' j (S.top + 1) c hl, ht,
            decodeOpt_append_below S _ _ c.next (S.edge S.top i) hbound,
            ← decode_cons S hb i S.top c hc, hst i]
          rfl
  | pop i =>
      cases hc : cellComp S S.top i with
      | none =>
          have hl : cellComp S' (S.top + 1) j = none := by
            rw [hlabel]; simp [newLabel, tailCell, hop, hc]
          have hempty : stacks i = [] := by
            rw [← hst i]; exact decode_empty S i S.top hc
          rw [decode_empty S' j (S.top + 1) hl, Derive.apply, hempty]
          rfl
      | some c =>
          have hsource : stacks i = c.symbol :: decodeOpt S c.next (S.edge S.top i) := by
            rw [← hst i]; exact decode_cons S hb i S.top c hc
          cases he : S.edge S.top i with
          | none =>
              have hl : cellComp S' (S.top + 1) j = none := by
                rw [hlabel]; simp [newLabel, tailCell, hop, hc, he, cellCompOpt]
              rw [decode_empty S' j (S.top + 1) hl, Derive.apply, hsource, he]
              rfl
          | some m =>
              have hmlt : m < S.top := hb S.top i m he
              cases hd : cellComp S m c.next with
              | none =>
                  have hl : cellComp S' (S.top + 1) j = none := by
                    rw [hlabel]
                    simp [newLabel, tailCell, hop, hc, he, cellCompOpt, hd]
                  rw [decode_empty S' j (S.top + 1) hl, Derive.apply, hsource, he,
                    decodeOpt_some, decode_empty S c.next m hd]
                  rfl
              | some d =>
                  have hl : cellComp S' (S.top + 1) j = some d := by
                    rw [hlabel]
                    simp [newLabel, tailCell, hop, hc, he, cellCompOpt, hd]
                  have ht : S'.edge (S.top + 1) j = S.edge m c.next := by
                    rw [hedge]
                    simp [newTarget, hop, hc, tailCell, he, cellCompOpt, hd,
                      resolve_mixed_two]
                  have hbound : ∀ m', S.edge m c.next = some m' → m' ≤ S.top := by
                    intro m' hm'
                    have hm'lt : m' < m := hb m c.next m' hm'
                    omega
                  rw [decode_cons S' hb' j (S.top + 1) d hl, ht,
                    decodeOpt_append_below S _ _ d.next (S.edge m c.next) hbound,
                    Derive.apply, hsource, he, decodeOpt_some,
                    ← decode_cons S hb c.next m d hd]
                  rfl

theorem cleanMissing_append (S : Scaffold s (RootLabel s A))
    (hc : CleanMissing S) (ops : Fin s → Derive s A) :
    CleanMissing (S.append (newLabel S ops) (newTarget S ops)) := by
  intro v j hv
  by_cases hnew : v = S.top + 1
  · subst hnew
    simp only [cellComp, Scaffold.append_label_new, Option.bind_some] at hv
    rw [Scaffold.append_edge_new]
    cases hop : ops j with
    | push i a => simp [newLabel, hop] at hv
    | keep i =>
        simp only [newLabel, hop] at hv
        simp [newTarget, hop, hv, Scaffold.resolve]
    | pop i =>
        simp only [newLabel, hop] at hv
        simp only [newTarget, hop]
        cases hs : cellComp S S.top i with
        | none => simp [Scaffold.resolve]
        | some c =>
            have ht : tailCell S i = none := hv
            simp [ht, Scaffold.resolve]
  · rw [S.append_edge_old (newLabel S ops) (newTarget S ops) hnew j]
    apply hc v j
    rw [cellComp] at hv ⊢
    rw [S.append_label_old (newLabel S ops) (newTarget S ops) hnew] at hv
    exact hv

/-! ## Exact shared-tail calibration -/

def firstOps (j : Fin 2) : Derive 2 Nat :=
  if j = 0 then .push 0 1 else .push 1 9

def rootBase : Scaffold 2 (RootLabel 2 Nat) :=
  Scaffold.base 2 (RootLabel 2 Nat)

def history1 : Scaffold 2 (RootLabel 2 Nat) :=
  rootBase.append (newLabel rootBase firstOps) (newTarget rootBase firstOps)

def secondOps (j : Fin 2) : Derive 2 Nat :=
  if j = 0 then .push 0 2 else .keep 1

def history : Scaffold 2 (RootLabel 2 Nat) :=
  history1.append (newLabel history1 secondOps) (newTarget history1 secondOps)

def splitOps (j : Fin 2) : Derive 2 Nat :=
  if j = 0 then .push 0 7 else .push 0 8

def splitHistory : Scaffold 2 (RootLabel 2 Nat) :=
  history.append (newLabel history splitOps) (newTarget history splitOps)

theorem rootBase_backward : rootBase.Backward :=
  Scaffold.base_backward 2 (RootLabel 2 Nat)

theorem rootBase_cleanMissing : CleanMissing rootBase := by
  intro v j h
  rfl

theorem history1_backward : history1.Backward :=
  append_backward rootBase rootBase_backward firstOps

theorem history1_cleanMissing : CleanMissing history1 :=
  cleanMissing_append rootBase rootBase_cleanMissing firstOps

theorem history_backward : history.Backward :=
  append_backward history1 history1_backward secondOps

theorem history_cleanMissing : CleanMissing history :=
  cleanMissing_append history1 history1_cleanMissing secondOps

theorem splitHistory_backward : splitHistory.Backward :=
  append_backward history history_backward splitOps

theorem splitHistory_cleanMissing : CleanMissing splitHistory :=
  cleanMissing_append history history_cleanMissing splitOps

theorem source_before_split : decode history 0 history.top = [2, 1] := by
  rw [decode_eq_decodeAux]
  native_decide

theorem destination_before_split : decode history 1 history.top = [9] := by
  rw [decode_eq_decodeAux]
  native_decide

theorem first_after_split : decode splitHistory 0 splitHistory.top = [7, 2, 1] := by
  rw [decode_eq_decodeAux]
  native_decide

theorem second_after_split : decode splitHistory 1 splitHistory.top = [8, 2, 1] := by
  rw [decode_eq_decodeAux]
  native_decide

#check decode_new
#check append_backward
#check cleanMissing_append
#check splitHistory_backward
#check splitHistory_cleanMissing
#check first_after_split
#check second_after_split

end FrontierDerivative.RootSharing