-- LeanMLIR/LeanMLIR/Transforms/Rewrite/Match.lean

import LeanMLIR.Framework

import LeanMLIR.Transforms.Rewrite.Mapping

open Ctxt (Var VarSet Valuation Hom)

variable {d} [DialectSignature d] [DecidableEq d.Ty]

variable {Γ : Ctxt d.Ty} {ty : d.Ty}

abbrev MatchVarM (Δ Γ : Ctxt d.Ty) := (StateT (Mapping Δ Γ) Option)

abbrev MatchVar (Δ Γ : Ctxt d.Ty)  := MatchVarM Δ Γ Unit

def MatchVarM.unifyVars {Δ Γ : Ctxt d.Ty} (v : Δ.Var t) (w : Γ.Var t) : MatchVar Δ Γ :=
  fun ma =>
    match ma.lookup ⟨_, v⟩ with
    | some v =>
      if v = w then
        some ((), ma)
      else
        none
    | none =>
      some ((), AList.insert ⟨_, v⟩ w ma)

open MatchVarM

variable [DecidableEq d.Op]

def matchArg [DecidableEq d.Op]
    (lets : Lets d Γ_in eff Γ_out) (matchLets : Lets d Δ_in .pure Δ_out) :
    {l : List d.Ty} → HVector (Var Γ_out) l → HVector (Var Δ_out) l →
    MatchVar Δ_in Γ_out
  | _, .nil, .nil => return
  | t::l, .cons vₗ vsₗ, .cons vᵣ vsᵣ => do
      matchVar (t := t) lets vₗ matchLets vᵣ
      matchArg lets matchLets vsₗ vsᵣ
  termination_by l => (sizeOf matchLets, l.length + 1)

def matchVar {Γ_in Γ_out Δ_in Δ_out : Ctxt d.Ty} {t : d.Ty} [DecidableEq d.Op]
    (lets : Lets d Γ_in eff Γ_out) (v : Var Γ_out t) :
    (matchLets : Lets d Δ_in .pure Δ_out) →
    (w : Var Δ_out t) →
    MatchVar Δ_in Γ_out
   
  | @Lets.var _ _ _ _ Δ_out ts matchLets matchExpr, w => by admit /- proof elided -/
      | right w =>
        exact matchVar lets v matchLets w
      | left w => exact do
        let ⟨ts', w', ie⟩ ← lets.getPureExpr v
        if hs : ∃ h : ie.op = matchExpr.op, ie.regArgs = (h ▸ matchExpr.regArgs) then
          have hts : Ctxt.ofList ts' = ts := by admit /- proof elided -/

end

section MatchVar

variable [DecidableEq d.Op] {Γ_in Γ_out Δ_in Δ_out t te}
          {lets : Lets d Γ_in eff Γ_out} {v : Var Γ_out t}
          {matchLets : Lets d Δ_in .pure Δ_out}
          {matchExpr : Expr d Δ_out .pure te}

variable (lets v matchLets w) (mapIn : Mapping _ _) in

variable (lets matchLets) {tys} (vs ws : HVector _ tys) (mapIn : Mapping _ _) in

namespace MatchVarResult

variable [TyDenote d.Ty] [∀ (t : d.Ty), Inhabited ⟦t⟧] in

section Left

variable {w : Δ_out.Var t}

variable {mapIn} (mapOut : MatchVarResult lets v (.var matchLets matchExpr) w.appendInr mapIn)

end Left

variable {w : Var ⟨te⟩ _} {mapIn}

end MatchVarResult

end MatchVar

section SubsetEntries

def MatchVar.IsMonotone (f : MatchVar Δ Γ) : Prop :=
    ∀ mapIn, ∀ mapOut ∈ f mapIn,
      mapIn.entries ⊆ mapOut.2.entries

open MatchVar

section UnifyVars

variable {Δ Γ : Ctxt d.Ty} {t} (w : Δ.Var t) (v : Γ.Var t)

end UnifyVars

variable [DecidableEq d.Op]

-- Repository definitions used:
content: abbrev Mapping (Γ Δ : Ctxt Ty) : Type :=
  @AList (Σ t, Var Γ t) (fun x => Var Δ x.1)

content: structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty
  
  deriving Repr, Lean.ToExpr, DecidableEq

content: def Var (Γ : Ctxt Ty) (t : Ty) : Type :=
  { i : Nat // Γ[i]? = some t }

content: inductive Lets (Γ_in : Ctxt d.Ty) (eff : EffectKind) :
    (Γ_out : Ctxt d.Ty) → Type where
  | nil : Lets Γ_in eff Γ_in
  | var (body : Lets Γ_in eff Γ_out) (e : Expr d Γ_out eff t) : Lets Γ_in eff e.outContext

content: protected instance Expr.decidableEq [DecidableEq d.Op] [DecidableEq d.Ty] :
    {Γ : Ctxt d.Ty} → {ty : List d.Ty} → DecidableEq (Expr d Γ eff ty)
  | Γ, _, .mk op₁ rfl eff_le₁ arg₁ regArgs₁, .mk op₂ eq eff_le₂ arg₂ regArgs₂ =>
    if ho : op₁ = op₂ then by
      subst ho
      letI := HVector.decidableEq arg₁ arg₂
      letI := HVector.decidableEqReg regArgs₁ regArgs₂
      exact decidable_of_iff (arg₁ = arg₂ ∧ regArgs₁ = regArgs₂) (by admit /- proof elided -/
      )
    else isFalse (by admit /- proof elided -/
    )

content: def Var (Γ : Context Ty) (ty : Ty) : Type := { v : VarName // Γ.hasType v ty }

content: def Valuation.nil : Ctxt.Valuation (∅ : Ctxt Ty) := fun _ v => v.emptyElim

content: inductive EffectKind
| pure 
| impure 
deriving Repr, DecidableEq, Lean.ToExpr

content: @[elab_as_elim]
def appendCases
    {motive : (Γ ++ Δ).Var t → Sort u}
    (left : (v : Var Γ t) → motive (appendInl v))
    (right : (v : Var Δ t) → motive (appendInr v)) :
    (v : (Γ ++ Δ).Var t) → motive v
  | ⟨idx, h⟩ =>
    if hv : idx < Γ.length then
      left ⟨idx, by admit /- proof elided -/
      ⟩
    else
      let v' : Var _ _ := ⟨idx - Γ.length, by admit /- proof elided -/
      ⟩
      have eq : v'.appendInr = ⟨idx, h⟩ := by admit /- proof elided -/

content: @[match_pattern]
def last (Γ : Ctxt Ty) (t : Ty) : Ctxt.Var (Ctxt.cons t Γ) t :=
  ⟨0, by admit /- proof elided -/
  ⟩

content: def Valuation.cons {Γ : Ctxt Ty} {t : Ty} (x : toType t) (V : Γ.Valuation) :
    (Γ.cons t).Valuation :=

content: protected instance HVector.decidableEqReg [DecidableEq d.Op] [DecidableEq d.Ty] :
    ∀ {l : RegionSignature d.Ty}, DecidableEq (HVector (fun t => Com d t.1 .impure t.2) l)
  | _, .nil, .nil => isTrue rfl
  | _, .cons x₁ v₁, .cons x₂ v₂ =>
    letI := HVector.decidableEqReg v₁ v₂
    letI := Com.decidableEq x₁ x₂
    decidable_of_iff (x₁ = x₂ ∧ v₁ = v₂) (by admit /- proof elided -/
    )

content: infixr:50 "::ₕ" => HVector.cons

content: infixr:67 "::ᵥ" => Valuation.cons

-- Repository lemmas used:
content: @[simp] theorem appendCases_appendInl (v : Γ.Var t) :
    appendCases (motive := motive) left right v.appendInl = (left v)

-- Local definitions used:
content: abbrev MatchVarM (Δ Γ : Ctxt d.Ty) := (StateT (Mapping Δ Γ) Option)

content: abbrev MatchVar (Δ Γ : Ctxt d.Ty)  := MatchVarM Δ Γ Unit

content: def MatchVarM.unifyVars {Δ Γ : Ctxt d.Ty} (v : Δ.Var t) (w : Γ.Var t) : MatchVar Δ Γ :=
  fun ma =>
    match ma.lookup ⟨_, v⟩ with
    | some v =>
      if v = w then
        some ((), ma)
      else
        none
    | none =>
      some ((), AList.insert ⟨_, v⟩ w ma)

content: def matchArg [DecidableEq d.Op]
    (lets : Lets d Γ_in eff Γ_out) (matchLets : Lets d Δ_in .pure Δ_out) :
    {l : List d.Ty} → HVector (Var Γ_out) l → HVector (Var Δ_out) l →
    MatchVar Δ_in Γ_out
  | _, .nil, .nil => return
  | t::l, .cons vₗ vsₗ, .cons vᵣ vsᵣ => do
      matchVar (t := t) lets vₗ matchLets vᵣ
      matchArg lets matchLets vsₗ vsᵣ
  termination_by l => (sizeOf matchLets, l.length + 1)

content: def matchVar {Γ_in Γ_out Δ_in Δ_out : Ctxt d.Ty} {t : d.Ty} [DecidableEq d.Op]
    (lets : Lets d Γ_in eff Γ_out) (v : Var Γ_out t) :
    (matchLets : Lets d Δ_in .pure Δ_out) →
    (w : Var Δ_out t) →
    MatchVar Δ_in Γ_out
   
  | @Lets.var _ _ _ _ Δ_out ts matchLets matchExpr, w => by admit /- proof elided -/
      | right w =>
        exact matchVar lets v matchLets w
      | left w => exact do
        let ⟨ts', w', ie⟩ ← lets.getPureExpr v
        if hs : ∃ h : ie.op = matchExpr.op, ie.regArgs = (h ▸ matchExpr.regArgs) then
          have hts : Ctxt.ofList ts' = ts := by admit /- proof elided -/

content: def MatchVar.IsMonotone (f : MatchVar Δ Γ) : Prop :=
    ∀ mapIn, ∀ mapOut ∈ f mapIn,
      mapIn.entries ⊆ mapOut.2.entries

-- Local lemmas used:
content: @[simp]
theorem unifyVars_eq_some_iff :
    unifyVars w v mapIn = some ((), mapOut)
    ↔ ( mapIn.lookup ⟨t, w⟩ = none ∧ mapIn.insert ⟨t, w⟩ v = mapOut
        ∨ mapIn.lookup ⟨t, w⟩ = v ∧ mapIn = mapOut
      )

content: @[simp]
theorem MatchVar.isMonotone_bind {f : MatchVar Δ Γ} {g : Unit → MatchVar Δ Γ} :
    f.IsMonotone → (g ()).IsMonotone → IsMonotone (f >>= g)

content: @[simp]
theorem MatchVar.isMonotone_bind_liftM {x? : Option α} {g : α → MatchVar Δ Γ} :
    IsMonotone (liftM x? >>= g) ↔ (∀ x ∈ x?, (g x).IsMonotone)

content: @[simp] theorem MatchVar.isMonotone_none : IsMonotone (none : MatchVar Δ Γ)

content: theorem MatchVar.isMonotone_unifyVars  : IsMonotone (unifyVars w v)

-- Target theorem and proof:
theorem isMonotone_matchVarArg_aux (lets : Lets d Γ_in eff Γ_out) :
    (
     ∀  (Δ_out : Ctxt d.Ty)
        (matchLets : Lets d Δ_in EffectKind.pure Δ_out) (l : List d.Ty)
        (argsl : HVector Γ_out.Var l) (argsr : HVector Δ_out.Var l),
        (matchArg lets matchLets argsl argsr).IsMonotone
    )
    ∧ (
      ∀ (Δ_out : Ctxt d.Ty) (t : d.Ty) (v : Γ_out.Var t)
        (matchLets : Lets d Δ_in EffectKind.pure Δ_out)
        (w : Var Δ_out t),
        (matchVar lets v matchLets w).IsMonotone
    )
:= by
  apply matchArg.mutual_induct (d:=d)
  · intro _ _ mapIn ⟨(), mapOut⟩ hvarMap
    obtain rfl : mapIn = mapOut := by
      simp only [matchArg] at hvarMap
      change some ((), _) = some ((), _) at hvarMap
      simp_all
    exact Set.Subset.refl _

  · intros
    simp only [matchArg]
    apply isMonotone_bind <;> assumption

  · intro _ _ Δ_out u matchLets matchExpr l h ih
    cases l using Var.appendCases with
    | right _ => simp [ih]
    | left _ =>
      simp only [matchVar, Var.appendCases_appendInl, isMonotone_bind_liftM,
      Option.mem_def]
      intro ⟨_, e⟩ h_getPureExpr
      split_ifs
      · apply h; assumption
      · exact isMonotone_none
      · exact isMonotone_none

  · simp [isMonotone_unifyVars]