import RoleFlow.Typing

/-! # Forward search over equations

An equation `binds ← reads` remembers its line and whether it is the swapped solved form of an
`add`. Starting from a known set `K₀`, the search fires
every equation whose reads are known and whose results are unknown, sweep after sweep, until
nothing fires (`fwdKnown`). The search is exact for a list `A` of recipes (`fwdKnown_exact`):
if every variable has at most one recipe, every equation of `E` is a recipe or has a back edge
to one, and every result of `A` is derivable from `K₀`, then the fired equations are exactly
`A`. Nothing here depends on the signature; the front end (`Front`, `FrontComplete`) supplies
the equations of the lines and the swapped solved form of `add`. -/

namespace RoleFlow.Search

/-- An equation `binds ← reads` of the line `line`; `swap` marks the swapped solved form. -/
structure Eqn where
  binds : List ℕ
  reads : List ℕ
  line : ℕ
  swap : Bool

/-- Fire `q` if its reads are known and its results unknown. -/
def fireStep (r : Std.HashSet ℕ × List Eqn) (q : Eqn) : Std.HashSet ℕ × List Eqn :=
  if q.binds.all (!r.1.contains ·) && q.reads.all r.1.contains then (insertAll r.1 q.binds, q :: r.2)
  else r

/-- One sweep: fire, in order, every equation that can. -/
def sweep (E : List Eqn) (K : Std.HashSet ℕ) (F : List Eqn) : Std.HashSet ℕ × List Eqn :=
  E.foldl fireStep (K, F)

/-- Sweep until nothing fires: the known set and the fired equations. Every equation binds
something, so each fires at most once and `E.length + 1` sweeps suffice. -/
def fwdKnown (E : List Eqn) : ℕ → Std.HashSet ℕ → List Eqn → Std.HashSet ℕ × List Eqn
  | 0, K, F => (K, F)
  | fuel + 1, K, F =>
    let r := sweep E K F
    if r.2.length = F.length then r else fwdKnown E fuel r.1 r.2


/-- `q` can fire at `K`: its results are unknown and its reads known. -/
def CanFire (K : Std.HashSet ℕ) (q : Eqn) : Prop := (∀ v ∈ q.binds, v ∉ K) ∧ ∀ w ∈ q.reads, w ∈ K

theorem canFire_iff {K : Std.HashSet ℕ} {q : Eqn} :
    (q.binds.all (!K.contains ·) && q.reads.all K.contains) = true ↔ CanFire K q := by
  simp only [CanFire, Bool.and_eq_true, List.all_eq_true, Bool.not_eq_eq_eq_not, Bool.not_true,
    Std.HashSet.contains_eq_false_iff_not_mem, Std.HashSet.contains_iff_mem]

/-- The known set is closed under `A`: a known variable's recipe reads only known variables. -/
def Closed (A : List Eqn) (K : Std.HashSet ℕ) : Prop :=
  ∀ u ∈ K, ∀ r ∈ A, u ∈ r.binds → ∀ w ∈ r.reads, w ∈ K

/-- `q` reads a variable whose recipe reads one of `q`'s results. -/
def BackEdge (A : List Eqn) (q : Eqn) : Prop :=
  ∃ u ∈ q.reads, ∃ v ∈ q.binds, ∃ r ∈ A, u ∈ r.binds ∧ v ∈ r.reads

/-- Every variable has at most one recipe. -/
def UniqB (A : List Eqn) : Prop :=
  ∀ r₁ ∈ A, ∀ r₂ ∈ A, ∀ v ∈ r₁.binds, v ∈ r₂.binds → r₁ = r₂

/-- Nothing in `E` can fire. -/
def Terminal (E : List Eqn) (K : Std.HashSet ℕ) : Prop := ∀ q ∈ E, ¬CanFire K q

theorem BackEdge.mono {e q : Eqn} {A : List Eqn} (h : BackEdge [e] q) (he : e ∈ A) :
    BackEdge A q := by
  obtain ⟨u, hu, v, hv, r, hr, h₁, h₂⟩ := h
  exact ⟨u, hu, v, hv, r, List.mem_singleton.1 hr ▸ he, h₁, h₂⟩

theorem BackEdge.mono' {A B : List Eqn} {q : Eqn} (h : BackEdge A q) (hAB : ∀ r ∈ A, r ∈ B) :
    BackEdge B q := by
  obtain ⟨u, hu, v, hv, r, hr, h₁, h₂⟩ := h
  exact ⟨u, hu, v, hv, r, hAB r hr, h₁, h₂⟩

variable {A E : List Eqn} {K₀ : Std.HashSet ℕ}

/-- The invariant of the search: the known set `K` is closed, the fired equations `F` are
distinct recipes, and a variable is known iff it is a start variable or a result of `F`. -/
structure Inv (A E : List Eqn) (K₀ K : Std.HashSet ℕ) (F : List Eqn) : Prop where
  closed : Closed A K
  fa : ∀ q ∈ F, q ∈ A
  fe : ∀ q ∈ F, q ∈ E
  nd : F.Nodup
  src : ∀ v ∈ K, v ∈ K₀ ∨ ∃ q ∈ F, v ∈ q.binds
  fk : ∀ q ∈ F, ∀ v ∈ q.binds, v ∈ K
  k0 : ∀ v ∈ K₀, v ∈ K

/-- The hypotheses the typing supplies. -/
structure Hyp (A E : List Eqn) (K₀ : Std.HashSet ℕ) : Prop where
  uniq : UniqB A
  back : ∀ q ∈ E, q ∈ A ∨ BackEdge A q
  sub : ∀ r ∈ A, r ∈ E
  out : ∀ r ∈ A, ∀ v ∈ r.binds, v ∉ K₀
  ne : ∀ r ∈ A, r.binds ≠ []

theorem inv_init (hd : ∀ r ∈ A, ∀ v ∈ r.binds, v ∉ K₀) : Inv A E K₀ K₀ [] :=
  ⟨fun u hu r hr hur => absurd hu (hd r hr u hur), by simp, by simp, List.nodup_nil,
    fun v hv => Or.inl hv, by simp, fun _ h => h⟩

/-- **Only recipes fire**: an equation with a back edge cannot fire at a closed known set. -/
theorem only_A {K : Std.HashSet ℕ} {q : Eqn} (hC : Closed A K) (hq : q ∈ A ∨ BackEdge A q)
    (hf : CanFire K q) : q ∈ A := by
  rcases hq with h | ⟨u, hu, v, hv, r, hr, hur, hvr⟩
  · exact h
  · exact absurd (hC u (hf.2 u hu) r hr hur v hvr) (hf.1 v hv)

theorem inv_fire (H : Hyp A E K₀) {K : Std.HashSet ℕ} {F : List Eqn} (hI : Inv A E K₀ K F)
    {q : Eqn} (hqE : q ∈ E) (hf : CanFire K q) : Inv A E K₀ (insertAll K q.binds) (q :: F) := by
  have hqA := only_A hI.closed (H.back q hqE) hf
  have hqF : q ∉ F := fun h => by
    obtain ⟨v, hv⟩ := List.exists_mem_of_ne_nil _ (H.ne q hqA)
    exact hf.1 v hv (hI.fk q h v hv)
  refine ⟨?_, ?_, ?_, List.nodup_cons.2 ⟨hqF, hI.nd⟩, ?_, ?_, ?_⟩
  · intro u hu r hr hur w hw
    rw [mem_insertAll] at hu ⊢
    rcases hu with h | h
    · cases H.uniq q hqA r hr u h hur
      exact Or.inr (hf.2 w hw)
    · exact Or.inr (hI.closed u h r hr hur w hw)
  · simpa using ⟨hqA, hI.fa⟩
  · simpa using ⟨hqE, hI.fe⟩
  · intro v hv
    rw [mem_insertAll] at hv
    rcases hv with h | h
    · exact Or.inr ⟨q, List.mem_cons_self .., h⟩
    · exact (hI.src v h).imp id fun ⟨p, hp, hvp⟩ => ⟨p, List.mem_cons_of_mem _ hp, hvp⟩
  · intro p hp v hv
    rw [mem_insertAll]
    rcases List.mem_cons.1 hp with rfl | h
    · exact Or.inl hv
    · exact Or.inr (hI.fk p h v hv)
  · exact fun v hv => mem_insertAll.2 (Or.inr (hI.k0 v hv))

/-- A sweep keeps the invariant, never loses fired equations, and if it fires nothing, nothing
could fire. -/
theorem foldl_inv (H : Hyp A E K₀) :
    ∀ (L : List Eqn), (∀ q ∈ L, q ∈ E) → ∀ {K : Std.HashSet ℕ} {F : List Eqn}, Inv A E K₀ K F →
      let r := L.foldl fireStep (K, F)
      Inv A E K₀ r.1 r.2 ∧ F.length ≤ r.2.length ∧
        (r.2.length = F.length → r.1 = K ∧ ∀ q ∈ L, ¬CanFire K q)
  | [], _, _, _, hI => ⟨hI, le_rfl, fun _ => ⟨rfl, by simp⟩⟩
  | q :: L, hL, K, F, hI => by
    have hL' : ∀ p ∈ L, p ∈ E := fun p hp => hL p (List.mem_cons_of_mem _ hp)
    simp only [List.foldl_cons, fireStep]
    split
    · rename_i hc
      obtain ⟨h₁, h₂, -⟩ := foldl_inv H L hL' (inv_fire H hI (hL q (List.mem_cons_self ..))
        (canFire_iff.1 hc))
      refine ⟨h₁, by simp at h₂; omega, fun h => ?_⟩
      simp at h₂; omega
    · rename_i hc
      obtain ⟨h₁, h₂, h₃⟩ := foldl_inv H L hL' hI
      refine ⟨h₁, h₂, fun h => ⟨(h₃ h).1, ?_⟩⟩
      intro p hp
      rcases List.mem_cons.1 hp with rfl | hp
      · exact fun hf => hc (canFire_iff.2 hf)
      · exact (h₃ h).2 p hp

theorem fwdKnown_inv (H : Hyp A E K₀) :
    ∀ (fuel : ℕ) {K : Std.HashSet ℕ} {F : List Eqn}, Inv A E K₀ K F →
      let r := fwdKnown E fuel K F
      Inv A E K₀ r.1 r.2 ∧ (Terminal E r.1 ∨ F.length + fuel ≤ r.2.length)
  | 0, _, _, hI => ⟨hI, Or.inr le_rfl⟩
  | fuel + 1, K, F, hI => by
    obtain ⟨h₁, h₂, h₃⟩ := foldl_inv H E (fun _ h => h) hI
    rw [show List.foldl fireStep (K, F) E = sweep E K F from rfl] at h₁ h₂ h₃
    simp only [fwdKnown]
    split
    · rename_i h
      obtain ⟨e, hT⟩ := h₃ h
      exact ⟨h₁, Or.inl (by rw [e]; exact hT)⟩
    · rename_i h
      obtain ⟨h₄, h₅⟩ := fwdKnown_inv H fuel h₁
      refine ⟨h₄, h₅.imp id fun h' => ?_⟩
      omega

/-- **The search stops at a fixpoint**: the fired equations are distinct members of `E`. -/
theorem fwdKnown_terminal (H : Hyp A E K₀) :
    let r := fwdKnown E (E.length + 1) K₀ []
    Inv A E K₀ r.1 r.2 ∧ Terminal E r.1 := by
  obtain ⟨hI, h⟩ := fwdKnown_inv H (E.length + 1) (inv_init (E := E) H.out)
  refine ⟨hI, h.resolve_right fun hl => ?_⟩
  have := (hI.nd.subperm hI.fe).length_le
  simp at hl; omega

/-- At a fixpoint the known set is closed under firing `A` forwards. -/
theorem fix_forward (H : Hyp A E K₀) {K : Std.HashSet ℕ} {F : List Eqn} (hI : Inv A E K₀ K F)
    (hT : Terminal E K) : ∀ r ∈ A, (∀ w ∈ r.reads, w ∈ K) → ∀ v ∈ r.binds, v ∈ K := by
  intro r hr hreads
  by_cases hk : ∃ v ∈ r.binds, v ∈ K
  · obtain ⟨v, hv, hvK⟩ := hk
    rcases hI.src v hvK with h | ⟨q, hq, hvq⟩
    · exact absurd h (H.out r hr v hv)
    · obtain rfl := H.uniq q (hI.fa q hq) r hr v hvq hv
      exact hI.fk _ hq
  · simp only [not_exists, not_and] at hk
    exact absurd ⟨hk, hreads⟩ (hT r (H.sub r hr))

/-- **Every recipe of `A` fires** once its results are known. -/
theorem fix_fires_all (H : Hyp A E K₀) {K : Std.HashSet ℕ} {F : List Eqn}
    (hI : Inv A E K₀ K F) {r : Eqn} (hr : r ∈ A) (hk : ∀ v ∈ r.binds, v ∈ K) : r ∈ F := by
  obtain ⟨v, hv⟩ := List.exists_mem_of_ne_nil _ (H.ne r hr)
  rcases hI.src v (hk v hv) with h | ⟨q, hq, hvq⟩
  · exact absurd h (H.out r hr v hv)
  · obtain rfl := H.uniq q (hI.fa q hq) r hr v hvq hv; exact hq

/-- **The search fires exactly `A`**, given that every result of `A` lies in every set that
contains `K₀` and is closed under `A`. -/
theorem fwdKnown_exact (H : Hyp A E K₀)
    (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) (q : Eqn) :
    q ∈ (fwdKnown E (E.length + 1) K₀ []).2 ↔ q ∈ A := by
  obtain ⟨hI, hT⟩ := fwdKnown_terminal H
  exact ⟨hI.fa q, fun hq => fix_fires_all H hI hq
    (hW {v | v ∈ (fwdKnown E (E.length + 1) K₀ []).1} hI.k0 (fix_forward H hI hT) q hq)⟩


end RoleFlow.Search