-- LeanMLIR/LeanMLIR/Framework/Basic.lean

import LeanMLIR.ErasedContext

import LeanMLIR.HVector

import LeanMLIR.EffectKind

import LeanMLIR.Framework.Dialect

import Mathlib.Data.Finset.Union

open Ctxt (Var VarSet Valuation Hom)

open TyDenote (toType)

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

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

abbrev Signature.mk (sig : List Ty) (regSig : RegionSignature Ty) (returnTypes : List Ty) : Signature Ty :=
 { sig, regSig, returnTypes }

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

namespace DialectSignature

variable {d} [s : DialectSignature d]

def sig          := Signature.sig ∘ s.signature

def regSig       := Signature.regSig ∘ s.signature

def returnTypes  := Signature.returnTypes ∘ s.signature

def effectKind   := Signature.effectKind ∘ s.signature

end DialectSignature

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))

section DataStructures

variable (d : Dialect) [DialectSignature d]

inductive Expr : (Γ : Ctxt d.Ty) → (eff : EffectKind) → (ty : List d.Ty) → Type where
  | mk {Γ} {ty} (op : d.Op)
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) <| DialectSignature.sig op)
     
    (regArgs : HVector (fun t : Ctxt d.Ty × List d.Ty => Com t.1 .impure t.2)
      (DialectSignature.regSig op)) : Expr Γ eff ty

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 β

end

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

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

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

variable {d} [DialectSignature d]

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 -/
    )

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 -/
    )

protected instance Com.decidableEq [DecidableEq d.Op] [DecidableEq d.Ty]
    {Γ : Ctxt d.Ty} {eff : EffectKind} {tys : List d.Ty} : DecidableEq (Com d Γ eff tys)
  | .rets v₁, .rets v₂ => decidable_of_iff (v₁ = v₂) (by admit /- proof elided -/
  )
  | .var (ty := ty₁) e₁ body₁, .var (ty := ty₂) e₂ body₂ =>
    if hα : ty₁ = ty₂
    then by
      subst hα
      letI := Expr.decidableEq e₁ e₂
      letI := Com.decidableEq body₁ body₂
      exact decidable_of_iff (e₁ = e₂ ∧ body₁ = body₂) (by admit /- proof elided -/
      )
    else isFalse (by admit /- proof elided -/
    )
  | .rets _, .var _ _ => isFalse (fun h => Com.noConfusion h)
  | .var _ _, .rets _ => isFalse (fun h => Com.noConfusion h)

end -- decEq

end DataStructures

variable {d : Dialect} [DialectSignature d]

section Rec

variable {eff t} {motive : ∀ {Γ}, Com d Γ eff t → Sort u}
          (rets : ∀ {Γ : Ctxt _} , (v : HVector Γ.Var t) → motive (Com.rets v))
          (var : ∀ {Γ} {u},
            (e : Expr d Γ eff u) → (body : Com d e.outContext eff t) →
              motive body → motive (Com.var e body))

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

variable {rets} {var} {Γ : Ctxt _}

end Rec

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

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

section Lemmas

namespace Com

end Com

end Lemmas

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

abbrev Expr.contextHom (e : Expr d Γ eff ts) : Γ.Hom e.outContext :=
  Hom.id.appendCodomain

section Lemmas

end Lemmas

variable [TyDenote d.Ty] [DialectDenote d] [DecidableEq d.Ty] [Monad d.m] [LawfulMonad d.m]

end

section Unfoldings

open EffectKind (liftEffect)

end Unfoldings

section Lemmas

end Lemmas

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⟩

section Lemmas

variable {Γ Γ' : Ctxt d.Ty} {t} (f : Γ.Hom Γ') (e : Expr d Γ eff t) (V : Γ'.Valuation)

end Lemmas

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

section Lemmas

end Lemmas

section toPureLemmas

variable {Γ eff ty} {e : Expr d Γ eff ty} (h : e.HasPureOp)

end toPureLemmas

section DenoteInsert

end DenoteInsert

def Lets.getPureExprAux {Γ₁ Γ₂ : Ctxt d.Ty} {t} : Lets d Γ₁ eff Γ₂ → (v : Var Γ₂ t) →
    Option (Σ ts, (Var ⟨ts⟩ t) × Expr d (Γ₂.dropUntil v) .pure ts)
  | .nil, _ => none
  | .var (Γ_out := Γ_out) (t := t) lets e, v => by admit /- proof elided -/
    | right v =>
        apply cast ?_ <| Lets.getPureExprAux lets v
        simp
    | left v =>
        have h : (Ctxt.dropUntil t v) ++ Γ_out = e.outContext.dropUntil v.appendInl := by admit /- proof elided -/

def Lets.getPureExpr {Γ₁ Γ₂ : Ctxt d.Ty} (lets : Lets d Γ₁ eff Γ₂) {t : d.Ty} (v : Var Γ₂ t) :
    Option (Σ ts, (Var ⟨ts⟩ t) × Expr d Γ₂ .pure ts) :=
  (getPureExprAux lets v).map fun ⟨_, v, e⟩ =>
    ⟨_, v, e.changeVars Ctxt.dropUntilHom⟩

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

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

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 Valuation.nil : Ctxt.Valuation (∅ : Ctxt Ty) := fun _ v => v.emptyElim

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

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

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

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

content: def dropUntil : Ctxt Ty :=
  ⟨Γ.toList.drop (v.val + 1)⟩

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: def Hom.castCodomain (h : Δ = Δ') (f : Γ.Hom Δ) : Γ.Hom Δ' :=
  fun _t v => (f v).castCtxt h

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

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

content: abbrev dropUntilHom : Hom (Γ.dropUntil v) Γ := dropUntilDiff.toHom

content: def dropUntilDiff : Diff (Γ.dropUntil v) Γ :=
  ⟨v.val+1, by admit /- proof elided -/
  ⟩

content: def emptyElim {α : Sort _} {t : Ty} : Ctxt.Var ∅ t → α :=
  fun ⟨_, h⟩ => by admit /- proof elided -/

content: @[match_pattern]
def cons (hd : Ty) : Ctxt Ty → Ctxt Ty
| ⟨tl⟩ => ⟨hd :: tl⟩

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 Hom.id {Γ : Ctxt Ty} : Γ.Hom Γ :=
  fun _ v => v

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

content: @[grind=]
def length (Γ : Ctxt Ty) : Nat := Γ.toList.length

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: structure HVectorLiteral where
  u : Level
  v : Level
  α : Q(Type $u)
  A : Q($α → Type $v)
  elems : Array ((a : Q($α)) × Q($A $a))

      instance : DialectSignature $dialect where
        signature := fun op => match op with $matchAlts:matchAlts
    )

content: noncomputable instance (q : ℕ) [Fact (q > 1)] : DialectDenote (ModArith q) where
denote
  | .add, arg, _ =>
      
      (fun args : R q × R q => args.1 + args.2) arg.toPair
  | .sub, arg, _ =>
      
      (fun args : R q × R q => args.1 - args.2) arg.toPair
  | .mul, arg, _ =>
      
      (fun args : R q × R q => args.1 * args.2) arg.toPair
  | .const _ c, _, _ =>
      
      c

content: @[simp, reducible]
def Op.signature : Op q n → Signature (Ty q n) :=
  fun o => {sig := Op.sig o, returnTypes := [Op.outTy o], regSig := []}

content: @[simp, reducible]
def Op.sig : Op q n → List (Ty q n)
| Op.add => [Ty.polynomialLike, Ty.polynomialLike]
| Op.sub => [Ty.polynomialLike, Ty.polynomialLike]
| Op.mul => [Ty.polynomialLike, Ty.polynomialLike]
| Op.mul_constant => [Ty.polynomialLike, Ty.integer]
| Op.leading_term => [Ty.polynomialLike]
| Op.monomial => [Ty.integer, Ty.index]
| Op.monomial_mul => [Ty.polynomialLike, Ty.index]
| Op.from_tensor => [Ty.tensor]
| Op.to_tensor => [Ty.polynomialLike]
| Op.const _ => []
| Op.const_int _ => []
| Op.const_idx _ => []

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: @[simp, reducible]
def Op.outTy : Op q n → Ty q n
| Op.add | Op.sub | Op.mul | Op.mul_constant | Op.leading_term | Op.monomial
| Op.monomial_mul | Op.from_tensor | Op.const _  => Ty.polynomialLike
| Op.to_tensor => Ty.tensor
| Op.const_int _ => Ty.integer
| Op.const_idx _ => Ty.index

content: @[reducible, simp]
def Op.regSig : Op → RegionSignature Ty
  | .map2d => [([Ty.int], [.int])]
  | _ => []

content: instance : DialectSignature Ex where
  signature
    | .add    => ⟨[.nat, .nat], [], [.nat], .pure⟩
    | .beq    => ⟨[.nat, .nat], [], [.bool], .pure⟩
    | .cst _  => ⟨[], [], [.nat], .pure⟩

content: inductive ExOp :  Type
  | add : ExOp
  | beq : ExOp
  | cst : ℕ → ExOp
  deriving DecidableEq

content: inductive ExTy
  | nat
  | bool
  deriving DecidableEq

content: def add {Γ : Ctxt _} (e₁ e₂ : Ctxt.Var Γ .nat) : Expr Ex Γ .pure [.nat] :=
  Expr.mk
    (op := .add)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

content: abbrev Ex : Dialect where
  Op := ExOp
  Ty := ExTy

content: def cst {Γ : Ctxt _} (n : ℕ) : Expr Ex Γ .pure [.nat]  :=
  Expr.mk
    (op := .cst n)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

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

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

content: @[reducible]
instance : DialectSignature Tensor2D where
  signature op := { sig := op.sig, regSig := op.regSig, returnTypes := [op.outTy] }

content: instance : DialectSignature Tensor1D where
  signature op := { sig := op.sig, regSig := op.regSig, returnTypes := [op.outTy], effectKind := .pure }

content: instance : DialectSignature RV64 where
  signature o := {sig := Op.sig o, returnTypes := [Op.outTy o], regSig := []}

content: @[simp, reducible]
def Op.sig : Op → List Ty
  | .li _ => []
  | .mulh  => [Ty.bv, Ty.bv]
  | .mulhu  => [Ty.bv, Ty.bv]
  | .mulhsu  => [Ty.bv, Ty.bv]
  | .divu =>  [Ty.bv, Ty.bv]
  | .remuw  => [Ty.bv, Ty.bv]
  | .remu  =>  [Ty.bv, Ty.bv]
  | .addiw (_imm : BitVec 12) => [Ty.bv]
  | .lui (_imm : BitVec 20) => [Ty.bv]
  | .auipc (_imm : BitVec 20)  => [Ty.bv]
  | .slliw (_shamt : BitVec 5)  => [Ty.bv]
  | .srliw (_shamt : BitVec 5) => [Ty.bv]
  | .sraiw (_shamt : BitVec 5) => [Ty.bv]
  | .slli (_shamt : BitVec 6) => [Ty.bv]
  | .srli (_shamt : BitVec 6) => [Ty.bv]
  | .srai (_shamt : BitVec 6) => [Ty.bv]
  | .addw => [Ty.bv, Ty.bv]
  | .subw => [Ty.bv, Ty.bv]
  | .sllw => [Ty.bv, Ty.bv]
  | .srlw => [Ty.bv, Ty.bv]
  | .sraw => [Ty.bv, Ty.bv]
  | .add => [Ty.bv, Ty.bv]
  | .slt => [Ty.bv, Ty.bv]
  | .sltu => [Ty.bv, Ty.bv]
  | .and => [Ty.bv, Ty.bv]
  | .or => [Ty.bv, Ty.bv]
  | .xor => [Ty.bv, Ty.bv]
  | .sll => [Ty.bv, Ty.bv]
  | .srl => [Ty.bv, Ty.bv]
  | .sub => [Ty.bv, Ty.bv]
  | .sra => [Ty.bv, Ty.bv]
  | .remw  => [Ty.bv, Ty.bv]
  | .rem  =>  [Ty.bv, Ty.bv]
  | .mul => [Ty.bv, Ty.bv]
  | .mulw => [Ty.bv, Ty.bv]
  | .div  =>  [Ty.bv, Ty.bv]
  | .divw  =>  [Ty.bv, Ty.bv]
  | .divuw  =>  [Ty.bv, Ty.bv]
  | .addi (_imm : BitVec 12) => [Ty.bv]
  | .slti (_imm : BitVec 12) => [Ty.bv]
  | .sltiu (_imm : BitVec 12) => [Ty.bv]
  | .andi (_imm : BitVec 12) => [Ty.bv]
  | .ori (_imm : BitVec 12) => [Ty.bv]
  | .xori (_imm : BitVec 12) => [Ty.bv]
  | .bclr => [Ty.bv, Ty.bv]
  | .bext => [Ty.bv, Ty.bv]
  | .binv => [Ty.bv, Ty.bv]
  | .bset  => [Ty.bv, Ty.bv]
  | .bclri (_shamt : BitVec 6) => [Ty.bv]
  | .bexti (_shamt : BitVec 6) => [Ty.bv]
  | .binvi (_shamt : BitVec 6) => [Ty.bv]
  | .bseti (_shamt : BitVec 6) => [Ty.bv]
  | .adduw => [Ty.bv, Ty.bv]
  | .sh1adduw => [Ty.bv, Ty.bv]
  | .sh2adduw => [Ty.bv, Ty.bv]
  | .sh3adduw => [Ty.bv, Ty.bv]
  | .sh1add => [Ty.bv, Ty.bv]
  | .sh2add => [Ty.bv, Ty.bv]
  | .sh3add => [Ty.bv, Ty.bv]
  | .slliuw (_shamt : BitVec 6) => [Ty.bv]
  | .andn => [Ty.bv, Ty.bv]
  | .orn => [Ty.bv, Ty.bv]
  | .xnor => [Ty.bv, Ty.bv]
  | .clz
  | .clzw
  | .ctz
  | .ctzw
  | .max => [Ty.bv, Ty.bv]
  | .maxu => [Ty.bv, Ty.bv]
  | .min  => [Ty.bv, Ty.bv]
  | .minu  => [Ty.bv, Ty.bv]
  | .sextb => [Ty.bv]
  | .sexth => [Ty.bv]
  | .zexth => [Ty.bv]
  | .rol => [Ty.bv, Ty.bv]
  | .rolw => [Ty.bv, Ty.bv]
  | .ror => [Ty.bv, Ty.bv]
  | .rori (_shamt : BitVec 6) =>[Ty.bv]
  | .roriw (_shamt : BitVec 5) =>[Ty.bv]
  | .rorw => [Ty.bv, Ty.bv]
  | .pack => [Ty.bv, Ty.bv]
  | .packh => [Ty.bv, Ty.bv]
  | .packw => [Ty.bv, Ty.bv]
  | .mv => [Ty.bv]
  | .not => [Ty.bv]
  | .neg => [Ty.bv]
  | .negw => [Ty.bv]
  | .sextw => [Ty.bv]
  | .zextb => [Ty.bv]
  | .zextw => [Ty.bv]
  | .seqz => [Ty.bv]
  | .snez => [Ty.bv]
  | .sltz => [Ty.bv]
  | .sgtz => [Ty.bv]

content: inductive Op
   
  | li : (val : BitVec 64) → Op
  | lui (imm : BitVec 20)
  | auipc (imm : BitVec 20)
  | addi (imm : BitVec 12)
  | andi (imm : BitVec 12)
  | ori (imm : BitVec 12)
  | xori (imm : BitVec 12)
  | addiw (imm : BitVec 12)
  | add
  | slli (shamt : BitVec 6)
  | sub
  | and
  | or
  | xor
  | sll
  | srl
  | sra
  | addw
  | subw
  | sllw
  | srlw
  | sraw
  | slti (imm : BitVec 12)
  | sltiu (imm : BitVec 12)
  | srli (shamt : BitVec 6)
  | srai (shamt : BitVec 6)
  | slliw (shamt : BitVec 5)
  | srliw (shamt : BitVec 5)
  | sraiw (shamt : BitVec 5)
  | slt
  | sltu
   
  | mul
  | mulw
  | mulh
  | mulhu
  | mulhsu
  | divw
  | divuw
  | div
  | divu
  | remw
  | rem
  | remuw
  | remu
   
   
  | adduw
  | sh1adduw
  | sh2adduw
  | sh3adduw
  | sh1add
  | sh2add
  | sh3add
  | slliuw (shamt : BitVec 6)
   
  | andn
  | orn
  | xnor
  | clz
  | clzw
  | ctz
  | ctzw
  | max
  | maxu
  | min
  | minu
  | sextb
  | sexth
  | zexth
  | rol
  | rolw
  | ror
  | rori (_shamt : BitVec 6)
  | roriw (_shamt : BitVec 5)
  | rorw
   
  | bclr
  | bclri (shamt : BitVec 6)
  | bext
  | bexti (shamt : BitVec 6)
  | binv
  | binvi (shamt : BitVec 6)
  | bset
  | bseti (shamt : BitVec 6)
   
  | pack
  | packh
  | packw
   
  | mv
  | not
  | neg
  | negw
  | sextw
  | zextb
  | zextw
  | seqz
  | snez
  | sltz
  | sgtz
  deriving DecidableEq, Repr, Lean.ToExpr

content: inductive Ty
  | bv : Ty
  deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

content: instance : DialectSignature (FHE q n) := ⟨Op.signature⟩

content: instance : DialectSignature LLVM where
  signature op := ⟨op.sig, [], [op.outTy], .pure⟩

content: instance : DialectSignature HSxComb where
  signature := fun op =>
    match op with
    | .comb o => liftSig (signature o) 
    
    
    | .hs o => MLIR2Handshake.instDialectSignatureHandshake.signature o

content: inductive Op : Type _
  | comb (o : MLIR2Comb.Comb.Op)
  | hs (o : MLIR2Handshake.Handshake.Op)
  deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

content: def liftSig (sig : Signature MLIR2Comb.Ty) : Signature MLIR2Handshake.Ty :=
  Signature.mk (sig.sig.map liftTy) [] (liftTy sig.outTy)

content: def liftTy : MLIR2Comb.Ty → MLIR2Handshake.Ty
| .bitvec w => .stream (.bitvec w)

content: inductive Ty
| stream (ty2 : Ty2) : Ty 
| stream2 (ty2 : Ty2) : Ty 
| stream2token (ty2 : Ty2) : Ty 
deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

content: inductive Ty
| bitvec (w : Nat) : Ty 
deriving DecidableEq, Repr, ToExpr

content: inductive Ty2
  | bitvec (w : Nat) : Ty2
deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

content: def map {α β : Type} (s : Stream α) (f : α → β) : Stream β :=
  fun i => (s i).map f

content: def Stream (β : Type) := Stream' (Option β)

content: instance : DialectSignature DCxComb where
  signature := fun op =>
    match op with
    | .comb o => liftSig (signature o) 
    
    
    | .dc o => MLIR2DC.instDialectSignatureDC.signature o

content: inductive Op : Type _
  | comb (o : MLIR2Comb.Comb.Op)
  | dc (o : MLIR2DC.DC.Op)
  deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

content: def liftSig (sig : Signature MLIR2Comb.Ty) : Signature MLIR2DC.Ty :=
  Signature.mk (sig.sig.map liftTy) [] (liftTy sig.outTy)

content: def liftTy : MLIR2Comb.Ty → MLIR2DC.Ty
| .bitvec w => .valuestream w

content: inductive Ty
| tokenstream : Ty
| tokenstream2 : Ty
| valuestream (w : Nat) : Ty 
| valuestream2 (w : Nat) : Ty 
| valuetokenstream (w : Nat) : Ty 
| variadicvaluetokenstream (w : Nat) : Ty 
deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

content: instance [SIG : DialectSignature d] [DENOTE : DialectDenote d] {Γ : Ctxt d.Ty} {t}
    (com : Com d Γ .pure t) : Inhabited (DCEType com) where
  default :=
    ⟨Γ, Hom.id, com, by admit /- proof elided -/
    ⟩

content: inductive ExTy
  | nat
  | bool
  deriving DecidableEq, Repr

content: def cst {Γ : Ctxt _} (n : ℕ) : Expr Ex Γ .pure [.nat] :=
  Expr.mk
    (op := .cst n)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

content: @[simp]
instance LLVMPlusRiscVSignature : DialectSignature LLVMPlusRiscV where
  signature
  | .llvm llvmOp => .llvm <$> DialectSignature.signature llvmOp
  | .riscv riscvOp => .riscv <$> DialectSignature.signature riscvOp
  | .castRiscv w =>
      {sig := [Ty.riscv .bv], returnTypes := [Ty.llvm (.bitvec w)], regSig := []}
  | .castLLVM w =>
      {sig := [Ty.llvm (.bitvec w)], returnTypes := [Ty.riscv .bv], regSig := []}

content: inductive Op where
  | llvm : LLVM.Op -> Op
  | riscv : RISCV64.RV64.Op -> Op
  | castRiscv : Nat → Op
  | castLLVM : Nat → Op
  deriving DecidableEq, Repr, Lean.ToExpr

content: inductive Ty where
  | llvm : LLVM.Ty -> Ty
  | riscv : RISCV64.RV64.Ty -> Ty
  deriving DecidableEq, Repr, Lean.ToExpr

content: instance : DialectSignature ExOp ExTy where
  signature
    | .add    => ⟨[.nat, .nat], [], .nat, .pure⟩
    | .beq    => ⟨[.nat, .nat], [], .bool, .pure⟩
    | .cst _  => ⟨[], [], .nat, .pure⟩

content: inductive ExOp :  Type
  | add : ExOp
  | beq : ExOp
  | cst : ℕ → ExOp
  deriving DecidableEq, Repr

content: def add {Γ : Ctxt _} (e₁ e₂ : Var Γ .nat) : Expr Γ .nat :=
  Expr.mk
    (op := .add)
    (ty_eq := rfl)
    (eff_le := EffectKind.le_refl _)
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

content: inductive ExTy
  | nat
  deriving DecidableEq, Repr

content: abbrev Expr (Γ) (ty) := _root_.Expr ExOp Γ .pure ty

content: def cst {Γ : Ctxt _} (n : ℕ) : Expr Γ .nat  :=
  Expr.mk
    (op := .cst n)
    (ty_eq := rfl)
    (eff_le := EffectKind.le_refl _)
    (args := .nil)
    (regArgs := .nil)

content: instance : DialectSignature ExOp ExTy where
  signature
  | .add    => ⟨[.nat, .nat], [], .nat, .pure⟩
  | .runK _ => ⟨[.nat], [([.nat], .nat)], .nat, .pure⟩

content: inductive ExOp :  Type
  | add : ExOp
  | runK : ℕ → ExOp
  deriving DecidableEq, Repr

content: def add {Γ : Ctxt _} (e₁ e₂ : Var Γ .nat) : Expr Γ .nat :=
  Expr.mk
    (op := .add)
    (ty_eq := rfl)
    (eff_le := EffectKind.pure_le _)
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

content: @[reducible]
instance [TyDenote d.Ty] [DialectSignature d] [DialectDenote d]
    [B : HasBool d] [N : HasNat d] [I : HasInt d] : DialectSignature (Scf d) where
   signature
   | .coe o => signature (d:=d) o
    | .if t t' => ⟨[B.ty, t], [(⟨[t]⟩, [t']), (⟨[t]⟩, [t'])], [t'], .impure⟩
      
      
      
      
      
    | .for t => ⟨[ I.ty,  I.ty,  N.ty, t], [(⟨[I.ty, t]⟩, [t])], [t], .impure⟩
    | .run t => ⟨[t], [(⟨[t]⟩, [t])], [t], .impure⟩
    | .iterate _k => ⟨[I.ty], [(⟨[I.ty]⟩, [I.ty])], [I.ty], .impure⟩

content: class HasTy (d : Dialect) (DenotedTy : Type) [TyDenote d.Ty] [DialectSignature d] where
    ty : d.Ty
    denote_eq : toType ty = DenotedTy := by admit /- proof elided -/

content: inductive Scf.Op (Op' Ty' : Type) (m') [TyDenote Ty'] [DialectSignature ⟨Op', Ty', m'⟩]
    [DialectDenote ⟨Op', Ty', m'⟩] : Type _
  | coe (o : Op')
  | iterate (k : ℕ) 
  | run (inputty : Ty')  
  | if (inputty retty' : Ty')  
  | for (ty : Ty')
  deriving DecidableEq, Repr

content: @[simp_denote] def iterate {Γ : Ctxt _} (k : Nat) (input : Var Γ Arith.Ty.int)
    (body : Com ScfArith ⟨[.int]⟩ .impure .int) : Expr ScfArith Γ .impure .int :=
  Expr.mk
    (op := .iterate k)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons input .nil)
    (regArgs := HVector.cons body HVector.nil)

content: abbrev ScfArith := Scf Arith

content: def Scf (d : Dialect) [TyDenote d.Ty] [DialectSignature d] [DialectDenote d] : Dialect where
  Op := Scf.Op d.Op d.Ty d.m
  Ty := d.Ty
  m  := d.m

content: inductive Op
  | add : Op  
  | add_nat : Op  
  | axpy : Op  
  | neg : Op  
  | const : (val : ℤ) → Op
  | const_nat : (val : ℕ) → Op

content: @[simp_denote]
def run {Γ : Ctxt _} {t : Arith.Ty} (v : Var Γ t) (body : Com ScfArith ⟨[t]⟩ .impure t) :
    Expr ScfArith Γ .impure t :=
  Expr.mk
    (op := .run t)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons v .nil)
    (regArgs := HVector.cons body <| HVector.nil)

content: inductive Ty
| int
| bool
| nat
 deriving DecidableEq, Repr

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: @[simp_denote] def neg {Γ : Ctxt _} (a : Var Γ .int) : Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .neg)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

content: @[simp_denote] def axpy {Γ : Ctxt _} (a : Var Γ .int) (x : Var Γ .nat) (b: Var Γ .int) :
    Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .axpy)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons x <| .cons b .nil)
    (regArgs := .nil)

content: @[simp_denote] def add_nat (e₁ e₂ : Var Γ .nat) : Expr ScfArith Γ .pure .nat :=
  Expr.mk
    (op := .coe <| .add_nat)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

content: inductive Ty
  | int
   
  | int2
  deriving DecidableEq, Lean.ToExpr

content: inductive Op
  | noop
  | mkPair
  | unPair
  deriving Lean.ToExpr

content: abbrev Arith : Dialect := {Op, Ty}

content: @[reducible]
instance : DialectSignature Arith where
  signature
    | .axpy => ⟨[.int, .nat, .int], [], [.int], .pure⟩
    | .neg => ⟨[.int], [], [.int], .pure⟩
    | .const _ => ⟨[], [], [.int], .pure⟩
    | .const_nat _ => ⟨[], [], [.nat], .pure⟩
    | .add   => ⟨[.int, .int], [], [.int], .pure⟩
    | .add_nat   => ⟨[.nat, .nat], [], [.nat], .pure⟩

content: @[simp_denote] def add {Γ : Ctxt _} (e₁ e₂ : Var Γ .int) : Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .add)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

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

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

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

-- Repository lemmas used:
content: theorem map_map {A B C : α → Type*} {l : List α} (t : HVector A l)
    (f : ∀ a, A a → B a) (g : ∀ a, B a → C a) :
    (t.map f).map g = t.map (fun a v => g a (f a v))

-- Local definitions used:
content: abbrev RegionSignature Ty := List (Ctxt Ty × List Ty)

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

content: abbrev Signature.mk (sig : List Ty) (regSig : RegionSignature Ty) (returnTypes : List Ty) : Signature Ty :=
 { sig, regSig, returnTypes }

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

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

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

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

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

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: inductive Expr : (Γ : Ctxt d.Ty) → (eff : EffectKind) → (ty : List d.Ty) → Type where
  | mk {Γ} {ty} (op : d.Op)
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) <| DialectSignature.sig op)
     
    (regArgs : HVector (fun t : Ctxt d.Ty × List d.Ty => Com t.1 .impure t.2)
      (DialectSignature.regSig op)) : Expr Γ eff ty

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: abbrev Expr.outContext (_ : Expr d Γ eff ts) : Ctxt d.Ty :=
  ts ++ Γ

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

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 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: 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: protected instance Com.decidableEq [DecidableEq d.Op] [DecidableEq d.Ty]
    {Γ : Ctxt d.Ty} {eff : EffectKind} {tys : List d.Ty} : DecidableEq (Com d Γ eff tys)
  | .rets v₁, .rets v₂ => decidable_of_iff (v₁ = v₂) (by admit /- proof elided -/
  )
  | .var (ty := ty₁) e₁ body₁, .var (ty := ty₂) e₂ body₂ =>
    if hα : ty₁ = ty₂
    then by
      subst hα
      letI := Expr.decidableEq e₁ e₂
      letI := Com.decidableEq body₁ body₂
      exact decidable_of_iff (e₁ = e₂ ∧ body₁ = body₂) (by admit /- proof elided -/
      )
    else isFalse (by admit /- proof elided -/
    )
  | .rets _, .var _ _ => isFalse (fun h => Com.noConfusion h)
  | .var _ _, .rets _ => isFalse (fun h => Com.noConfusion h)

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 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: def Com.outContext {Γ} : Com d Γ eff ts → Ctxt d.Ty :=
  Com.rec' (motive := fun _ => Ctxt d.Ty)
    (@fun Γ _ => Γ) 
    (fun _ _ r => r)

content: abbrev Expr.contextHom (e : Expr d Γ eff ts) : Γ.Hom e.outContext :=
  Hom.id.appendCodomain

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: 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: def Lets.getPureExprAux {Γ₁ Γ₂ : Ctxt d.Ty} {t} : Lets d Γ₁ eff Γ₂ → (v : Var Γ₂ t) →
    Option (Σ ts, (Var ⟨ts⟩ t) × Expr d (Γ₂.dropUntil v) .pure ts)
  | .nil, _ => none
  | .var (Γ_out := Γ_out) (t := t) lets e, v => by admit /- proof elided -/
    | right v =>
        apply cast ?_ <| Lets.getPureExprAux lets v
        simp
    | left v =>
        have h : (Ctxt.dropUntil t v) ++ Γ_out = e.outContext.dropUntil v.appendInl := by admit /- proof elided -/

content: def Lets.getPureExpr {Γ₁ Γ₂ : Ctxt d.Ty} (lets : Lets d Γ₁ eff Γ₂) {t : d.Ty} (v : Var Γ₂ t) :
    Option (Σ ts, (Var ⟨ts⟩ t) × Expr d Γ₂ .pure ts) :=
  (getPureExprAux lets v).map fun ⟨_, v, e⟩ =>
    ⟨_, v, e.changeVars Ctxt.dropUntilHom⟩

-- Local lemmas used:
content: @[simp] theorem Expr.changeVars_changeVars (e : Expr d Γ eff ty) (f : Γ.Hom Δ) (g : Δ.Hom Ξ) :
    (e.changeVars f).changeVars g = e.changeVars (f.comp g)

content: theorem Expr.changeVars_castCodomain (e : Expr d Γ eff t)
    (f : Hom Γ Δ) (h : Δ = Δ') :
    e.changeVars (f.castCodomain h) = cast (by simp [h]) (e.changeVars f)

content: @[simp] theorem Lets.getPureExprAux_var_appendInr (lets : Lets d Γ_in eff Γ_out)
    (e : Expr d Γ_out eff ty₁) (v : Var Γ_out ty₂) :
    getPureExprAux (lets.var e) v.appendInr
    = (getPureExprAux lets v).map fun ⟨_, w, e⟩ =>
        ⟨_, w, e.changeVars <| Hom.id.castCodomain (by simp)⟩

-- Target theorem and proof:
@[simp] theorem Lets.getPureExpr_var_appendInr (lets : Lets d Γ_in eff Γ_out) (e : Expr d Γ_out _ ty₁)
    (v : Var Γ_out ty₂):
    getPureExpr (lets.var e) v.appendInr
    = (fun ⟨_, w, e'⟩ => ⟨_, w,  e'.changeVars <| e.contextHom⟩) <$> (getPureExpr lets v)
:= by
  simp only [getPureExpr, getPureExprAux_var_appendInr, Option.map_eq_map, Option.map_map]
  congr 1
  funext ⟨_, e⟩
  simp only [Function.comp_apply, Expr.changeVars_changeVars, Sigma.mk.injEq, heq_eq_eq, true_and]
  congr 2
  funext t v
  apply Subtype.ext
  simp [Hom.castCodomain]
  grind