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

import LeanMLIR.Framework

import LeanMLIR.Framework.Zipper

import LeanMLIR.Transforms.Rewrite.Match

open Ctxt (Var VarSet Valuation)

variable {d} [DialectSignature d] [TyDenote d.Ty] [DialectDenote d] [Monad d.m]

variable [DecidableEq d.Ty] [DecidableEq d.Op]

variable [∀ (t : d.Ty), Inhabited (toType t)]

section SplitProgram

variable (d Γ eff t) in

structure SplitProgramResult extends Zipper d Γ eff t where
  {midTypes : List d.Ty}
  midRet : HVector toZipper.Γ_mid.Var midTypes

def splitProgramAtAux : (pos : ℕ) → (lets : Lets d Γ₁ eff Γ₂) →
    (prog : Com d Γ₂ eff t) →
    Option (SplitProgramResult d Γ₁ eff t)
  | 0, lets, .var e body => some {
      top := lets.var e
      bot := body
      midRet := e.returnVars
    }
  | _, _, .rets _ => none
  | n+1, lets, .var e body =>
    splitProgramAtAux n (lets.var e) body

def splitProgramAt (pos : ℕ) (prog : Com d Γ eff t) :
    Option (SplitProgramResult d Γ eff t) :=
  splitProgramAtAux pos .nil prog

end SplitProgram

def rewriteAt
    (lhs rhs : Com d Γ₁ .pure ts₁)
    (hlhs : ∀ t (v : Var Γ₁ t), ⟨t, v⟩ ∈ lhs.vars)
    (pos : ℕ) (target : Com d Γ₂ eff t₂) :
    Option (Com d Γ₂ eff t₂) := do
  let splitRes ← splitProgramAt pos target
  if h : ts₁ = splitRes.midTypes then
    let m ← matchArgRes splitRes.top lhs.toLets splitRes.midRet (h ▸ lhs.returnVars)
    let m := m.toHom <| by
      subst h; exact hlhs
    let rhs := rhs.changeVars m
    let zip := splitRes.insertPureCom splitRes.midRet (cast (by admit /- proof elided -/
    ) rhs)
    return zip.toCom
  else none

-- Repository definitions used:
content: syntax "neg" : MLIR.Pretty.uniform_op

syntax "llvm.and"     : MLIR.Pretty.uniform_op

syntax "llvm.ashr"    : MLIR.Pretty.exact_op

syntax "llvm.add"     : MLIR.Pretty.overflow_op

syntax "llvm.return"  : MLIR.Pretty.uniform_op

content: notation:50 x " ≤ₛ " y => BitVec.sle x y

content: notation:50 x " >ᵤ " y => BitVec.ult y x

content: notation:50 x " ≥ᵤ " y => BitVec.ule y x

content: notation:50 x " <ᵤ " y => BitVec.ult x y

content: notation:50 x " ≥ₛ " y => BitVec.sle y x

content: notation:50 x " <ₛ " y => BitVec.slt x y

content: notation:50 x " >ₛ " y => BitVec.slt y x

content: notation:50 x " ≤ᵤ " y => BitVec.ule x y

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = const ($x)
      $[: $outer_type]? ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _)
      `(mlir_op| $res:mlir_op_operand = "const"()
          {val = $x:num : $outer_type} : ($outer_type) -> ($outer_type) )

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = li ($x)
     $[: $outer_type]?  ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _ )
      `(mlir_op| $res:mlir_op_operand = "li"()
          {imm = $x:num : $outer_type } : ($outer_type) -> ($outer_type))

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithImmediate $reg1 , $x : $t)  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {imm = $x:num : $t} : ($t) -> ($t) )

content: macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithShamt $reg1 , $x  : $t )  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {shamt = $x:num : $t}  : ($t) -> ($t) )

content: macro_rules
  | `(mlir_op| $resName:mlir_op_operand = $name:InstCombine.cmp_op_name $x, $y $[: $t]?) => do
    let some opName := extractOpName name.raw
      | Macro.throwUnsupported
    let t ← t.getDM `(mlir_type| _)
    `(mlir_op| $resName:mlir_op_operand = $opName ($x, $y) : ($t, $t) -> (i1) )

content: macro_rules
  | `(mlir_op| $resName:mlir_op_operand = $name:InstCombine.int_cast_op $x : $t to $t') => do
    let some opName := extractOpName name.raw
      | Macro.throwUnsupported
    `(mlir_op| $resName:mlir_op_operand = $opName ($x) : ($t) -> $t')

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant( $x $[: $inner_type]?)
      $[: $outer_type]? ) => do
       
      let outer_type ← outer_type.getDM `(mlir_type| _)
      let inner_type := inner_type.getD outer_type
      `(mlir_op| $res:mlir_op_operand = "llvm.mlir.constant"()
          {value = $x:neg_num : $inner_type} : () -> ($outer_type) )
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant( ${ $x:term }) $[: $t]?) => do
      let t ← t.getDM `(mlir_type| _)
      let x ← `(MLIR.AST.AttrValue.int $x [mlir_type| $t])
      `(mlir_op| $res:mlir_op_operand = "llvm.mlir.constant"() {value = $$($x) } : () -> ($t) )

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (true) $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (1 : i1) : i1)
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (false) $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (0 : i1) : i1)
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant $x $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant($x $[: $t]?) $[: $t]?)
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant ${ $x:term } $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant($$($x) $[: $t]?) $[: $t]?)

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.icmp $p $x, $y $[: $t]?) => do
    let t ← t.getDM `(mlir_type| _)
    match p.getString with
      | "eq" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.eq" ($x, $y) : ($t, $t) -> (i1))
      | "ne" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ne" ($x, $y) : ($t, $t) -> (i1))
      | "slt" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.slt" ($x, $y) : ($t, $t) -> (i1))
      | "sle" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.sle" ($x, $y) : ($t, $t) -> (i1))
      | "sgt" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.sgt" ($x, $y) : ($t, $t) -> (i1))
      | "sge" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.sge" ($x, $y) : ($t, $t) -> (i1))
      | "ult" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ult" ($x, $y) : ($t, $t) -> (i1))
      | "ule" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ule" ($x, $y) : ($t, $t) -> (i1))
      | "ugt" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ugt" ($x, $y) : ($t, $t) -> (i1))
      | "uge" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.uge" ($x, $y) : ($t, $t) -> (i1))
      | _ => Macro.throwErrorAt p s!"unexpected predicate {p.getString}"

content: macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.select $c, $x, $y $[: $t]?) => do
    let t ← t.getDM `(mlir_type| _)
    `(mlir_op| $res:mlir_op_operand = "llvm.select" ($c, $x, $y) : (i1, $t, $t) -> ($t))

content: structure Zipper (Γ_in : Ctxt d.Ty) (eff : EffectKind) (tys : List d.Ty) where
   
  {Γ_mid : Ctxt d.Ty}
   
  top : Lets d Γ_in eff Γ_mid
   
  bot : Com d Γ_mid eff tys

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: structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty
  
  deriving Repr, Lean.ToExpr, DecidableEq

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

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

content: inductive Com : Ctxt d.Ty → EffectKind → List d.Ty → Type where
  | rets {Γ} {tys} {eff : EffectKind} (vs : HVector Γ.Var tys) : Com Γ eff tys
  | var (e : Expr Γ eff ty) (body : Com (ty ++ Γ) eff β) : Com Γ eff β

content: structure FlatCom (d : Dialect) [DialectSignature d]  (Γ_in : Ctxt d.Ty) (eff : EffectKind)
    (Γ_out : Ctxt d.Ty) (ts : List d.Ty) where
  lets : Lets d Γ_in eff Γ_out
  rets : HVector Γ_out.Var ts

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: def Var (Γ : Ctxt Ty) (t : Ty) : Type :=
  { i : Nat // Γ[i]? = some t }

content: def matchArgRes (lets : Lets d Γ_in eff Γ_out)
    (matchLets : Lets d Δ_in .pure Δ_out)
    (vs : HVector Γ_out.Var ts)
    (ws : HVector Δ_out.Var ts) :
    Option (MatchArgResult lets matchLets vs ws ∅) := do
  (matchArg lets matchLets vs ws ∅).attach.map fun ⟨⟨_, _⟩, h⟩ => .mk h

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: abbrev MatchVar (Δ Γ : Ctxt d.Ty)  := MatchVarM Δ Γ Unit

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

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

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: @[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: def MatchArgResult := { mapOut : Mapping _ _ //
  ∃ (mapIn' mapOut' : Mapping _ _),
    mapIn.entries ⊆ mapIn'.entries
    ∧ mapOut'.entries ⊆ mapOut.entries
    ∧ matchArg lets matchLets vs ws mapIn' = some ((), mapOut') }

content: def toHom (d : Diff Γ₁ Γ₂) : Hom Γ₁ Γ₂ :=
  fun _ v => ⟨v.val + d.val, d.property v.property⟩

content: abbrev Hom (Γ Γ' : Ctxt Ty) := ⦃t : Ty⦄ → Γ.Var t → Γ'.Var t

content: def Diff (Γ₁ Γ₂ : Ctxt Ty) : Type :=
  {d : Nat // Diff.Valid Γ₁ Γ₂ d}

content: @[simp]
abbrev Diff.Valid (Γ₁ Γ₂ : Ctxt Ty) (d : Nat) : Prop :=
  ∀ {i t}, Γ₁[i]? = some t → Γ₂[i+d]? = some t

content: @[simp] abbrev Expr.pdenoteOp :
    Expr d Γ .pure ty → Γ.Valuation → (HVector toType ty) :=
  Expr.denoteOp

content: def Expr.ty : Expr d Γ eff [t] → d.Ty := fun _ => t

content: def Expr.denoteOp (e : Expr d Γ eff ty) (V : Γ.Valuation) :
    eff.toMonad d.m (HVector toType ty) :=
  EffectKind.liftEffect e.eff_le <| cast (by admit /- proof elided -/
  ) <|
    DialectDenote.denote e.op (e.args.map V) e.regArgs.denote

content: class DialectDenote (d : Dialect) [TyDenote d.Ty] [DialectSignature d] where
  denote : (op : d.Op) → HVector toType (DialectSignature.sig op) →
    (HVector (fun t : Ctxt d.Ty × List d.Ty =>
      t.1.Valuation
      → EffectKind.impure.toMonad d.m (HVector toType t.2))
            (DialectSignature.regSig op)) →
    ((DialectSignature.effectKind op).toMonad d.m
      (HVector toType <| DialectSignature.returnTypes op))

content: def Lets.denote [DialectSignature d] [DialectDenote d] {Γ₂}
    (lets : Lets d Γ₁ eff Γ₂) (V : Valuation Γ₁) : (eff.toMonad d.m <| Valuation Γ₂) :=
  match lets with
  | .nil          => return V
  | .var lets' e  => lets'.denote V >>= e.denote

content: def sig          := Signature.sig ∘ s.signature

content: structure Signature (Ty : Type) where
  mkEffectful ::
  sig         : List Ty
  regSig      : RegionSignature Ty
  returnTypes : List Ty
  effectKind  : EffectKind := .pure

content: class DialectSignature (d : Dialect) where
  signature : d.Op → Signature d.Ty

content: def regSig       := Signature.regSig ∘ s.signature

content: abbrev RegionSignature Ty := List (Ctxt Ty × List Ty)

content: inductive Ty (q : Nat) (n : Nat)
  | index : Ty q n
  | integer : Ty q n
  | tensor : Ty q n
  | polynomialLike : Ty q n
  deriving DecidableEq, Repr

content: def effectKind   := Signature.effectKind ∘ s.signature

content: def returnTypes  := Signature.returnTypes ∘ s.signature

content: structure Dialect where
  (Op : Type)
  (Ty : Type)
  (m : Type → Type := Id)

content: inductive Op (q : Nat) (n : Nat)
  | add : Op q n
  | sub : Op q n
  | mul : Op q n
  | mul_constant : Op q n
    
    
  | leading_term : Op q n
  | monomial : Op q n
  | monomial_mul : Op q n
  | from_tensor : Op q n
  | to_tensor : Op q n
  | const (c : R q n) : Op q n
  | const_int (c : Int) : Op q n
  | const_idx (i : Nat) : Op q n

content: class TyDenote (β : Type) : Type 1 where
  toType : β → Type

content: @[simp]
instance Valuation.instAppendHVector (Γ : Ctxt Ty) (ts : List Ty) :
    HAppend (HVector toType ts) (Valuation Γ) (Valuation <| ⟨ts⟩ ++ Γ) where
  hAppend vals V :=
    (Valuation.ofHVector vals) ++ V

content: def HVector.denote :
    {l : RegionSignature d.Ty} → (T : HVector (fun t => Com d t.1 .impure t.2) l) →
    HVector (fun t => t.1.Valuation → EffectKind.impure.toMonad d.m (HVector toType t.2)) l
  | _, .nil => HVector.nil
  | _, .cons v vs => HVector.cons (v.denote) (HVector.denote vs)

content: @[simp] abbrev FlatCom.denote [DialectDenote d]
    (flatCom : FlatCom d Γ eff Γ_out ts)
    (V : Γ.Valuation) : eff.toMonad d.m (HVector toType ts) :=
  flatCom.lets.denote V >>= (return flatCom.rets.map ·)

content: def RegionSignature.map (f : Ty → Ty') : RegionSignature Ty → RegionSignature Ty' :=
  List.map fun ⟨Γ, ty⟩ => (Γ.map f, ty.map f)

content: def Signature.map (f : Ty → Ty') : Signature Ty → Signature Ty' :=
  fun sig => {
    sig         := sig.sig.map f
    regSig      := sig.regSig.map f
    returnTypes := sig.returnTypes.map f
  }

content: def map (f : Ty₁ → Ty₂) : Ctxt Ty₁ → Ctxt Ty₂ :=
  ofList ∘ (List.map f) ∘ toList

content: def Expr.denote {ty} (e : Expr d Γ eff ty) (V : Valuation Γ) :
    eff.toMonad d.m (e.outContext.Valuation) :=
  match e with
  | ⟨op, ty_eq, heff, args, regArgs⟩ => do
      let argsDenote := args.map V
      let val ← EffectKind.liftEffect heff <| DialectDenote.denote op argsDenote regArgs.denote
      return (val ++ V).cast (by admit /- proof elided -/
      )

content: def Expr.op {Γ : Ctxt d.Ty} {eff : EffectKind} {ty} (e : Expr d Γ eff ty) : d.Op :=
  Expr.casesOn e (fun op _ _ _ _ => op)

content: def Com.outContext {Γ} : Com d Γ eff ts → Ctxt d.Ty :=
  Com.rec' (motive := fun _ => Ctxt d.Ty)
    (@fun Γ _ => Γ) 
    (fun _ _ r => r) 

content: def Com.rec' {Γ} (com : Com d Γ eff t) : motive com :=
   
  Com.rec
    (motive_1 := fun _ _ _ _ => PUnit)
    (motive_2 := fun _ eff' t' c =>
      (h₁ : eff = eff') → (h₂ : t = t') → motive (h₁ ▸ h₂ ▸ c))
    (motive_3 := fun _ _ => PUnit)
    (fun _ _ _ _ _ _ => ⟨⟩) 
    (fun v h₁ h₂ => 
      cast (by admit /- proof elided -/
      ) <| rets (h₂ ▸ v))
    (fun e' body' _ r' h₁ h₂ => 
      let e := h₁ ▸ e'
      let body : Com _ _ eff t := cast (by admit /- proof elided -/
      ) body'
      let r : motive body := cast (by admit /- proof elided -/
      ) (r' h₁ h₂)
      cast (by admit /- proof elided -/
      ) <| var e body r)
    ⟨⟩
    (fun _ _ _ _ => ⟨⟩)
    com
    rfl
    rfl

content: def Valuation.cast {Γ Δ : Ctxt Ty} (h : Γ = Δ) (V : Valuation Γ) : Valuation Δ :=
  fun _ v => V <| v.castCtxt h.symm

content: abbrev Expr.outContext (_ : Expr d Γ eff ts) : Ctxt d.Ty :=
  ts ++ Γ

content: def liftEffect [Pure m] {e1 e2 : EffectKind} {α : Type}
    (hle : e1 ≤ e2) (v1 : e1.toMonad m α) : e2.toMonad m α :=
  match e1, e2, hle with
    | .pure, .pure, _ | .impure, .impure, _ => v1
    | .pure, .impure, _ => Pure.pure v1

content: def toMonad (e : EffectKind) (m : Type → Type) : Type → Type :=
  match e with
  | pure => Id
  | impure => m

content: def Com.denote : Com d Γ eff ty → (Γv : Valuation Γ) →
    eff.toMonad d.m (HVector toType ty)
  | .rets vs, Γv     => pure (vs.map Γv)
  | .var e body, V => e.denote V >>= body.denote

content: def Com.ty : Com d Γ eff [t] → d.Ty := fun _ => t

content: def Expr.args {Γ ts} (e : Expr d Γ eff ts) :
    HVector (Var Γ) (DialectSignature.sig e.op) :=
  Expr.casesOn e (fun _ _ _ args _ => args)

content: def Expr.regArgs {Γ ts} (e : Expr d Γ eff ts) :
    Regions d (DialectSignature.regSig e.op) :=
  Expr.casesOn e (fun _ _ _ _ regArgs => regArgs)

content: abbrev Regions (regSig : RegionSignature d.Ty) : Type :=
  HVector (fun t => Com d t.1 .impure t.2) regSig

content: def FlatCom.denoteLets (flatCom : FlatCom d Γ eff Γ_out t) (Γv : Γ.Valuation) :
    eff.toMonad d.m <| Γ_out.Valuation :=
  flatCom.lets.denote Γv

content: def Com.toLets (com : Com d Γ eff t) : Lets d Γ eff com.outContext :=
  Lets.nil.addComToEnd com

content: def Com.denoteLets : (com : Com d Γ eff ty) → (Γv : Valuation Γ) →
    eff.toMonad d.m (com.outContext.Valuation)
  | .rets _, V => pure V
  | .var e body, V =>
      e.denote V >>= body.denoteLets >>= fun V =>
        return V.cast (by admit /- proof elided -/
        )

content: def cast (h_eq : ty₁ = ty₂) : Γ.Var ty₁ → Γ.Var ty₂
  | ⟨i, h⟩ => ⟨i, h_eq ▸ h⟩

content: def cast (h₁ : Γ = Γ') (h₂ : Δ = Δ') : Diff Γ Δ → Diff Γ' Δ'
  | ⟨n, h⟩ => ⟨n, by admit /- proof elided -/
  ⟩

content: def toCom (zip : Zipper d Γ_in eff ty) : Com d Γ_in eff ty :=
  go zip.top zip.bot
  where
    go : {Γ_mid : _} → Lets d Γ_in eff Γ_mid → Com d Γ_mid eff ty → Com d Γ_in eff ty
      | _, .nil, com          => com
      | _, .var body e, com  => go body (.var e com)

content: def denote (zip : Zipper d Γ_in eff tys) (V_in : Valuation Γ_in) :
    eff.toMonad d.m (HVector toType tys) :=
  (zip.top.denote V_in) >>= zip.bot.denote

content: def Expr.returnVars (e : Expr d Γ eff tys) : HVector e.outContext.Var tys :=
  .ofFn _ _ <| fun i => (Var.ofFin i).appendInl

content: def ofFin (i : Fin Γ.length) : Γ.Var (Γ[i]) :=
  ⟨i.val, by admit /- proof elided -/
  ⟩

content: def Com.returnVars : (com : Com d Γ eff ts) → HVector (Var com.outContext) ts
  | .rets vs => vs
  | .var _ body => body.returnVars

content: def MatchVarResult.toHom
    (map : MatchVarResult lets v matchLets w mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.vars w) :
    Δ_in.Hom Γ_out :=
  map.val.toHom <| map.isTotal_of hvars

content: def MatchVarResult.isTotal_of
    (map : MatchVarResult lets v matchLets w mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.vars w) :
    map.val.IsTotal :=

content: def MatchArgResult.isTotal_of
    (map : MatchArgResult lets matchLets vs ws mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.varsOfVec ws) :
    map.val.IsTotal :=

content: def MatchArgResult.toHom
    (map : MatchArgResult lets matchLets vs ws mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.varsOfVec ws) :
    Δ_in.Hom Γ_out :=
  map.val.toHom <| map.isTotal_of hvars

variable
  {Γ_in Γ_out Δ_in Δ_out : Ctxt d.Ty}
  {lets : Lets d Γ_in eff Γ_out}
  {matchTy}
  {v : Var Γ_out matchTy}
  {matchLets : Lets d Δ_in .pure Δ_out}
  {w : Var Δ_out matchTy}
in

content: def Valuation.comap {Γi Γo : Ctxt Ty} (Γiv: Γi.Valuation) (hom : Ctxt.Hom Γo Γi) : Γo.Valuation :=
  fun _to vo => Γiv (hom vo)

content: def map (f : ∀ (a : α), A a → B a) :
    ∀ {l : List α}, HVector A l → HVector B l
  | [],   .nil        => .nil
  | t::_, .cons a as  => .cons (f t a) (map f as)

content: def MatchVarResult := { mapOut : Mapping _ _ //
  ∃ (mapIn' mapOut' : Mapping _ _),
    mapIn.entries ⊆ mapIn'.entries
    ∧ mapOut'.entries ⊆ mapOut.entries
    ∧ matchVar lets v matchLets w mapIn' = some ((), mapOut') }

content: noncomputable def toArgResult
    (mapOut : MatchVarResult lets v (.var matchLets matchExpr) w.appendInl mapIn) :
    let args := mapOut.getPureExpr_eq_some.choose
    MatchArgResult lets matchLets args matchExpr.args mapIn :=
  ⟨mapOut.1, by admit /- proof elided -/
  ⟩

content: def appendInl (v : Γ.Var t) : (Γ ++ Δ).Var t :=
  ⟨v.val, by admit /- proof elided -/
  ⟩

content: def eqvVarLeft  :
    MatchVarResult lets v (.var matchLets matchExpr) w.appendInr ma
    ≃ MatchVarResult lets v matchLets w ma where
  toFun := fun ⟨x, h⟩ => ⟨x, by admit /- proof elided -/
  ⟩
  invFun := fun ⟨x, h⟩ => ⟨x, by admit /- proof elided -/
  ⟩

content: def Expr.changeVars (varsMap : Γ.Hom Γ') {ty} (e : Expr d Γ eff ty) :
    Expr d Γ' eff ty :=
  ⟨e.op, e.ty_eq, e.eff_le, e.args.map varsMap, e.regArgs⟩

content: def Com.changeVars : Com d Γ eff ty →
    (varsMap : Γ.Hom Γ') →
    Com d Γ' eff ty
  |  .rets e => fun varsMap => .rets (e.map varsMap)
  |  .var e body => fun varsMap => .var (e.changeVars varsMap)
      (body.changeVars (fun _ v => varsMap.append v))

content: def Hom.append {ζ : Ctxt Ty} (f : Γ.Hom Δ) : Hom (ζ ++ Γ) (ζ ++ Δ) :=
  fun _ => Var.appendCases
    (fun v => v.appendInl)
    (fun v => (f v).appendInr)

content: def append (d₁ : Diff Γ₁ Γ₂) (d₂ : Diff Γ₂ Γ₃) : Diff Γ₁ Γ₃ :=
  {val := d₁.val + d₂.val,  property := append_valid d₁.property d₂.property}

content: def insertPureCom (zip : Zipper d Γ_in eff ty)
    (vs : HVector zip.Γ_mid.Var newTy)
    (newCom : Com d zip.Γ_mid .pure newTy) : Zipper d Γ_in eff ty :=
  zip.insertCom vs (newCom.castPureToEff eff)

content: def insertCom (zip : Zipper d Γ_in eff ty)
    (vs : HVector zip.Γ_mid.Var newTy) (newCom : Com d zip.Γ_mid eff newTy) :
    Zipper d Γ_in eff ty :=
  let top := zip.top.addComToEnd newCom
  
  let bot := zip.bot.changeVars <| newCom.outContextHom.with vs newCom.returnVars
  
  
  { top, bot }

content: def Hom.with [DecidableEq Ty] {Γ₁ Γ₂ : Ctxt Ty} (f : Γ₁.Hom Γ₂) {ts}
    (v₁ : HVector Γ₁.Var ts) (v₂ : HVector Γ₂.Var ts) : Γ₁.Hom Γ₂ :=
  fun _ w =>
    match v₁.idxOf? w with
    | none => f w
    | some ⟨i, h⟩ => (v₂.get i).cast h

content: def Com.outContextHom (com : Com d Γ eff t) : Γ.Hom com.outContext :=
  com.outContextDiff.toHom

content: def Com.outContextDiff (com : Com d Γ eff ts) : Γ.Diff com.outContext :=
  ⟨com.bvars, by admit /- proof elided -/
      ⟩

content: @[simp, grind=] def Expr.bvars (e : Expr d Γ eff Δ) : Nat :=
  (DialectSignature.returnTypes e.op).length

content: def Com.bvars : Com d Γ eff t → Nat :=
  Com.rec'
           (fun _ => 0)
      (fun e _body bodySize => e.bvars + bodySize)

content: def castCtxt (h_eq : Γ = Δ) : Γ.Var ty → Δ.Var ty
  | ⟨i, h⟩ => ⟨i, h_eq ▸ h⟩

content: def Com.castPureToEff (eff : EffectKind) : Com d Γ .pure t → Com d Γ eff t :=
  changeEffect (EffectKind.pure_le eff)

content: def Com.changeEffect {eff₁ eff₂ : EffectKind} (h : eff₁ ≤ eff₂) :
    Com d Γ eff₁ t → Com d Γ eff₂ t := fun com =>
  Com.rec' (motive := @fun Γ _ => eff₁ ≤ eff₂ → Com d Γ eff₂ t)
            (fun v _h               => rets v)
      (fun e _body castBody h => var (e.changeEffect h) (castBody h))
    com h

content: def Expr.changeEffect {eff₁ eff₂ : EffectKind} (h : eff₁ ≤ eff₂) :
    Expr d Γ eff₁ t → Expr d Γ eff₂ t
  | Expr.mk op ty_eq eff_le args regArgs =>
    have heff : DialectSignature.effectKind op ≤ eff₂ := by admit /- proof elided -/

content: def com := mkCom (d := InstCombine.MetaLLVM 0) bb0 |>.toOption |>.get (by admit /- proof elided -/
)

content: def bb0 : Region 0 := [mlir_region|
{
  ^bb0(%arg0: i32):
    %0 = llvm.mlir.constant(8) : i32
    %1 = llvm.mlir.constant(31) : i32
    %2 = llvm.ashr %arg0, %1 : i32
    %3 = llvm.and %2, %0 : i32
    %4 = llvm.add %3, %2 : i32
    llvm.return %4 : i32
  }]

content: structure Region where
  (name: String)
  (args: List <| TypedSSAVal φ)
  (ops: List Op)

content: abbrev MetaLLVM (φ : Nat) : Dialect where
  Op := MOp φ
  Ty := MTy φ

content: @[deprecated "Use `LLVM.Ty` instead" (since:="2025-04-30")] abbrev Ty := LLVM.Ty

content: @[deprecated "Use `LLVM.Op` instead" (since:="2025-04-30")] abbrev Op := LLVM.Op

content: inductive MOp (φ : Nat) : Type
  | unary   (w : Width φ) (op : MOp.UnaryOp φ) :  MOp φ
  | binary  (w : Width φ) (op : MOp.BinaryOp) :  MOp φ
  | select  (w : Width φ) : MOp φ
  | icmp    (c : IntPred) (w : Width φ) : MOp φ
   
  | const (w : Width φ) (val : ℤ) : MOp φ
deriving Repr, DecidableEq, Inhabited, Lean.ToExpr

content: @[match_pattern] abbrev binary (w : Nat) (op : MOp.BinaryOp) : LLVM.Op :=
  MOp.binary (.concrete w) op

content: inductive MOp.BinaryOp : Type
  | and
  | or   (disjoint : DisjointFlag := {disjoint := false} )
  | xor
  | shl  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | lshr (exact : ExactFlag := {exact := false} )
  | ashr (exact : ExactFlag := {exact := false} )
  | urem
  | srem
  | add  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | mul  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | sub  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | sdiv (exact : ExactFlag := {exact := false} )
  | udiv (exact : ExactFlag := {exact := false} )
deriving DecidableEq, Inhabited, Lean.ToExpr

content: def LLVM : Dialect where
  Op := MOp 0
  Ty := MTy 0

content: inductive MTy (φ : Nat)
  | bitvec (w : Width φ) : MTy φ
  deriving DecidableEq, Inhabited, Lean.ToExpr

content: abbrev Width φ := ConcreteOrMVar Nat φ

content: inductive ConcreteOrMVar (α : Type u) (φ : Nat)
  | concrete (a : α)
  | mvar (i : Fin φ)
  deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

content: @[match_pattern] abbrev const (w : Nat) (val : ℤ) : LLVM.Op        := MOp.const (.concrete w) val

content: inductive MOp.UnaryOp (φ : Nat) : Type
  | neg
  | not
  | copy
  | freeze
  | trunc (w' : Width φ) (noWrapFlags : NoWrapFlags := {nsw := false, nuw := false} )
  | zext  (w' : Width φ) (nneg : NonNegFlag := {nneg := false} )
  | sext  (w' : Width φ)
deriving Repr, DecidableEq, Inhabited, Lean.ToExpr

content: @[simp_llvm_option]
def select {w : Nat} (c? : IntW 1) (x? y? : IntW w ) : IntW w := do
  let c ← c?
  if c = 1#1 then x? else y?

content: def IntW w := PoisonOr <| BitVec w

content: structure PoisonOr (α : Type) where
  val : α
  poisonous : Bool
deriving Inhabited, DecidableEq

content: @[simp_llvm_option]
def icmp {w : Nat} (c : IntPred) (x y : IntW w) : IntW 1 := do
  let x' ← x
  let y' ← y
  icmp? c x' y'

content: @[simp_llvm]
def icmp? {w : Nat} (c : IntPred) (x y : BitVec w) : IntW 1 :=
  .value ↑(icmp' c x y)

content: inductive IntPred where
  | eq
  | ne
  | ugt
  | uge
  | ult
  | ule
  | sgt
  | sge
  | slt
  | sle
deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

content: @[simp_llvm]
def icmp' {w : Nat} (c : IntPred) (x y : BitVec w) : Bool :=
  match c with
    | .eq => (x == y)
    | .ne => (x != y)
    | .sgt => (x >ₛ y)
    | .sge => (x ≥ₛ y)
    | .slt => (x <ₛ y)
    | .sle => (x ≤ₛ y)
    | .ugt => (x >ᵤ y)
    | .uge => (x ≥ᵤ y)
    | .ult => (x <ᵤ y)
    | .ule => (x ≤ᵤ y)

content: def mkCom [TransformTy d φ] [TransformExpr d φ] [TransformReturn d φ]
  (reg : MLIR.AST.Region φ) :
  ExceptM d (Σ (Γ : Ctxt d.Ty) (eff : EffectKind) (ty : _), Com d Γ eff ty) :=
  match reg.ops with
  | [] => throw <| .generic "Ill-formed region (empty)"
  | coms => BuilderM.runWithEmptyMapping <| do
    let Γ ← declareBindings ∅ reg.args
    let com ← mkComHelper Γ coms
    return ⟨Γ, com⟩

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

content: notation "⟦" x "⟧" => TyDenote.toType x

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

-- Repository lemmas used:
content: theorem Com.denoteLets_eq {com : Com d Γ eff t} : com.denoteLets = com.toLets.denote

content: @[simp] theorem Lets.denote_var {lets : Lets d Γ_in eff Γ_out} {e : Expr d Γ_out eff t} :
    (lets.var e).denote = fun V_in => lets.denote V_in >>= e.denote

content: @[simp] theorem denote_toCom [LawfulMonad d.m] (zip : Zipper d Γ_in eff ty) :
    zip.toCom.denote = zip.denote

content: @[simp] theorem Com.denoteLets_returnVars (c : Com d Γ .pure tys) (V : Valuation Γ) :
    c.returnVars.map (c.denoteLets V) = c.denote V

content: theorem Id.bind_eq' (x : Id α) (f : α → id β) : x >>= f = f x

content: theorem Id.pure_eq' (a : α) : (pure a : Id α) = a

content: theorem denote_matchLets_of
    (map : MatchArgResult lets matchLets vs ws mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.varsOfVec ws)
    (V : lets.ValidDenotation) :
    ws.map (matchLets.denote (V.val.comap <| map.toHom h)) = vs.map V.val

content: theorem denote_matchArg
    {vs ws : HVector (Var _) ts}
    (mapOut : MatchArgResult lets matchLets vs ws mapIn)
    (V : lets.ValidDenotation) :
    HVector.map (matchLets.denote (mapOut.val.mapValuation V.val)) ws = HVector.map (V.val) vs

content: theorem HVector.map_eq_map_of_matchArg
    {lets : Lets d Γ_in eff Γ_out}
    {matchLets : Lets d Δ_in .pure Δ_out}
    {ma : Mapping Δ_in Γ_out}
    {l : List d.Ty} {args₁ : HVector _ l} {args₂ : HVector _ l}
    (mapOut : MatchArgResult lets matchLets args₁ args₂ ma)
    (f₁ f₂ : (t : d.Ty) → Var _ t → ⟦t⟧)
    (hf : ∀ {t v₁ v₂},
      (mapOut' : MatchVarResult lets v₁ matchLets v₂ ma)
      → mapOut'.val = mapOut.val
      → f₂ t v₂ = f₁ t v₁) :
    HVector.map f₂ args₂ = HVector.map f₁ args₁

content: theorem denote_matchVar
    {v w : Var _ t}
    (mapOut : MatchVarResult lets v matchLets w mapIn)
    (V : lets.ValidDenotation) :
    (matchLets.denote (mapOut.val.mapValuation V.val) w)
    = V.val v

content: theorem getPureExpr_eq_some
    (mapOut : MatchVarResult lets v (.var matchLets matchExpr) w.appendInl mapIn) :
    ∃ args, lets.getPureExpr v = some ⟨te, w, ⟨
        matchExpr.op,
        matchExpr.ty_eq,
        matchExpr.eff_le,
        args,
        matchExpr.regArgs
      ⟩⟩

content: theorem matchVar_appendInl {w : Var ⟨te⟩ t} :
    matchVar lets v (.var matchLets matchExpr) w.appendInl ma = some ma' →
    ∃ args,
      lets.getPureExpr v
        = some ⟨_, w, matchExpr.op, matchExpr.ty_eq, matchExpr.eff_le, args, matchExpr.regArgs⟩
      ∧ matchArg lets matchLets args matchExpr.args ma = some ma'

content: @[simp]
theorem MatchVar.liftM_bind_eq_some_iff (x? : Option α)
    (f : α → MatchVarM Δ Γ β) :
    ((liftM x? >>= f) mapIn = some mapOut)
    ↔ ( ∃ h : x?.isSome,
        f (x?.get h) mapIn = some mapOut )

content: @[simp] theorem Com.denote_changeVars
    (varsMap : Γ.Hom Γ') (c : Com d Γ eff ty) :
    (c.changeVars varsMap).denote =
    fun V => c.denote (V.comap varsMap)

content: theorem denote_insertPureCom_eq_of [LawfulMonad d.m]
    {zip : Zipper d Γ_in eff tys} {vs}
    {newCom : Com d zip.Γ_mid .pure newTys} {V_in : Valuation Γ_in}
    (h : ∀ V : zip.top.ValidDenotation,
        newCom.denote V.val = vs.map V.val) :
    (zip.insertPureCom vs newCom).denote V_in = zip.denote V_in

content: theorem denote_insertPureCom {zip : Zipper d Γ_in eff t₁} [LawfulMonad d.m]
    {newCom : Com d zip.Γ_mid .pure newTys} {vs : HVector zip.Γ_mid.Var newTys} :
    (zip.insertPureCom vs newCom).denote = (fun (V_in : Valuation Γ_in) => do
      let V_mid ← zip.top.denote V_in
      zip.bot.denote
        ((Com.denoteLets newCom V_mid).comap <| newCom.outContextHom.with vs newCom.returnVars)
      )

content: theorem denote_insertCom {zip : Zipper d Γ_in eff t₁} [LawfulMonad d.m]
    {newCom : Com d zip.Γ_mid eff newTys} {vs : HVector zip.Γ_mid.Var newTys} :
    (zip.insertCom vs newCom).denote = (fun (V_in : Valuation Γ_in) => do
      let V_mid ← zip.top.denote V_in
      let V_newMid ← newCom.denoteLets V_mid
      zip.bot.denote
        (V_newMid.comap <| newCom.outContextHom.with vs newCom.returnVars)
      )

-- Local definitions used:
content: structure SplitProgramResult extends Zipper d Γ eff t where
  {midTypes : List d.Ty}
  midRet : HVector toZipper.Γ_mid.Var midTypes

content: def splitProgramAtAux : (pos : ℕ) → (lets : Lets d Γ₁ eff Γ₂) →
    (prog : Com d Γ₂ eff t) →
    Option (SplitProgramResult d Γ₁ eff t)
  | 0, lets, .var e body => some {
      top := lets.var e
      bot := body
      midRet := e.returnVars
    }
  | _, _, .rets _ => none
  | n+1, lets, .var e body =>
    splitProgramAtAux n (lets.var e) body

content: def splitProgramAt (pos : ℕ) (prog : Com d Γ eff t) :
    Option (SplitProgramResult d Γ eff t) :=
  splitProgramAtAux pos .nil prog

content: def rewriteAt
    (lhs rhs : Com d Γ₁ .pure ts₁)
    (hlhs : ∀ t (v : Var Γ₁ t), ⟨t, v⟩ ∈ lhs.vars)
    (pos : ℕ) (target : Com d Γ₂ eff t₂) :
    Option (Com d Γ₂ eff t₂) := do
  let splitRes ← splitProgramAt pos target
  if h : ts₁ = splitRes.midTypes then
    let m ← matchArgRes splitRes.top lhs.toLets splitRes.midRet (h ▸ lhs.returnVars)
    let m := m.toHom <| by
      subst h; exact hlhs
    let rhs := rhs.changeVars m
    let zip := splitRes.insertPureCom splitRes.midRet (cast (by admit /- proof elided -/
    ) rhs)
    return zip.toCom
  else none

-- Local lemmas used:
content: theorem denote_splitProgramAtAux [LawfulMonad d.m] :
    {pos : ℕ} → {lets : Lets d Γ₁ eff Γ₂} →
    {prog : Com d Γ₂ eff t} →
    {res : _} → (hres : res ∈ splitProgramAtAux pos lets prog) →
    (V : Valuation Γ₁) →
    res.denote V = (lets.denote V) >>= prog.denote
  | 0, lets, .var e body, res, hres, V => by
    obtain rfl

content: @[simp]
theorem denote_splitProgramAt [LawfulMonad d.m] {pos : ℕ} {prog : Com d Γ eff t}
    {res : _} (hres : res ∈ splitProgramAt pos prog) :
    res.denote = prog.denote

-- Target theorem and proof:
theorem denote_rewriteAt [LawfulMonad d.m]
    {lhs rhs : Com d Γ₁ .pure t₁}
    (hl : lhs.denote = rhs.denote)
    {hlhs : ∀ t (v : Var Γ₁ t), ⟨t, v⟩ ∈ lhs.vars}
    {pos : ℕ} {target : Com d Γ₂ eff t₂}
    {rew : Com d Γ₂ eff t₂}
    (hrew : rew ∈ rewriteAt lhs rhs hlhs pos target) :
    rew.denote = target.denote
:= by
  funext V
  simp only [rewriteAt, Option.pure_def, Option.bind_eq_bind, Option.mem_def,
    Option.bind_eq_some_iff, Option.dite_none_right_eq_some, Option.some.injEq] at hrew
  rcases hrew with ⟨res, h_split, rfl, varMap', -, rfl⟩
  dsimp only at varMap'
  simp only [cast_eq, Zipper.denote_toCom, ← denote_splitProgramAt h_split]
  apply Zipper.denote_insertPureCom_eq_of
  intro V
  simp only [Expr.pdenoteOp, Com.denote_changeVars, ← hl]
  rw [← Com.denoteLets_returnVars, Com.denoteLets_eq]
  apply denote_matchLets_of varMap' hlhs