import RoleFlow.Typing

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

Inference is a single backward pass (`backward`): starting from the returned variables `N`, a
line is a spine step if one of its results is in `N`, and then its consumed arguments are added
to `N`; otherwise it is a conditioner. `Prog.infer` runs the pass and checks the result.

Inference is exact (`Prog.infer_iff`): it finds the roles of every typing, so typings are
unique (`Prog.Typed.unique`). The reason (`backward_typed`) is that at every node of a
derivation the needed set below the scope is exactly the live set: a live variable is either
returned or consumed by exactly one later spine step, and a conditioner's results are never
live. Nothing here uses measure theory. -/

namespace RoleFlow

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

/-- The backward pass over the lines `ls` at scope `n`, from the needed set `N₀` at the end:
the roles, and the needed set at the start. A line is a spine step if one of its results is
needed and it has a semi-inverse, and a conditioner otherwise. A needed result of a conditioner
is never live, so the typing then fails. -/
def backward : ℕ → List (Line Sg) → Finset ℕ → List Role × Finset ℕ
  | _, [], N => ([], N)
  | n, l :: ls, N₀ =>
    let r := backward (n + l.bs.length) ls N₀
    match (l.ys n).any (· ∈ r.2), Sg.inv l.op with
    | true, some si => (.spine :: r.1, r.2 ∪ (l.args.take si.k).toFinset)
    | _, _ => (.cond :: r.1, r.2)

/-- Infer the roles of a program: the backward pass, confirmed by deciding the typing. -/
def Prog.infer [DecidableEq S] (P : Prog Sg) : Option (List Role) :=
  let ρ := (backward P.ctx₀.length P.lines P.ret.toFinset).1
  if P.Typed ρ then some ρ else none

/-- **The backward pass finds the roles of a derivation**, and below the scope its needed set
is the live set. -/
theorem backward_typed {O Γ : List S} {Λ : Finset ℕ} {D : List (Finset ℕ)} {ls : List (Line Sg)}
    {xs : List ℕ} {ρ : List Role} (hT : Typed Sg O Γ Λ D ls xs ρ) (hΛ : ∀ v ∈ Λ, v < Γ.length) :
    (backward Γ.length ls xs.toFinset).1 = ρ ∧
      ∀ v < Γ.length, v ∈ (backward Γ.length ls xs.toFinset).2 ↔ v ∈ Λ := by
  induction ls generalizing Γ Λ D ρ with
  | nil => cases ρ with
    | nil => exact ⟨rfl, fun v _ => by simp [backward, hT.2.1]⟩
    | cons => exact hT.elim
  | cons l ls ih =>
    match ρ, hT with
    | .cond :: ρ, ⟨_, hT⟩ =>
      obtain ⟨h₁, h₂⟩ := ih hT fun v hv => by have := hΛ v hv; simp; omega
      simp only [List.length_append] at h₁ h₂
      set n := Γ.length
      have hany : (l.ys n).any (· ∈ (backward (n + l.bs.length) ls xs.toFinset).2) = false := by
        refine List.any_eq_false.2 fun y hy hyN => ?_
        have hy' := Line.mem_ys.1 hy
        have := hΛ y ((h₂ y hy'.2).1 (by simpa using hyN))
        omega
      simp only [backward, hany, h₁, true_and]
      exact fun v hv => h₂ v (by omega)
    | .spine :: ρ, ⟨_, si, hsi, _, hcΛ, _, hT⟩ =>
      obtain ⟨h₁, h₂⟩ := ih hT fun v hv => by
        rcases Finset.mem_union.1 hv with hv | hv
        · have := hΛ v (Finset.mem_sdiff.1 hv).1; simp; omega
        · simpa using (Finset.mem_Ico.1 hv).2
      simp only [List.length_append] at h₁ h₂
      set n := Γ.length
      have hbs : 0 < l.bs.length := List.length_pos_iff.2 (Sg.bs_ne l.op)
      have hany : (l.ys n).any (· ∈ (backward (n + l.bs.length) ls xs.toFinset).2) = true := by
        refine List.any_eq_true.2 ⟨n, Line.mem_ys.2 ⟨le_rfl, by omega⟩, ?_⟩
        simpa using (h₂ n (by omega)).2
          (Finset.mem_union_right _ (Finset.mem_Ico.2 ⟨le_rfl, by omega⟩))
      simp only [backward, hany, hsi, h₁, true_and, Finset.mem_union, List.mem_toFinset]
      intro v hv
      rw [h₂ v (by omega)]
      simp only [Finset.mem_union, Finset.mem_sdiff, List.mem_toFinset, Finset.mem_Ico]
      constructor
      · rintro ((⟨h, _⟩ | h) | h)
        · exact h
        · omega
        · exact hcΛ v h
      · intro h
        by_cases hc : v ∈ l.args.take si.k
        · exact Or.inr hc
        · exact Or.inl (Or.inl ⟨h, hc⟩)

/-- **Inference is exact** (`prop:infer`). -/
theorem Prog.infer_iff [DecidableEq S] (P : Prog Sg) (ρ : List Role) :
    P.infer = some ρ ↔ P.Typed ρ := by
  unfold Prog.infer
  constructor
  · intro h
    simp only at h
    split_ifs at h with hc
    cases h
    exact hc
  · intro hT
    have e := (backward_typed hT fun v hv => by
      simp only [Prog.inputs, Finset.mem_Ico] at hv; simp; omega).1
    simp only [e, if_pos hT]

/-- A program has at most one typing. -/
theorem Prog.Typed.unique {P : Prog Sg} {ρ₁ ρ₂ : List Role} (h₁ : P.Typed ρ₁)
    (h₂ : P.Typed ρ₂) : ρ₁ = ρ₂ := by
  classical exact Option.some.inj (((P.infer_iff _).2 h₁).symm.trans ((P.infer_iff _).2 h₂))

end RoleFlow