import RoleFlow.Recipes

/-! # Role typing (`docs/roleflow.tex`, `sec:typing`)

The judgement `Typed Sg O Γ Λ D ls xs ρ` has the output sorts `O`, the context `Γ` (its length
is the next free position), the live set `Λ` and the supports `D`; `ls` are the remaining lines,
`xs` the returned positions and `ρ` the roles. Every rule checks the sorts of its line, and
**Ret** those of the returned variables, so a typing includes well-formedness (`Typed.wf`). The support `D.getD v ∅` of `v` stands for its live ancestors: empty for a
weight and `{x}` for an input; a conditioner's results get the union of its arguments'
supports (`condSupp`), and a spine step replaces the consumed arguments in every support that
meets them by its results and its parameters' supports (`spineSupp`).

The invariant `Recov` (every variable is recomputable from the weights and its support) is
preserved from each rule's conclusion to its premise (`Recov.init`, `Recov.cond`,
`Recov.spine`), for a fixed list `R` of recipes that contains those of the lines (`xrecs`).
`Typed.induction_recov` threads `Recov` through a whole derivation.

`Typed` is defined by recursion on the lines, so it is decidable as it stands (`decTyped`,
`Prog.decTyped`, for decidable equality of sorts). Nothing here uses measure theory. -/

namespace RoleFlow

/-! ## The judgement -/

variable {S : Type} {Sg : Sig S}

/-- The live set after a spine step at scope `n` solving the first `k` arguments: the consumed
arguments leave, the results `[n, n + |ys|)` enter. -/
abbrev spineLive (Λ : Finset ℕ) (l : Line Sg) (k n : ℕ) : Finset ℕ :=
  (Λ \ (l.args.take k).toFinset) ∪ Finset.Ico n (n + l.bs.length)

/-- The supports at the start: none for a weight, itself for an input. -/
def initSupp (nw n : ℕ) : List (Finset ℕ) := (List.range n).map fun v => if v < nw then ∅ else {v}

/-- The supports after a conditioner: each result has the union of its arguments' supports. -/
def condSupp (D : List (Finset ℕ)) (l : Line Sg) : List (Finset ℕ) :=
  D ++ List.replicate l.bs.length (l.args.toFinset.biUnion (D.getD · ∅))

/-- The supports after a spine step at scope `n` solving the first `k` arguments: a support that
meets the consumed arguments drops them and gains the results and the parameters' supports;
each result supports itself. -/
def spineSupp (D : List (Finset ℕ)) (l : Line Sg) (k n : ℕ) : List (Finset ℕ) :=
  let c := (l.args.take k).toFinset
  let new := Finset.Ico n (n + l.bs.length) ∪ (l.args.drop k).toFinset.biUnion (D.getD · ∅)
  D.map (fun s => if Disjoint s c then s else s \ c ∪ new) ++
    (l.ys n).map fun y => {y}

/-- `Γ; Λ; D ⊢ e : O` with roles `ρ`, for the lines `ls` and the returned positions `xs`: the
rules **Ret**, **Cond** and **Spine**, by recursion on the lines. Each rule also checks the sorts
of its line, so a typing includes well-formedness (`Typed.wf`). -/
def Typed (Sg : Sig S) (O : List S) :
    List S → Finset ℕ → List (Finset ℕ) → List (Line Sg) → List ℕ → List Role → Prop
  | Γ, Λ, _, [], xs, [] => xs.Nodup ∧ xs.toFinset = Λ ∧ Args Γ xs O
  | Γ, Λ, D, l :: ls, xs, .cond :: ρ => Args Γ l.args l.as ∧ Typed Sg O (Γ ++ l.bs) Λ (condSupp D l) ls xs ρ
  | Γ, Λ, D, l :: ls, xs, .spine :: ρ => Args Γ l.args l.as ∧ ∃ si, Sg.inv l.op = some si ∧
      (l.args.take si.k).Nodup ∧ (∀ c ∈ l.args.take si.k, c ∈ Λ) ∧
      (∀ p ∈ l.args.drop si.k, ∀ c ∈ l.args.take si.k, c ∉ D.getD p ∅) ∧
      Typed Sg O (Γ ++ l.bs) (spineLive Λ l si.k Γ.length) (spineSupp D l si.k Γ.length) ls xs ρ
  | _, _, _, _, _, _ => False

instance decTyped [DecidableEq S] (O : List S) :
    ∀ Γ Λ D (ls : List (Line Sg)) xs ρ, Decidable (Typed Sg O Γ Λ D ls xs ρ)
  | _, _, _, [], _, [] => inferInstanceAs (Decidable (_ ∧ _))
  | _, _, _, [], _, _ :: _ | _, _, _, _ :: _, _, [] => inferInstanceAs (Decidable False)
  | Γ, Λ, D, l :: ls, xs, .cond :: ρ =>
    have := decTyped O (Γ ++ l.bs) Λ (condSupp D l) ls xs ρ
    inferInstanceAs (Decidable (_ ∧ _))
  | Γ, Λ, D, l :: ls, xs, .spine :: ρ =>
    have := fun k => decTyped O (Γ ++ l.bs) (spineLive Λ l k Γ.length) (spineSupp D l k Γ.length) ls xs ρ
    inferInstanceAs (Decidable (_ ∧ ∃ si ∈ Sg.inv l.op, _))

/-- A typing includes well-formedness. -/
theorem Typed.wf {O : List S} : ∀ {Γ Λ D} {ls : List (Line Sg)} {xs ρ},
    Typed Sg O Γ Λ D ls xs ρ → LinesWF Γ ls ∧ Args (ctxAfter Γ ls) xs O
  | _, _, _, [], _, [], h => ⟨trivial, h.2.2⟩
  | _, _, _, _ :: _, _, .cond :: _, h => let ih := h.2.wf; ⟨⟨h.1, ih.1⟩, ih.2⟩
  | _, _, _, _ :: _, _, .spine :: _, ⟨ha, _, _, _, _, _, h⟩ => let ih := h.wf; ⟨⟨ha, ih.1⟩, ih.2⟩

/-- `⊢ P : ρ`: `P` is typed from its inputs. -/
def Prog.Typed (P : Prog Sg) (ρ : List Role) : Prop :=
  RoleFlow.Typed Sg P.outs P.ctx₀ P.inputs (initSupp P.nw P.ctx₀.length) P.lines P.ret ρ

instance Prog.decTyped [DecidableEq S] (P : Prog Sg) (ρ : List Role) : Decidable (P.Typed ρ) :=
  RoleFlow.decTyped ..

theorem Prog.Typed.wf {P : Prog Sg} {ρ : List Role} (h : P.Typed ρ) : P.WF := RoleFlow.Typed.wf h

/-! ## Recoverability -/

/-- `Recov`: there are `n` supports, and every variable in scope is recomputable from the weights
and its support, which is live. -/
def Recov (nw n : ℕ) (Λ : Finset ℕ) (R : List (XRecipe Sg)) (D : List (Finset ℕ)) : Prop :=
  D.length = n ∧ ∀ v < n, D.getD v ∅ ⊆ Λ ∧ Known R ({w | w < nw} ∪ ↑(D.getD v ∅)) v

/-- At the start, with the inputs live. -/
theorem Recov.init (nw n : ℕ) (R : List (XRecipe Sg)) :
    Recov nw n (Finset.Ico nw n) R (initSupp nw n) := by
  refine ⟨by simp [initSupp], fun v hv => ?_⟩
  have e : (initSupp nw n).getD v ∅ = if v < nw then ∅ else {v} := by
    simp [initSupp, List.getD_eq_getElem?_getD, hv]
  rw [e]
  split_ifs with hw
  · exact ⟨Finset.empty_subset _, .base (Or.inl hw)⟩
  · exact ⟨by simp; omega, .base (Or.inr (by simp))⟩

theorem Prog.recov_init (P : Prog Sg) (R : List (XRecipe Sg)) :
    Recov P.nw P.ctx₀.length P.inputs R (initSupp P.nw P.ctx₀.length) := by
  simpa [Prog.inputs] using Recov.init P.nw P.ctx₀.length R

variable {nw n : ℕ} {Λ : Finset ℕ} {R : List (XRecipe Sg)} {D : List (Finset ℕ)} {l : Line Sg}

/-- **Cond** preserves `Recov`: the new variables are settled by their own equation. -/
theorem Recov.cond (hR : Recov nw n Λ R D) (hargs : ∀ x ∈ l.args, x < n)
    (hr : ⟨l, l.ys n⟩ ∈ R) : Recov nw (n + l.bs.length) Λ R (condSupp D l) := by
  obtain ⟨hlen, hR⟩ := hR
  refine ⟨by simp [condSupp, hlen], fun v hv => ?_⟩
  by_cases hvn : v < n
  · have e : (condSupp D l).getD v ∅ = D.getD v ∅ := by
      simp [condSupp, List.getD_eq_getElem?_getD, List.getElem?_append_left, hlen, hvn]
    rw [e]
    exact hR v hvn
  · have e : (condSupp D l).getD v ∅ = l.args.toFinset.biUnion (D.getD · ∅) := by
      have : v - D.length < l.bs.length := by omega
      rw [condSupp, List.getD_eq_getElem?_getD, List.getElem?_append_right (by omega)]
      simp [this]
    rw [e]
    refine ⟨fun x hx => ?_, .fire hr (fun w hw => ?_)
      (Line.mem_ys.2 ⟨by omega, hv⟩)⟩
    · obtain ⟨a, ha, hx⟩ := Finset.mem_biUnion.1 hx
      exact (hR a (hargs a (List.mem_toFinset.1 ha))).1 hx
    · refine (hR w (hargs w hw)).2.mono (Set.union_subset_union_right _ ?_)
      exact Finset.coe_subset.2 fun x hx => Finset.mem_biUnion.2 ⟨w, List.mem_toFinset.2 hw, hx⟩

/-- **Spine** preserves `Recov`: the consumed arguments are settled by the solved equation,
whose sources are the new live variables and parameters, whose supports avoid them. -/
theorem Recov.spine {si : SemiInv Sg.Op l.as l.bs} (hR : Recov nw n Λ R D)
    (hargs : ∀ x ∈ l.args, x < n)
    (hacyc : ∀ p ∈ l.args.drop si.k, ∀ c ∈ l.args.take si.k, c ∉ D.getD p ∅)
    (hr : ⟨⟨si.op, l.ys n ++ l.args.drop si.k⟩, l.args.take si.k⟩ ∈ R) :
    Recov nw (n + l.bs.length) (spineLive Λ l si.k n) R (spineSupp D l si.k n) := by
  obtain ⟨hlen, hR⟩ := hR
  set c := (l.args.take si.k).toFinset
  set Y := Finset.Ico n (n + l.bs.length)
  set P := (l.args.drop si.k).toFinset.biUnion (D.getD · ∅)
  have hP : P ⊆ Λ \ c := fun x hx => by
    obtain ⟨p, hp, hx⟩ := Finset.mem_biUnion.1 hx
    have hp := List.mem_toFinset.1 hp
    exact Finset.mem_sdiff.2 ⟨(hR p (hargs p (List.mem_of_mem_drop hp))).1 hx,
      fun hxc => hacyc p hp x (List.mem_toFinset.1 hxc) hx⟩
  refine ⟨by simp [spineSupp, hlen, Line.ys], fun v hv => ?_⟩
  by_cases hvn : v < n
  · have e : (spineSupp D l si.k n).getD v ∅ =
        if Disjoint (D.getD v ∅) c then D.getD v ∅ else D.getD v ∅ \ c ∪ (Y ∪ P) := by
      simp [spineSupp, List.getD_eq_getElem?_getD, List.getElem?_append_left, hlen, hvn, c, Y, P]
    obtain ⟨hΛ, hK⟩ := hR v hvn
    rw [e]
    split_ifs with hd
    · refine ⟨fun x hx => Finset.mem_union_left _ ?_, hK⟩
      exact Finset.mem_sdiff.2 ⟨hΛ hx, Finset.disjoint_left.1 hd hx⟩
    · refine ⟨Finset.union_subset (fun x hx => Finset.mem_union_left _
        (Finset.sdiff_subset_sdiff hΛ subset_rfl hx)) (Finset.union_subset Finset.subset_union_right
        (hP.trans Finset.subset_union_left)), Known.trans (fun w hw => ?_) hK⟩
      rcases hw with hw | hw
      · exact .base (Or.inl hw)
      by_cases hwc : w ∈ c
      · -- `w` is consumed: fire the solved equation
        refine .fire hr (fun u hu => ?_) (List.mem_toFinset.1 hwc)
        rcases List.mem_append.1 hu with hy | hp
        · exact .base (Or.inr (Finset.mem_union_right _
            (Finset.mem_union_left _ (Finset.mem_Ico.2 (Line.mem_ys.1 hy)))))
        · refine (hR u (hargs u (List.mem_of_mem_drop hp))).2.mono
            (Set.union_subset_union_right _ (Finset.coe_subset.2 fun x hx => ?_))
          exact Finset.mem_union_right _ (Finset.mem_union_right _
            (Finset.mem_biUnion.2 ⟨u, List.mem_toFinset.2 hp, hx⟩))
      · exact .base (Or.inr (Finset.mem_union_left _ (Finset.mem_sdiff.2 ⟨hw, hwc⟩)))
  · have e : (spineSupp D l si.k n).getD v ∅ = {v} := by
      have : v - D.length < l.bs.length := by omega
      rw [spineSupp, List.getD_eq_getElem?_getD, List.getElem?_append_right (by simp; omega)]
      simp [Line.ys, this]
      omega
    rw [e]
    exact ⟨by simp; omega, .base (Or.inr (by simp))⟩

/-! ## Recoverability along a derivation -/

/-- **`Recov` along a derivation**: induction on a typing in which every node may use `Recov`,
given `Recov` at the root and recipes `R` that contain those of the lines. -/
theorem Typed.induction_recov {O : List S}
    {motive : List S → Finset ℕ → List (Line Sg) → List ℕ → List Role → Prop}
    (ret : ∀ {Γ Λ D xs}, Recov nw Γ.length Λ R D → xs.Nodup → xs.toFinset = Λ → Args Γ xs O →
      motive Γ Λ [] xs [])
    (cond : ∀ {Γ Λ D l ls xs ρ}, Recov nw Γ.length Λ R D → Args Γ l.args l.as →
      motive (Γ ++ l.bs) Λ ls xs ρ → motive Γ Λ (l :: ls) xs (.cond :: ρ))
    (spine : ∀ {Γ Λ D l ls xs ρ} (si : SemiInv Sg.Op l.as l.bs), Recov nw Γ.length Λ R D →
      Args Γ l.args l.as → Sg.inv l.op = some si → (l.args.take si.k).Nodup →
      (∀ c ∈ l.args.take si.k, c ∈ Λ) →
      (∀ p ∈ l.args.drop si.k, ∀ c ∈ l.args.take si.k, c ∉ D.getD p ∅) →
      motive (Γ ++ l.bs) (spineLive Λ l si.k Γ.length) ls xs ρ →
      motive Γ Λ (l :: ls) xs (.spine :: ρ))
    {Γ : List S} {Λ : Finset ℕ} {D : List (Finset ℕ)} {ls : List (Line Sg)} {xs : List ℕ}
    {ρ : List Role} (hT : Typed Sg O Γ Λ D ls xs ρ)
    (hR : Recov nw Γ.length Λ R D) (hX : ∀ x ∈ xrecs Γ.length ls ρ, x ∈ R) : motive Γ Λ ls xs ρ := by
  induction ls generalizing Γ Λ D ρ with
  | nil => cases ρ with
    | nil => exact ret hR hT.1 hT.2.1 hT.2.2
    | cons => exact hT.elim
  | cons l ls ih =>
    match ρ, hT with
    | .cond :: _, ⟨ha, hT⟩ =>
      rw [xrecs_cond] at hX
      exact cond hR ha (ih hT (by rw [List.length_append]; exact hR.cond (fun _ => ha.lt) (hX _ (List.mem_cons_self ..)))
        (by rw [List.length_append]; exact fun x hx => hX x (List.mem_cons_of_mem _ hx)))
    | .spine :: _, ⟨ha, si, hsi, h₁, h₂, h₃, hT⟩ =>
      rw [xrecs_spine _ _ _ _ hsi] at hX
      exact spine si hR ha hsi h₁ h₂ h₃ (ih hT
        (by rw [List.length_append]; exact hR.spine (fun _ => ha.lt) h₃ (hX _ (List.mem_cons_self ..)))
        (by rw [List.length_append]; exact fun x hx => hX x (List.mem_cons_of_mem _ hx)))

/-! ## Congruence and the returned variables -/

theorem condSupp_congr {D : List (Finset ℕ)} {l l' : Line Sg} (hb : l'.bs.length = l.bs.length)
    (ha : ∀ x, x ∈ l'.args ↔ x ∈ l.args) : condSupp D l' = condSupp D l := by
  have : l'.args.toFinset = l.args.toFinset := Finset.ext fun x => by simp [ha]
  simp only [condSupp, this, hb]

/-- The returned variables are live or introduced later. -/
theorem typed_final {O Γ : List S} {Λ : Finset ℕ} {D : List (Finset ℕ)} {ls : List (Line Sg)}
    {xs : List ℕ} {ρ : List Role} (hT : Typed Sg O Γ Λ D ls xs ρ) :
    ∀ v ∈ xs, v ∈ Λ ∨ Γ.length ≤ v := by
  induction ls generalizing Γ Λ D ρ with
  | nil =>
    cases ρ with
    | nil => exact fun v hv => Or.inl (hT.2.1 ▸ List.mem_toFinset.2 hv)
    | cons => exact hT.elim
  | cons l ls ih =>
    match ρ, hT with
    | [], hT => exact hT.elim
    | .cond :: _, ⟨_, hT⟩ => exact fun v hv => (ih hT v hv).imp id fun h => by simp at h; omega
    | .spine :: _, ⟨_, si, _, _, _, _, hT⟩ =>
      intro v hv
      rcases ih hT v hv with h | h
      · simp only [spineLive, Finset.mem_union, Finset.mem_sdiff, Finset.mem_Ico] at h
        rcases h with h | h
        · exact Or.inl h.1
        · exact Or.inr h.1
      · exact Or.inr (by simp at h; omega)

/-! ## Hash sets -/

/-- Insert a list into a hash set. -/
def insertAll (H : Std.HashSet ℕ) (xs : List ℕ) : Std.HashSet ℕ := xs.foldl (·.insert ·) H

theorem mem_insertAll {H : Std.HashSet ℕ} {xs : List ℕ} {a : ℕ} :
    a ∈ insertAll H xs ↔ a ∈ xs ∨ a ∈ H := by
  induction xs generalizing H with
  | nil => simp [insertAll]
  | cons x xs ih =>
    simp only [insertAll] at ih ⊢
    rw [List.foldl_cons, ih, Std.HashSet.mem_insert, beq_iff_eq, List.mem_cons]
    tauto

end RoleFlow