import RoleFlow.FrontSem
import RoleFlow.Infer

/-! # The front end is complete (module `FrontComplete`)

If some swap of `P` (`SwapLine`) is typed, `swapArgs P` is typed with the same roles
(`swapArgs_complete`); with `Prog.infer_iff`, inference after the front end is exact up to swaps
(`infer_swapArgs_iff`).

Fix the typed swap and its roles. Its recipes `recs`, one per line, are equations of the search
(`RecsOK.sub`): the forward form of a conditioner, the solved or swapped solved form of a spine
step. The search fires only recipes: every other equation of a line has a back edge to the
line's recipe (`line_back_cond`, `line_back_spine`), reading a variable whose recipe reads one of
its results, so it cannot fire while the known set is closed under the recipes (`only_A`). The
search stops at a fixpoint (`fwdKnown_terminal`), where everything the recipes derive from the
weights and the returned variables is known (`fix_forward`); that is every variable
(`Typed.known_xrecs`), so every recipe has fired (`fwdKnown_exact`). Hence the swapped forms that fired
are exactly the swapped spine steps of the typed swap, and the swaps `swapArgs` makes differ from
it only on conditioners, whose argument order the typing ignores (`typed_congr`). The typing
supplies unique binders and results outside the weights and the returned variables
(`recs_ok`).

The search itself is `Search`; this module supplies its hypotheses for the lines. -/

namespace RoleFlow

open Search

variable {S : Type} {Sg : Sig S} [Swap Sg]

section Lines
variable {s : Bool} {l : Line Sg} {i n k : ℕ}

theorem swapIf_bs : (swapIf s l).bs = l.bs := by unfold swapIf; split <;> rfl

theorem swapIf_ys : (swapIf s l).ys n = l.ys n := by simp only [Line.ys, swapIf_bs]

theorem kOf_swapIf : kOf (swapIf s l) = kOf l := by unfold swapIf; split <;> rfl

theorem mem_swapIf_args {x : ℕ} : x ∈ (swapIf s l).args ↔ x ∈ l.args := by
  unfold swapIf; split <;> simp

theorem length_swapIf_args : (swapIf s l).args.length = l.args.length := by
  unfold swapIf; split <;> simp

theorem swapIf_of_not (h : ¬(Swap.swappable l.op = true ∧ l.args.length = 2)) : swapIf s l = l := by
  unfold swapIf; split
  · rename_i h'; simp only [Bool.and_eq_true, beq_iff_eq] at h'; exact absurd ⟨h'.1.2, h'.2⟩ h
  · rfl

theorem swapIf_args_add {a b : ℕ} (hA : Swap.swappable l.op = true) (hab : l.args = [a, b]) :
    (swapIf s l).args = if s then [b, a] else [a, b] := by
  cases s <;> simp [swapIf, hA, hab]

theorem kOf_add (h : Swap.swappable l.op = true) : kOf l = some 1 := by
  obtain ⟨si, hsi, hk⟩ := Swap.k_one l.op h
  simp [kOf, hsi, hk]

omit [Swap Sg] in
theorem kOf_spec (h : kOf l = some k) : ∃ si, Sg.inv l.op = some si ∧ si.k = k := by
  unfold kOf at h; obtain ⟨si, hsi, rfl⟩ := Option.map_eq_some_iff.1 h; exact ⟨si, hsi, rfl⟩

omit [Swap Sg] in
theorem kOf_of_inv {si : SemiInv Sg.Op l.as l.bs} (h : Sg.inv l.op = some si) :
    kOf l = some si.k := by simp [kOf, h]

omit [Swap Sg] in
theorem kOf_pos (h : kOf l = some k) : 1 ≤ k ∧ k ≤ l.as.length := by
  obtain ⟨si, -, rfl⟩ := kOf_spec h
  have h₁ := congrArg List.length si.split
  have h₂ := List.length_pos_of_ne_nil si.cs_ne
  rw [List.length_append] at h₁
  show 1 ≤ si.cs.length ∧ si.cs.length ≤ l.as.length
  omega

omit [Swap Sg] in
theorem ys_ne : l.ys n ≠ [] := by
  have := Sg.bs_ne l.op
  simp only [Line.ys, ne_eq, List.range'_eq_nil_iff]; exact fun h => this (List.eq_nil_of_length_eq_zero h)

theorem be {e q : Eqn} {u v : ℕ} (h₁ : u ∈ q.reads) (h₂ : v ∈ q.binds) (h₃ : u ∈ e.binds)
    (h₄ : v ∈ e.reads) : BackEdge [e] q :=
  ⟨u, h₁, v, h₂, e, List.mem_singleton_self _, h₃, h₄⟩

/-- The recipe of the line `l`, numbered `i`, at scope `n`, under a role, swapped if `s`: the forward form of a conditioner, the (swapped) solved form of a spine step. -/
def lineRecipe (s : Bool) (i n : ℕ) (l : Line Sg) : Role → List Eqn
  | .cond => [⟨l.ys n, l.args, i, false⟩]
  | .spine => match kOf l with
    | some k => [⟨(swapIf s l).args.take k, l.ys n ++ (swapIf s l).args.drop k, i,
        s && Swap.swappable l.op && l.args.length == 2⟩]
    | none => []

/-- **A conditioner's other equations have back edges to its forward form.** -/
theorem line_back_cond (hlen : ∀ k, kOf l = some k → k ≤ l.args.length) :
    ∀ q ∈ lineEqns i n l, q = ⟨l.ys n, l.args, i, false⟩ ∨ BackEdge [⟨l.ys n, l.args, i, false⟩] q := by
  obtain ⟨y, hy⟩ := List.exists_mem_of_ne_nil _ (ys_ne (l := l) (n := n))
  intro q hq
  simp only [lineEqns, List.cons_append, List.mem_cons, List.mem_append] at hq
  rcases hq with rfl | hq | hq
  · exact Or.inl rfl
  · unfold solvedEqns at hq; split at hq
    · rename_i k hk
      obtain ⟨h₁, -⟩ := kOf_pos hk
      have := hlen k hk
      obtain ⟨x, hx⟩ := List.exists_mem_of_ne_nil (l.args.take k)
        (List.ne_nil_of_length_pos (by rw [List.length_take]; omega))
      simp only [List.mem_singleton] at hq; subst hq
      exact Or.inr (be (u := y) (v := x) (by simp [hy]) hx hy (List.mem_of_mem_take hx))
    · simp at hq
  · unfold swappedEqns at hq; split at hq
    · rename_i a b _ hab
      simp only [List.mem_singleton] at hq; subst hq
      exact Or.inr (be (u := y) (v := b) (by simp [hy]) (by simp) hy (by simp [hab]))
    · simp at hq

/-- **A spine step's other equations have back edges to its (swapped) solved form.** -/
theorem line_back_spine (hk : kOf l = some k) (hlen : k ≤ l.args.length) :
    let e : Eqn := ⟨(swapIf s l).args.take k, l.ys n ++ (swapIf s l).args.drop k, i,
      s && Swap.swappable l.op && l.args.length == 2⟩
    e ∈ lineEqns i n l ∧ ∀ q ∈ lineEqns i n l, q = e ∨ BackEdge [e] q := by
  intro e
  obtain ⟨y, hy⟩ := List.exists_mem_of_ne_nil _ (ys_ne (l := l) (n := n))
  by_cases hadd : Swap.swappable l.op = true ∧ l.args.length = 2
  · obtain ⟨hA, h2⟩ := hadd
    obtain ⟨a, b, hab⟩ : ∃ a b, l.args = [a, b] := by
      match h : l.args, h2 with
      | [a, b], _ => exact ⟨a, b, rfl⟩
    cases (kOf_add hA).symm.trans hk
    have hE : lineEqns i n l = [⟨l.ys n, [a, b], i, false⟩, ⟨[a], l.ys n ++ [b], i, false⟩,
        ⟨[b], l.ys n ++ [a], i, true⟩] := by
      simp [lineEqns, solvedEqns, swappedEqns, kOf_add hA, hA, hab]
    have he : e = if s then ⟨[b], l.ys n ++ [a], i, true⟩ else ⟨[a], l.ys n ++ [b], i, false⟩ := by
      cases s <;> simp [e, swapIf_args_add hA hab, hA, hab]
    rw [hE]
    cases s <;> simp only [Bool.false_eq_true, ↓reduceIte] at he <;> rw [he] <;>
      refine ⟨by simp, fun q hq => ?_⟩ <;>
      simp only [List.mem_cons, List.not_mem_nil, or_false] at hq <;>
      rcases hq with rfl | rfl | rfl
    · exact Or.inr (be (u := a) (v := y) (by simp) hy (by simp) (by simp [hy]))
    · exact Or.inl rfl
    · exact Or.inr (be (u := a) (v := b) (by simp) (by simp) (by simp) (by simp))
    · exact Or.inr (be (u := b) (v := y) (by simp) hy (by simp) (by simp [hy]))
    · exact Or.inr (be (u := b) (v := a) (by simp) (by simp) (by simp) (by simp))
    · exact Or.inl rfl
  · have hsw := swapIf_of_not (s := s) hadd
    have hflag : (s && Swap.swappable l.op && l.args.length == 2) = false := by
      simp only [Bool.and_eq_false_iff, beq_eq_false_iff_ne]
      by_cases h₁ : Swap.swappable l.op = true
      · exact Or.inr fun h₂ => hadd ⟨h₁, h₂⟩
      · exact Or.inl (Or.inr (by simpa using h₁))
    have hS : swappedEqns i n l = [] := by
      unfold swappedEqns; split
      · rename_i a b hA hab; exact absurd ⟨hA, by simp [hab]⟩ hadd
      · rfl
    have he : e = ⟨l.args.take k, l.ys n ++ l.args.drop k, i, false⟩ := by
      simp only [e, hsw, hflag]
    have hE : lineEqns i n l = [⟨l.ys n, l.args, i, false⟩, e] := by
      simp [lineEqns, solvedEqns, hk, hS, he]
    rw [hE]
    refine ⟨by simp, fun q hq => ?_⟩
    simp only [List.mem_cons, List.not_mem_nil, or_false] at hq
    rcases hq with rfl | rfl
    · obtain ⟨h₁, -⟩ := kOf_pos hk
      obtain ⟨x, hx⟩ := List.exists_mem_of_ne_nil (l.args.take k)
        (List.ne_nil_of_length_pos (by rw [List.length_take]; omega))
      exact Or.inr (be (u := x) (v := y) (List.mem_of_mem_take hx) hy (by simp [he, hx])
        (by simp [he, hy]))
    · exact Or.inl rfl

end Lines

theorem swapLine_swapIf₂ (s t : Bool) (l : Line Sg) : SwapLine (swapIf s l) (swapIf t l) := by
  by_cases hadd : Swap.swappable l.op = true ∧ l.args.length = 2
  · rcases l with ⟨op, args⟩
    obtain ⟨hA, h2⟩ := hadd
    match args, h2 with
    | [a, b], _ => cases s <;> cases t <;> simp [swapIf, hA] <;> first
      | exact .refl _ | exact .swap op hA a b | exact .swap op hA b a
  · rw [swapIf_of_not hadd, swapIf_of_not hadd]; exact .refl _

/-! ## The typing's recipes -/

/-- The recipes of the lines numbered from `i` at scope `n` under the roles `ρ`, with the lines
selected by `sw` swapped. -/
def recs (sw : ℕ → Bool) : ℕ → ℕ → List (Line Sg) → List Role → List Eqn
  | i, n, l :: ls, r :: ρ => lineRecipe (sw i) i n l r ++ recs sw (i + 1) (n + l.bs.length) ls ρ
  | _, _, _, _ => []

/-- Two line lists that agree on the spine steps and differ at most by a swap on conditioners. -/
def SwapAgree : List (Line Sg) → List (Line Sg) → List Role → Prop
  | l₁ :: ls₁, l₂ :: ls₂, .cond :: ρ => SwapLine l₁ l₂ ∧ SwapAgree ls₁ ls₂ ρ
  | l₁ :: ls₁, l₂ :: ls₂, .spine :: ρ => l₁ = l₂ ∧ SwapAgree ls₁ ls₂ ρ
  | [], [], [] => True
  | _, _, _ => False

/-- **The typing ignores the argument order of conditioners.** -/
theorem typed_congr {O Γ : List S} {Λ : Finset ℕ} {D : List (Finset ℕ)}
    {ls₁ ls₂ : List (Line Sg)} {xs : List ℕ} {ρ : List Role}
    (hT : Typed Sg O Γ Λ D ls₁ xs ρ) (ha : SwapAgree ls₁ ls₂ ρ) : Typed Sg O Γ Λ D ls₂ xs ρ := by
  induction ls₁ generalizing Γ Λ D ls₂ ρ with
  | nil =>
    rcases ls₂ with _ | ⟨l₂, ls₂⟩ <;> rcases ρ with _ | ⟨r, ρ⟩
    · exact hT
    · exact hT.elim
    · exact ha.elim
    · cases r <;> exact ha.elim
  | cons l₁ ls₁ ih =>
    rcases ls₂ with _ | ⟨l₂, ls₂⟩ <;> rcases ρ with _ | ⟨r, ρ⟩
    · exact hT.elim
    · cases r <;> exact ha.elim
    · exact hT.elim
    · cases r with
      | cond =>
        obtain ⟨hl, ha⟩ := ha
        obtain ⟨e₁, e₂⟩ : condSupp D l₂ = condSupp D l₁ ∧ l₂.bs = l₁.bs := by
          cases hl with
          | refl => exact ⟨rfl, rfl⟩
          | swap => exact ⟨condSupp_congr rfl fun x => by simp; tauto, rfl⟩
        refine ⟨hl.args hT.1, ?_⟩
        rw [e₁, e₂]; exact ih hT.2 ha
      | spine =>
        obtain ⟨rfl, ha⟩ := ha
        obtain ⟨hA, si, hsi, h₁, h₂, h₃, hT⟩ := hT
        exact ⟨hA, si, hsi, h₁, h₂, h₃, ih hT ha⟩

/-- What the typing of a swap gives about its recipes. -/
structure RecsOK (sw : ℕ → Bool) (nw i n : ℕ) (Λ : Finset ℕ) (ls : List (Line Sg))
    (xs : List ℕ) (ρ : List Role) : Prop where
  uniq : UniqB (recs sw i n ls ρ)
  binds : ∀ r ∈ recs sw i n ls ρ, ∀ v ∈ r.binds, (v ∈ Λ ∨ n ≤ v) ∧ nw ≤ v ∧ v ∉ xs
  ne : ∀ r ∈ recs sw i n ls ρ, r.binds ≠ []
  back : ∀ q ∈ eqns i n ls, q ∈ recs sw i n ls ρ ∨ BackEdge (recs sw i n ls ρ) q
  sub : ∀ r ∈ recs sw i n ls ρ, r ∈ eqns i n ls
  line : ∀ r ∈ recs sw i n ls ρ, i ≤ r.line
  agree : ∀ S : ℕ → Bool, (∀ j, i ≤ j → (S j = true ↔ ∃ q ∈ recs sw i n ls ρ, q.swap = true ∧ q.line = j)) →
    SwapAgree (swapLines sw i ls) (swapLines S i ls) ρ

theorem exists_cons_line {e : Eqn} {A : List Eqn} {j : ℕ} (he : e.line ≠ j) :
    (∃ q ∈ e :: A, q.swap = true ∧ q.line = j) ↔ ∃ q ∈ A, q.swap = true ∧ q.line = j := by
  constructor
  · rintro ⟨q, hq, h₁, h₂⟩
    rcases List.mem_cons.1 hq with rfl | hq
    · exact absurd h₂ he
    · exact ⟨q, hq, h₁, h₂⟩
  · rintro ⟨q, hq, h⟩; exact ⟨q, List.mem_cons_of_mem _ hq, h⟩

/-- **The recipes of a typed swap**: unique binders, results live or later and outside the
weights and the returned variables, back edges, and the swap flags. -/
theorem recs_ok (sw : ℕ → Bool) (nw : ℕ) : ∀ (ls : List (Line Sg)) {i n : ℕ} {Λ : Finset ℕ}
    {D : List (Finset ℕ)} {xs : List ℕ} {ρ : List Role} {O Γ : List S},
    Typed Sg O Γ Λ D (swapLines sw i ls) xs ρ → LinesWF Γ ls → Γ.length = n →
    (∀ v ∈ Λ, nw ≤ v ∧ v < n) → nw ≤ n → RecsOK sw nw i n Λ ls xs ρ
  | [], i, n, Λ, D, xs, ρ, O, Γ, hT, _, _, _, _ => by
    cases ρ with
    | nil =>
      exact ⟨by simp [UniqB, recs], by simp [recs], by simp [recs], by simp [eqns],
        by simp [recs], by simp [recs], fun _ _ => by simp [swapLines, SwapAgree]⟩
    | cons => simp [swapLines, Typed] at hT
  | l :: ls, i, n, Λ, D, xs, ρ, O, Γ, hT, hwf, hn, hΛ, hnw => by
    have hlen : ∀ k, kOf l = some k → k ≤ l.args.length := fun k hk => by
      rw [hwf.1.length_eq]; exact (kOf_pos hk).2
    have hb : (swapIf (sw i) l).bs.length = l.bs.length := by rw [swapIf_bs]
    have heqns : eqns i n (l :: ls) = lineEqns i n l ++ eqns (i + 1) (n + l.bs.length) ls := rfl
    cases ρ with
    | nil => simp [swapLines, Typed] at hT
    | cons r ρ =>
    cases r with
    | cond =>
      have hT' : Typed Sg O (Γ ++ l.bs) Λ (condSupp D (swapIf (sw i) l))
          (swapLines sw (i + 1) ls) xs ρ := by
        have h := hT; simp only [swapLines, Typed] at h; rw [swapIf_bs] at h; exact h.2
      have ih := recs_ok sw nw ls (n := n + l.bs.length) hT' hwf.2 (by simp [hn])
        (fun v hv => ⟨(hΛ v hv).1, by have := (hΛ v hv).2; omega⟩) (by omega)
      have hfin := typed_final hT'
      simp only [List.length_append, hn] at hfin
      set e : Eqn := ⟨l.ys n, l.args, i, false⟩
      have hrecs : recs sw i n (l :: ls) (.cond :: ρ) = e :: recs sw (i + 1) (n + l.bs.length) ls ρ := rfl
      have hey : ∀ v ∈ e.binds, n ≤ v ∧ v < n + l.bs.length := fun v hv => Line.mem_ys.1 hv
      refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩ <;> simp only [hrecs, heqns]
      · intro r₁ h₁ r₂ h₂ v hv₁ hv₂
        rcases List.mem_cons.1 h₁ with rfl | h₁ <;> rcases List.mem_cons.1 h₂ with rfl | h₂
        · rfl
        · have := hey v hv₁
          rcases (ih.binds r₂ h₂ v hv₂).1 with h | h
          · have := (hΛ v h).2; omega
          · omega
        · have := hey v hv₂
          rcases (ih.binds r₁ h₁ v hv₁).1 with h | h
          · have := (hΛ v h).2; omega
          · omega
        · exact ih.uniq r₁ h₁ r₂ h₂ v hv₁ hv₂
      · intro r hr v hv
        rcases List.mem_cons.1 hr with rfl | hr
        · have := hey v hv
          refine ⟨Or.inr this.1, by omega, fun hx => ?_⟩
          rcases hfin v hx with h | h
          · have := (hΛ v h).2; omega
          · omega
        · obtain ⟨h₁, h₂, h₃⟩ := ih.binds r hr v hv
          exact ⟨h₁.imp id fun h => by omega, h₂, h₃⟩
      · intro r hr
        rcases List.mem_cons.1 hr with rfl | hr
        · exact ys_ne
        · exact ih.ne r hr
      · intro q hq
        rcases List.mem_append.1 hq with hq | hq
        · rcases line_back_cond hlen q hq with rfl | h
          · exact Or.inl (List.mem_cons_self ..)
          · exact Or.inr (h.mono (List.mem_cons_self ..))
        · exact (ih.back q hq).imp (List.mem_cons_of_mem _)
            fun h => h.mono' fun _ => List.mem_cons_of_mem _
      · intro r hr
        rcases List.mem_cons.1 hr with rfl | hr
        · exact List.mem_append_left _ (List.mem_cons_self ..)
        · exact List.mem_append_right _ (ih.sub r hr)
      · intro r hr
        rcases List.mem_cons.1 hr with rfl | hr
        · exact le_rfl
        · have := ih.line r hr; omega
      · intro S hS
        refine ⟨swapLine_swapIf₂ _ _ l, ih.agree S fun j hj => ?_⟩
        rw [hS j (by omega), exists_cons_line (by simp [e]; omega)]
    | spine =>
      obtain ⟨_, si, hsi, h₁, h₂, h₃, hT'⟩ := hT
      rw [show Γ ++ (swapIf (sw i) l).bs = Γ ++ l.bs by rw [swapIf_bs], hn] at hT'
      have hk : kOf l = some si.k := kOf_swapIf (l := l) ▸ kOf_of_inv hsi
      have hlive : ∀ v ∈ spineLive Λ (swapIf (sw i) l) si.k n,
          (v ∈ Λ ∧ v ∉ (swapIf (sw i) l).args.take si.k) ∨ (n ≤ v ∧ v < n + l.bs.length) := by
        intro v hv
        simpa [spineLive, swapIf_bs] using hv
      have ih := recs_ok sw nw ls (n := n + l.bs.length) hT' hwf.2 (by simp [hn])
        (fun v hv => by
          rcases hlive v hv with h | h
          · have := hΛ v h.1; omega
          · omega) (by omega)
      have hfin := typed_final hT'
      simp only [List.length_append, hn] at hfin
      set e : Eqn := ⟨(swapIf (sw i) l).args.take si.k, l.ys n ++ (swapIf (sw i) l).args.drop si.k,
        i, sw i && Swap.swappable l.op && l.args.length == 2⟩
      have hrecs : recs sw i n (l :: ls) (.spine :: ρ) =
          e :: recs sw (i + 1) (n + l.bs.length) ls ρ := by
        simp only [recs, lineRecipe, hk, e, List.singleton_append]
      have hey : ∀ v ∈ e.binds, v ∈ Λ ∧ v < n := fun v hv => ⟨h₂ v hv, (hΛ v (h₂ v hv)).2⟩
      -- a variable consumed here is not live, nor introduced, afterwards
      have hgone : ∀ v ∈ e.binds, v ∈ spineLive Λ (swapIf (sw i) l) si.k n ∨
          n + l.bs.length ≤ v → False := by
        intro v hv h
        have := (hey v hv).2
        rcases h with h | h
        · rcases hlive v h with h | h
          · exact h.2 hv
          · omega
        · omega
      obtain ⟨heE, hback⟩ := line_back_spine (s := sw i) (i := i) (n := n) hk (hlen _ hk)
      refine ⟨?_, ?_, ?_, ?_, ?_, ?_, ?_⟩ <;> simp only [hrecs, heqns]
      · intro r₁ hr₁ r₂ hr₂ v hv₁ hv₂
        rcases List.mem_cons.1 hr₁ with rfl | hr₁ <;> rcases List.mem_cons.1 hr₂ with rfl | hr₂
        · rfl
        · exact (hgone v hv₁ (ih.binds r₂ hr₂ v hv₂).1).elim
        · exact (hgone v hv₂ (ih.binds r₁ hr₁ v hv₁).1).elim
        · exact ih.uniq r₁ hr₁ r₂ hr₂ v hv₁ hv₂
      · intro r hr v hv
        rcases List.mem_cons.1 hr with rfl | hr
        · exact ⟨Or.inl (hey v hv).1, (hΛ v (hey v hv).1).1, fun hx => hgone v hv (hfin v hx)⟩
        · obtain ⟨h₄, h₅, h₆⟩ := ih.binds r hr v hv
          refine ⟨?_, h₅, h₆⟩
          rcases h₄ with h | h
          · rcases hlive v h with h | h
            · exact Or.inl h.1
            · exact Or.inr h.1
          · exact Or.inr (by omega)
      · intro r hr
        rcases List.mem_cons.1 hr with rfl | hr
        · obtain ⟨h₄, -⟩ := kOf_pos hk
          have := hlen _ hk
          exact List.ne_nil_of_length_pos (by rw [List.length_take, length_swapIf_args]; omega)
        · exact ih.ne r hr
      · intro q hq
        rcases List.mem_append.1 hq with hq | hq
        · rcases hback q hq with rfl | h
          · exact Or.inl (List.mem_cons_self ..)
          · exact Or.inr (h.mono (List.mem_cons_self ..))
        · exact (ih.back q hq).imp (List.mem_cons_of_mem _)
            fun h => h.mono' fun _ => List.mem_cons_of_mem _
      · intro r hr
        rcases List.mem_cons.1 hr with rfl | hr
        · exact List.mem_append_left _ heE
        · exact List.mem_append_right _ (ih.sub r hr)
      · intro r hr
        rcases List.mem_cons.1 hr with rfl | hr
        · exact le_rfl
        · have := ih.line r hr; omega
      · intro S hS
        refine ⟨?_, ih.agree S fun j hj => ?_⟩
        · by_cases hadd : Swap.swappable l.op = true ∧ l.args.length = 2
          · have hSi : S i = sw i := by
              have h := hS i le_rfl
              simp only [List.mem_cons, exists_eq_or_imp] at h
              have hno : ¬∃ q ∈ recs sw (i + 1) (n + l.bs.length) ls ρ, q.swap = true ∧ q.line = i :=
                fun ⟨q, hq, _, hl⟩ => by have := ih.line q hq; omega
              simp only [e, hadd.1, hadd.2, beq_self_eq_true, Bool.and_true, and_true, hno,
                or_false] at h
              cases h₁ : S i <;> cases h₂ : sw i <;> simp_all
            rw [hSi]
          · rw [swapIf_of_not hadd, swapIf_of_not hadd]
        · rw [hS j (by omega), exists_cons_line (by simp [e]; omega)]


/-! ## Every result of the recipes is derivable -/

/-- Each executable recipe of the swap is covered by a recipe of `recs`: it reads at most the
recipe's arguments and computes at most its results. -/
theorem xrecs_recs (sw : ℕ → Bool) : ∀ (ls : List (Line Sg)) {i n : ℕ} {ρ : List Role},
    ∀ x ∈ xrecs n (swapLines sw i ls) ρ, ∃ a ∈ recs sw i n ls ρ,
      (∀ w ∈ a.reads, w ∈ x.line.args) ∧ ∀ v ∈ x.tgt, v ∈ a.binds
  | [], _, _, ρ, x, hx => by cases ρ <;> simp [swapLines, xrecs] at hx
  | l :: ls, i, n, ρ, x, hx => by
    have hb : (swapIf (sw i) l).bs.length = l.bs.length := by rw [swapIf_bs]
    have tail : ∀ {r : Role} {ρ : List Role}, x ∈ xrecs (n + l.bs.length) (swapLines sw (i + 1) ls) ρ →
        ∃ a ∈ recs sw i n (l :: ls) (r :: ρ), (∀ w ∈ a.reads, w ∈ x.line.args) ∧
          ∀ v ∈ x.tgt, v ∈ a.binds := fun hx => by
      obtain ⟨a, ha, h⟩ := xrecs_recs sw ls x hx
      exact ⟨a, List.mem_append_right _ ha, h⟩
    cases ρ with
    | nil => simp [swapLines, xrecs] at hx
    | cons r ρ =>
      cases r with
      | cond =>
        simp only [swapLines, xrecs_cond, hb, List.mem_cons] at hx
        rcases hx with rfl | hx
        · exact ⟨_, List.mem_append_left _ (List.mem_singleton_self _),
            fun w hw => mem_swapIf_args.2 hw, fun v hv => by rwa [swapIf_ys] at hv⟩
        · exact tail hx
      | spine =>
        simp only [swapLines] at hx
        cases hsi : Sg.inv (swapIf (sw i) l).op with
        | none =>
          simp only [xrecs, hsi, hb] at hx
          exact tail hx
        | some si =>
          rw [xrecs_spine _ _ _ _ hsi, hb, List.mem_cons] at hx
          have hk : kOf l = some si.k := kOf_swapIf (l := l) ▸ kOf_of_inv hsi
          rcases hx with rfl | hx
          · have hm : (⟨(swapIf (sw i) l).args.take si.k, l.ys n ++ (swapIf (sw i) l).args.drop si.k,
                i, sw i && Swap.swappable l.op && l.args.length == 2⟩ : Eqn) ∈ recs sw i n (l :: ls) (.spine :: ρ) := by
              simp [recs, lineRecipe, hk]
            exact ⟨_, hm, fun w hw => by simpa only [swapIf_ys] using hw, fun v hv => hv⟩
          · exact tail hx

/-- Each recipe of `recs` computes results of an executable recipe of the swap. -/
theorem recs_xrecs (sw : ℕ → Bool) : ∀ (ls : List (Line Sg)) {i n : ℕ} {ρ : List Role},
    ∀ a ∈ recs sw i n ls ρ, ∃ x ∈ xrecs n (swapLines sw i ls) ρ, ∀ v ∈ a.binds, v ∈ x.tgt
  | [], _, _, ρ, a, ha => by cases ρ <;> simp [recs] at ha
  | l :: ls, i, n, ρ, a, ha => by
    have hb : (swapIf (sw i) l).bs.length = l.bs.length := by rw [swapIf_bs]
    cases ρ with
    | nil => simp [recs] at ha
    | cons r ρ =>
      simp only [recs, List.mem_append] at ha
      rcases ha with ha | ha
      · cases r with
        | cond =>
          simp only [lineRecipe, List.mem_singleton] at ha; subst ha
          exact ⟨_, List.mem_cons_self .., fun v hv => by rwa [swapIf_ys]⟩
        | spine =>
          cases hsi : Sg.inv (swapIf (sw i) l).op with
          | none =>
            have hk : kOf l = none := by rw [← kOf_swapIf (s := sw i)]; simp [kOf, hsi]
            simp [lineRecipe, hk] at ha
          | some si =>
            have hk : kOf l = some si.k := kOf_swapIf (l := l) ▸ kOf_of_inv hsi
            simp only [lineRecipe, hk, List.mem_singleton] at ha; subst ha
            simp only [swapLines]; rw [xrecs_spine _ _ _ _ hsi]
            exact ⟨_, List.mem_cons_self .., fun v hv => hv⟩
      · obtain ⟨x, hx, h⟩ := recs_xrecs sw ls a ha
        refine ⟨x, ?_, h⟩
        simp only [swapLines]
        cases r with
        | cond => rw [xrecs_cond, hb]; exact List.mem_cons_of_mem _ hx
        | spine =>
          cases hsi : Sg.inv (swapIf (sw i) l).op with
          | none => simp only [xrecs, hsi, hb]; exact hx
          | some si => rw [xrecs_spine _ _ _ _ hsi, hb]; exact List.mem_cons_of_mem _ hx

/-! ## Completeness -/

theorem swapLines_congr {sw sw' : ℕ → Bool} : ∀ {i : ℕ} (ls : List (Line Sg)),
    (∀ j, i ≤ j → sw j = sw' j) → swapLines sw i ls = swapLines sw' i ls
  | _, [], _ => rfl
  | i, l :: ls, h => by
    simp only [swapLines, h i le_rfl]
    rw [swapLines_congr ls fun j hj => h j (by omega)]

theorem swapLines_upd (b : Bool) (sw : ℕ → Bool) (i : ℕ) (l : Line Sg) {ls : List (Line Sg)} :
    swapLines (fun j => if j = i then b else sw j) i (l :: ls) = swapIf b l :: swapLines sw (i + 1) ls := by
  simp only [swapLines, ↓reduceIte]
  rw [swapLines_congr ls fun j hj => if_neg (by omega)]

/-- Every swap is `swapLines sw` for some selection `sw`. -/
theorem exists_swapLines {ls ls' : List (Line Sg)} (hs : List.Forall₂ SwapLine ls ls') :
    ∀ i, ∃ sw : ℕ → Bool, swapLines sw i ls = ls' := by
  induction hs with
  | nil => exact fun _ => ⟨fun _ => false, rfl⟩
  | cons hl _ ih =>
    intro i
    obtain ⟨sw, hsw⟩ := ih (i + 1)
    cases hl with
    | refl l =>
      exact ⟨_, by rw [swapLines_upd false sw i _, hsw]; simp [swapIf]⟩
    | swap o hs a b =>
      exact ⟨_, by rw [swapLines_upd true sw i _, hsw]; simp [swapIf, hs]⟩

variable {P : Prog Sg} {ρ : List Role}

/-- Completeness for the swap selected by `sw`. -/
theorem swapArgs_complete_sw (sw : ℕ → Bool) (hT : (P.swap (swapLines sw 0 P.lines)).Typed ρ) :
    (swapArgs P).Typed ρ := by
  set P' := P.swap (swapLines sw 0 P.lines)
  have hP : P.WF := Prog.wf_of_swap (swapLines_swap sw 0 P.lines) hT.wf
  have hΛ : ∀ v ∈ P.inputs, P.nw ≤ v ∧ v < P.ctx₀.length := by
    intro v hv; simp only [Prog.inputs, Finset.mem_Ico] at hv; simp only [List.length_append]; omega
  have H := recs_ok sw P.nw P.lines hT hP.1 rfl hΛ
    (by show P.wts.length ≤ (P.wts ++ P.ins).length; simp)
  set A := recs sw 0 P.ctx₀.length P.lines ρ
  set E := eqns 0 P.ctx₀.length P.lines
  set K₀ := insertAll ∅ (List.range P.wts.length ++ P.ret)
  have hK₀ : ∀ v, v ∈ K₀ ↔ v < P.nw ∨ v ∈ P.ret := by
    intro v; simp only [K₀, mem_insertAll, List.mem_append, List.mem_range]; simp
  have hyp : Hyp A E K₀ := ⟨H.uniq, H.back, H.sub, fun r hr v hv hv₀ => by
    obtain ⟨-, h₁, h₂⟩ := H.binds r hr v hv
    rcases (hK₀ v).1 hv₀ with h | h
    · omega
    · exact h₂ h, H.ne⟩
  have hW : ∀ S : Set ℕ, (∀ v ∈ K₀, v ∈ S) →
      (∀ r ∈ A, (∀ w ∈ r.reads, w ∈ S) → ∀ v ∈ r.binds, v ∈ S) → ∀ r ∈ A, ∀ v ∈ r.binds, v ∈ S := by
    intro S hS₀ hcl
    have key : ∀ v, Known (P'.xrecs ρ) ({w | w < P.nw} ∪ ↑P.ret.toFinset) v → v ∈ S := by
      intro v hk
      induction hk with
      | base hb =>
        refine hS₀ _ ((hK₀ _).2 ?_)
        simpa using hb
      | fire hx _ hv ih =>
        obtain ⟨a, ha, h₁, h₂⟩ := xrecs_recs sw P.lines _ hx
        exact hcl a ha (fun w hw => ih w (h₁ w hw)) _ (h₂ _ hv)
    intro r hr v hv
    obtain ⟨x, hx, hvx⟩ := recs_xrecs sw P.lines r hr
    have hlt := ((P'.wf_xrecs hT.wf) x hx).2.lt (hvx v hv)
    rw [show P'.ctx = ctxAfter P'.ctx₀ P'.lines from rfl, length_ctxAfter] at hlt
    exact key v (hT.known_xrecs (P'.recov_init (P'.xrecs ρ)) (fun _ h => h) v hlt)
  have hex := fwdKnown_exact hyp hW
  have hS : ∀ j, 0 ≤ j → ((swapSet P).contains j = true ↔
      ∃ q ∈ A, q.swap = true ∧ q.line = j) := by
    intro j _
    simp only [swapSet, Std.HashSet.contains_ofList, List.contains_iff_mem, List.mem_map,
      List.mem_filter]
    exact ⟨fun ⟨q, ⟨hq, h₁⟩, h₂⟩ => ⟨q, (hex q).1 hq, h₁, h₂⟩,
      fun ⟨q, hq, h₁, h₂⟩ => ⟨q, ⟨(hex q).2 hq, h₁⟩, h₂⟩⟩
  exact typed_congr hT (H.agree _ hS)

/-- **The front end is complete**: if some swap of `P` is typed, `swapArgs P` is typed with the
same roles. -/
theorem swapArgs_complete {ls : List (Line Sg)} (hs : List.Forall₂ SwapLine P.lines ls)
    (hT : (P.swap ls).Typed ρ) : (swapArgs P).Typed ρ := by
  obtain ⟨sw, rfl⟩ := exists_swapLines hs 0
  exact swapArgs_complete_sw sw hT

/-- **Inference after the front end is exact up to swaps.** -/
theorem infer_swapArgs_iff [DecidableEq S] (P : Prog Sg) (ρ : List Role) :
    (swapArgs P).infer = some ρ ↔ ∃ ls, List.Forall₂ SwapLine P.lines ls ∧ (P.swap ls).Typed ρ := by
  rw [Prog.infer_iff]
  exact ⟨fun h => ⟨_, swapArgs_swap P, h⟩, fun ⟨_, hs, hT⟩ => swapArgs_complete hs hT⟩

/-- All typed swaps of a program have the same roles. -/
theorem swap_roles_unique {ls₁ ls₂ : List (Line Sg)} {ρ₁ ρ₂ : List Role}
    (hs₁ : List.Forall₂ SwapLine P.lines ls₁) (hT₁ : (P.swap ls₁).Typed ρ₁)
    (hs₂ : List.Forall₂ SwapLine P.lines ls₂) (hT₂ : (P.swap ls₂).Typed ρ₂) : ρ₁ = ρ₂ :=
  (swapArgs_complete hs₁ hT₁).unique (swapArgs_complete hs₂ hT₂)

end RoleFlow