-- 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.op {Γ : Ctxt d.Ty} {eff : EffectKind} {ty} (e : Expr d Γ eff ty) : d.Op :=
  Expr.casesOn e (fun op _ _ _ _ => op)

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) 

section Lemmas

end Lemmas

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

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)

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

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

end

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

section Unfoldings

open EffectKind (liftEffect)

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

end Unfoldings

section Lemmas

end Lemmas

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

@[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 ·)

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

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

section Lemmas

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

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: class TyDenote (β : Type) : Type 1 where
  toType : β → Type

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

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: 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: 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: def map (f : Ty₁ → Ty₂) : Ctxt Ty₁ → Ctxt Ty₂ :=
  ofList ∘ (List.map f) ∘ toList

    @[simp_denote]
    instance : DialectDenote $dialect where
      denote := 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: @[reducible]
instance : DialectDenote Ex where
  denote
    | .cst n, _, _ => n ::ₕ .nil
    | .add, .cons (a : Nat) (.cons b .nil), _ => a + b ::ₕ .nil
    | .beq, .cons (a : Nat) (.cons b .nil), _ => (a == b) ::ₕ .nil

content: inductive ExOp :  Type
  | add : ExOp
  | beq : ExOp
  | cst : ℕ → ExOp
  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: inductive ExTy
  | nat
  | bool
  deriving DecidableEq

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: @[match_pattern]
def cons (hd : Ty) : Ctxt Ty → Ctxt Ty
| ⟨tl⟩ => ⟨hd :: tl⟩

content: @[simp, reducible]
instance : DialectDenote RV64 where
  denote o args _ := [o.denote args]ₕ

content: @[simp, simp_denote]
abbrev Op.denote : (o : RV64.Op) → HVector toType o.sig → ⟦o.outTy⟧
  | .li imm, _  => imm
  | .addiw imm, regs => RV64.addiw imm (regs.getN 0)
  | .lui imm, regs  => RV64.lui imm
  | .auipc imm, regs => RV64.auipc imm (regs.getN 0)
  | .slliw shamt, regs => RV64.slliw shamt (regs.getN 0)
  | .srliw shamt, regs => RV64.srliw shamt (regs.getN 0)
  | .sraiw shamt, regs => RV64.sraiw shamt (regs.getN 0)
  | .slli shamt, regs => RV64.slli shamt (regs.getN 0)
  | .srli shamt, regs => RV64.srli shamt (regs.getN 0)
  | .srai shamt, regs => RV64.srai shamt (regs.getN 0)
  | .addw, regs => RV64.addw (regs.getN 1) (regs.getN 0)
  | .subw, regs => RV64.subw (regs.getN 1) (regs.getN 0)
  | .sllw, regs => RV64.sllw (regs.getN 1) (regs.getN 0)
  | .srlw, regs => RV64.srlw (regs.getN 1) (regs.getN 0)
  | .sraw, regs => RV64.sraw (regs.getN 1) (regs.getN 0)
  | .add, regs => RV64.add (regs.getN 1) (regs.getN 0)
  | .slt, regs => RV64.slt (regs.getN 1) (regs.getN 0)
  | .sltu, regs => RV64.sltu (regs.getN 1) (regs.getN 0)
  | .and, regs => RV64.and (regs.getN 1) (regs.getN 0)
  | .or, regs => RV64.or (regs.getN 1) (regs.getN 0)
  | .xor, regs => RV64.xor (regs.getN 1) (regs.getN 0)
  | .sll, regs => RV64.sll (regs.getN 1) (regs.getN 0)
  | .srl, regs => RV64.srl (regs.getN 1) (regs.getN 0)
  | .sub, regs => RV64.sub (regs.getN 1) (regs.getN 0)
  | .sra, regs => RV64.sra (regs.getN 1) (regs.getN 0)
  | .remw, regs => RV64.remw (regs.getN 1) (regs.getN 0)
  | .remuw, regs => RV64.remuw (regs.getN 1) (regs.getN 0)
  | .rem, regs => RV64.rem (regs.getN 1) (regs.getN 0)
  | .remu, regs => RV64.remu (regs.getN 1) (regs.getN 0)
  | .mulhu,regs => RV64.mulhu (regs.getN 1) (regs.getN 0)
  | .mul ,regs => RV64.mul (regs.getN 1) (regs.getN 0)
  | .mulhsu ,regs => RV64.mulhsu (regs.getN 1) (regs.getN 0)
  | .mulh,regs => RV64.mulh (regs.getN 1) (regs.getN 0)
  | .mulw,  regs => RV64.mulw (regs.getN 1) (regs.getN 0)
  | .div, regs => RV64.div (regs.getN 1) (regs.getN 0)
  | .divu, regs => RV64.divu (regs.getN 1) (regs.getN 0)
  | .divw, regs => RV64.divw (regs.getN 1) (regs.getN 0)
  | .divuw, regs => RV64.divuw (regs.getN 1) (regs.getN 0)
  | .addi imm, reg => RV64.addi imm (reg.getN 0)
  | .slti imm, reg => RV64.slti imm (reg.getN 0)
  | .sltiu imm, reg => RV64.sltiu imm (reg.getN 0)
  | .andi imm, reg => RV64.andi imm (reg.getN 0)
  | .ori imm, reg => RV64.ori imm (reg.getN 0)
  | .xori imm, reg => RV64.xori imm (reg.getN 0)
  | .bclr, regs => RV64.bclr (regs.getN 1) (regs.getN 0)
  | .bext, regs => RV64.bext (regs.getN 1) (regs.getN 0)
  | .binv, regs => RV64.binv (regs.getN 1) (regs.getN 0)
  | .bset, regs => RV64.bset (regs.getN 1) (regs.getN 0)
  | .bclri shamt , reg => RV64.bclri shamt (reg.getN 0)
  | .bexti shamt, reg => RV64.bexti shamt (reg.getN 0)
  | .binvi shamt, reg => RV64.binvi shamt (reg.getN 0)
  | .bseti shamt, reg => RV64.bseti shamt (reg.getN 0)
  | .adduw, regs => RV64.adduw (regs.getN 1) (regs.getN 0)
  | .sh1adduw , regs => RV64.sh1adduw (regs.getN 1) (regs.getN 0)
  | .sh2adduw, regs => RV64.sh2adduw (regs.getN 1) (regs.getN 0)
  | .sh3adduw, regs => RV64.sh3adduw (regs.getN 1) (regs.getN 0)
  | .sh1add, regs => RV64.sh1add (regs.getN 1) (regs.getN 0)
  | .sh2add, regs => RV64.sh2add (regs.getN 1) (regs.getN 0)
  | .sh3add, regs => RV64.sh3add (regs.getN 1) (regs.getN 0)
  | .slliuw shamt, regs => RV64.slliuw shamt (regs.getN 0)
  | .andn, regs => RV64.andn (regs.getN 1) (regs.getN 0)
  | .orn, regs => RV64.orn (regs.getN 1) (regs.getN 0)
  | .xnor, regs => RV64.xnor (regs.getN 1) (regs.getN 0)
  | .clz, regs => RV64.clz (regs.getN 1)
  | .clzw, regs => RV64.clzw (regs.getN 1)
  | .ctz, regs => RV64.ctz (regs.getN 1)
  | .ctzw, regs => RV64.ctzw (regs.getN 1)
  | .max, regs => RV64.max (regs.getN 1) (regs.getN 0)
  | .maxu, regs => RV64.maxu (regs.getN 1) (regs.getN 0)
  | .min, regs => RV64.min (regs.getN 1) (regs.getN 0)
  | .minu, regs => RV64.minu (regs.getN 1) (regs.getN 0)
  | .sextb, reg => RV64.sextb (reg.getN 0)
  | .sexth, reg => RV64.sexth (reg.getN 0)
  | .zexth, reg => RV64.zexth (reg.getN 0)
  | .rol, regs => RV64.rol (regs.getN 1) (regs.getN 0)
  | .rolw, regs => RV64.rolw (regs.getN 1) (regs.getN 0)
  | .ror, regs => RV64.ror (regs.getN 1) (regs.getN 0)
  | .rori shamt, regs => RV64.rori shamt (regs.getN 0)
  | .roriw shamt, regs => RV64.roriw shamt (regs.getN 0)
  | .rorw, regs => RV64.rorw (regs.getN 1) (regs.getN 0)
  | .pack, regs => RV64.pack (regs.getN 1) (regs.getN 0)
  | .packh, regs => RV64.packh (regs.getN 1) (regs.getN 0)
  | .packw, regs => RV64.packw (regs.getN 1) (regs.getN 0)
  
  | .mv, regs  => RV64.mv_pseudo (regs.getN 0)
  | .not, regs => RV64.not_pseudo (regs.getN 0)
  | .neg, regs => RV64.neg_pseudo (regs.getN 0)
  | .negw, regs => RV64.negw_pseudo (regs.getN 0)
  | .sextw, regs => RV64.sextw_pseudo (regs.getN 0)
  | .zextb, regs => RV64.zextb_pseudo (regs.getN 0)
  | .zextw, regs => RV64.zextw_pseudo (regs.getN 0)
  | .seqz, regs => RV64.seqz_pseudo (regs.getN 0)
  | .snez, regs => RV64.snez_pseudo (regs.getN 0)
  | .sltz, regs => RV64.sltz_pseudo (regs.getN 0)
  | .sgtz, regs => RV64.sgtz_pseudo (regs.getN 0)

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: @[simp]
abbrev RV64 : Dialect where
  Op := Op
  Ty := Ty

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

content: @[simp_riscv]
def neg_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.sub rs1_val 0

content: @[simp_riscv]
def seqz_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.sltiu 1 rs1_val

content: @[simp_riscv]
def sgtz_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.slt rs1_val 0

content: @[simp_riscv]
def zextb_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.andi 255 rs1_val

content: @[simp_riscv]
def snez_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.sltu rs1_val 0

content: @[simp_riscv]
def sltz_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.slt 0 rs1_val

content: @[simp_riscv]
def zextw_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.adduw 0 rs1_val

content: @[simp_riscv]
def sextw_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.addiw 0 rs1_val

content: @[simp_riscv]
def not_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.xori (-1) rs1_val

content: @[simp_riscv]
def mv_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.addi 0 rs1_val

content: @[simp_riscv]
def negw_pseudo (rs1_val : BitVec 64) : BitVec 64 :=
  RV64.subw rs1_val 0

content: @[simp]
noncomputable instance : DialectDenote (FHE q n) where
    denote
    | Op.add, arg, _ => [(fun args : R q n × R q n => args.1 + args.2) arg.toPair]ₕ
    | Op.sub, arg, _ => [(fun args : R q n × R q n => args.1 - args.2) arg.toPair]ₕ
    | Op.mul, arg, _ => [(fun args : R q n × R q n => args.1 * args.2) arg.toPair]ₕ
    | Op.mul_constant, arg, _ => [(fun args : R q n × Int => args.1 * ↑(args.2)) arg.toPair]ₕ
    | Op.leading_term, arg, _ => [R.leadingTerm arg.toSingle]ₕ
    | Op.monomial, arg, _ => [(fun args => R.monomial ↑(args.1) args.2) arg.toPair]ₕ
    | Op.monomial_mul, arg, _ => [(fun args : R q n × Nat => args.1 * R.monomial 1 args.2) arg.toPair]ₕ
    | Op.from_tensor, arg, _ => [R.fromTensor arg.toSingle]ₕ
    | Op.to_tensor, arg, _ => [R.toTensor' arg.toSingle]ₕ
    | Op.const c, _arg, _
    | Op.const_int c, _, _
    | Op.const_idx c, _, _ => [c]ₕ

content: noncomputable def R.fromTensor {q n} (coeffs : List Int) : R q n :=
  coeffs.zipIdx.foldl (init := 0) fun res (c, i) =>
    res + R.monomial ↑c i

content: noncomputable def R.monomial {q n : Nat} (c : ZMod q) (i : Nat): R q n :=
  R.fromPoly (Polynomial.monomial i c)

content: abbrev R.fromPoly {q n : Nat} : (ZMod q)[X] →+* R q n := Ideal.Quotient.mk (Ideal.span {f q n})

content: noncomputable def f : (ZMod q)[X] := X^(2^n) + 1

content: noncomputable def R.leadingTerm {q n} (a : R q n) : R q n :=
  let deg? := Polynomial.degree a.representative
  match deg? with
    | .none => 0
    | .some deg =>  R.monomial (a.coeff deg) deg

content: noncomputable def R.representative :
    R q n → (ZMod q)[X] := fun x => R.representative' q n x %ₘ (f q n)

content: private noncomputable def R.representative' :
    R q n → (ZMod q)[X] := Function.surjInv (R.surjective_fromPoly q n)

content: noncomputable def R.coeff {q n} (a : R q n) (i : Nat) : ZMod q :=
  Polynomial.coeff a.representative i

content: noncomputable def R.toTensor' {q n} (a : R q n) : List Int :=
  let t := a.toTensor
  t ++ List.replicate (2^n - t.length + 1) 0

content: noncomputable def R.toTensor {q n} (a : R q n) : List Int :=
  List.range a.repLength |>.map fun i =>
        a.coeff i |>.toInt

content: noncomputable def R.repLength {q n} (a : R q n) : Nat := match
  Polynomial.degree a.representative with
    | none => 0
    | some d => d + 1

content: instance : DialectDenote LLVM := ⟨
  fun o args _ => [Op.denote o args]ₕ
⟩

content: @[reducible]
instance : DialectDenote SimpleReg where
  denote
    | .const n, _, _ => BitVec.ofInt 32 n ::ₕ .nil
    | .add, [(a : BitVec 32), (b : BitVec 32)]ₕ , _ => a + b ::ₕ .nil
    | .iterate k, [(x : BitVec 32)]ₕ, [(f : _ → _)]ₕ =>
      let f := fun y => (f y).getN 0
      let f' (v :  BitVec 32) : BitVec 32 := f (Ctxt.Valuation.nil.cons v)
      let y := k.iterate f' x
      [y]ₕ

content: inductive Op :  Type
  | add : Op
  | const : (val : ℤ) → Op
  | iterate (k : ℕ) : Op
  deriving DecidableEq

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

content: abbrev SimpleReg.int : SimpleReg.Ty := .int

content: abbrev SimpleReg : Dialect where
  Op := Op
  Ty := Ty

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

content: inductive Ty
  | int
  deriving DecidableEq

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

content: @[reducible]
instance : DialectDenote Ex where
  denote
    | .cst n, _, _ => n ::ₕ .nil
    | .add, (a : Nat) ::ₕ b ::ₕ .nil, _ => a + b    ::ₕ .nil
    | .beq, (a : Nat) ::ₕ b ::ₕ .nil, _ => (a == b) ::ₕ .nil

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: inductive ExTy
  | nat
  | bool
  deriving DecidableEq, Repr

content: @[simp, reducible]
instance : DialectDenote (LLVMPlusRiscV) where
  denote
  | .llvm llvmOp, args, .nil => do
      let xs ← DialectDenote.denote llvmOp (llvmArgsFromHybrid args) .nil
      return xs.map' Ty.llvm (fun t x => x)
  | .riscv (riscvOp), args, .nil => do
      let xs ← DialectDenote.denote riscvOp (riscvArgsFromHybrid args) .nil
      return xs.map' Ty.riscv (fun t x => x)
  | .castRiscv _ , elemToCast, _ =>
    let toCast : BitVec 64 :=
      elemToCast.getN 0 (by admit /- proof elided -/
      )
    [castriscvToLLVM toCast]ₕ
  | .castLLVM _,
    (elemToCast : HVector TyDenote.toType [Ty.llvm (.bitvec _)]), _ =>
    let toCast : PoisonOr (BitVec _) :=
      elemToCast.getN 0 (by admit /- proof elided -/
      )
    [castLLVMToriscv toCast]ₕ

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: @[simp_riscv]
def castLLVMToriscv (toCast : PoisonOr (BitVec w)) : BitVec 64 :=
  BitVec.signExtend 64 (toCast.toOption.getD 0#w)

content: structure PoisonOr (α : Type) where
  ofOption :: toOption : Option α
  deriving DecidableEq

content: @[simp_riscv]
def castriscvToLLVM (toCast : BitVec 64) : PoisonOr (BitVec w) :=
  .value (BitVec.signExtend w toCast)

content: @[simp_denote]
def llvmArgsFromHybrid : {tys : List LLVM.Ty} →
  HVector TyDenote.toType (tys.map LLVMRiscV.Ty.llvm) → HVector TyDenote.toType tys
  | [], .nil => .nil
  | _ :: _, .cons x xs => .cons x (llvmArgsFromHybrid xs)

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

content: @[simp_denote]
def riscvArgsFromHybrid : {tys : List RISCV64.RV64.Ty} →
  HVector TyDenote.toType (tys.map LLVMRiscV.Ty.riscv) → HVector TyDenote.toType tys
  | [], .nil => .nil
  | _ :: _, .cons x xs => .cons x (riscvArgsFromHybrid xs)

content: @[reducible]
instance : DialectDenote ExOp ExTy where
  denote
    | .cst n, _, _ => n
    | .add, .cons (a : Nat) (.cons b .nil), _ => a + b
    | .beq, .cons (a : Nat) (.cons b .nil), _ => a == b

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: @[reducible]
instance : DialectDenote ExOp ExTy where
  denote
    | .add, .cons (a : Nat) (.cons b .nil), _ => a + b
    | .runK (k : Nat), (.cons (v : Nat) .nil), (.cons rgn _nil) =>
      k.iterate (fun val => rgn (fun _ty _var => val)) v

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 [Monad d.m] : DialectDenote (Scf d) where
  denote
    | .coe o', args', regArgs' =>
        let denote' := DialectDenote.denote o'
        by
         exact denote' args' regArgs'
    | .if t t', (.cons (cond ) (.cons v .nil)),
         (.cons (f : Ctxt.Valuation ⟨[t]⟩ → d.m (HVector toType [t'])) (.cons (g : _ → _) .nil)) =>
         let body := if B.denote_eq ▸ cond then f else g
      body (Ctxt.Valuation.nil.cons v)
    | .run _t, (.cons v .nil), (.cons (f : _ → _) .nil) =>
        f (Ctxt.Valuation.nil.cons v)
    | .for ty, (.cons istart (.cons istep (.cons niter (.cons vstart .nil)))),
        (.cons (f : _  → _) .nil) => do
        let istart : ℤ := Z.denote_eq ▸ istart
        let istep : ℤ := Z.denote_eq ▸ istep
        let niter : ℕ := N.denote_eq ▸ niter
        let f' : LoopBody (d.m ⟦ty⟧) := fun i v => do
          let v ← v
          let i := Z.denote_eq.symm ▸ i
          let xs ← f (Valuation.ofPair i v)
          return xs.get (0 : Fin 1)
        let to_iterate := f'.counterDecorator (α := d.m ⟦ty⟧) (δ := istep)
        let loop_fn := niter.iterate (op := to_iterate)
        let x ← (loop_fn (istart, pure vstart)).2
        return [x]ₕ

    | .iterate k, (.cons (x) .nil), (.cons (f : _ → _) .nil) => do
      let x : ℤ := Z.denote_eq ▸ x
      let coe : ℤ = toType Z.ty := Z.denote_eq.symm
      let f' (v : d.m ℤ) : d.m ℤ := do
        let v ← v
        let xs ← f (Ctxt.Valuation.nil.cons (cast coe v))
        let x := xs.getN 0
        return coe ▸ x
      let y ← (k.iterate f' (pure x))
      return [cast Z.denote_eq.symm y]ₕ

content: def counterDecorator (δ : Int) (f : LoopBody α) : Int × α → Int × α :=
  fun (i, v) => (i + δ, f i v)

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: 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: abbrev LoopBody (t : Type) : Type := Int → t → t

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_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: def Valuation.ofPair  {t₁ t₂ : Ty} (v₁: ⟦t₁⟧) (v₂ : ⟦t₂⟧) :
    Valuation (Ctxt.ofList [t₁, t₂]) :=
  Valuation.ofHVector (.cons v₁ <| .cons v₂ <| .nil )

content: @[reducible]
instance : DialectDenote Arith where
  denote
    | .const n, _, _ => [n]ₕ
    | .const_nat n, _, _ => [n]ₕ
    | .neg, .cons (a : ℤ ) .nil, _ => [-a]ₕ
    | .axpy, .cons (a : ℤ) (.cons (x : ℕ) (.cons (b : ℤ) .nil)), _ => [a * (x : ℤ) + b]ₕ
    | .add, .cons (a : ℤ) (.cons (b : ℤ) .nil), _ => [a + b]ₕ
    | .add_nat, .cons (a : ℕ) (.cons (b : ℕ) .nil), _ => [a + b]ₕ

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)

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

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: @[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: @[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: 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: @[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: instance : DialectSignature ExOp ExTy where
  signature
    | .add    => ⟨[.nat, .nat], [], .nat, .pure⟩
    | .beq    => ⟨[.nat, .nat], [], .bool, .pure⟩
    | .cst _  => ⟨[], [], .nat, .pure⟩

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

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: @[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: infixr:50 "::ₕ" => HVector.cons

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

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

-- Repository lemmas used:
content: @[simp]
theorem pure_liftEffect {eff₁ eff₂ : EffectKind}
    (hle : eff₁ ≤ .pure) [Monad m] (x : eff₁.toMonad m α) :
    (Pure.pure (liftEffect hle x) : eff₂.toMonad m α)
    = liftEffect (by cases hle; constructor) x

content: @[simp]
theorem eq_of_le_pure {e : EffectKind}
    (he : e ≤ pure) : e = pure

content: theorem pure_map (f : α → β) (x : pure.toMonad m α) (eff : EffectKind) :
    (Pure.pure (f <$> x : pure.toMonad m _) : eff.toMonad m _) = f <$> (Pure.pure x)

-- 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.op {Γ : Ctxt d.Ty} {eff : EffectKind} {ty} (e : Expr d Γ eff ty) : d.Op :=
  Expr.casesOn e (fun op _ _ _ _ => op)

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: 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: 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 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 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 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: 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: @[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 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 Expr.castPureToEff (eff : EffectKind) : Expr d Γ .pure t → Expr d Γ eff t :=
  changeEffect (EffectKind.pure_le eff)

-- Local lemmas used:
content: @[simp]
theorem Expr.op_mk {Γ : Ctxt d.Ty} {ty} {eff : EffectKind} (op : d.Op)
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) (DialectSignature.sig op))
    (regArgs) :
    (Expr.mk op ty_eq eff_le args regArgs).op = op

content: @[simp]
theorem Expr.args_mk {Γ : Ctxt d.Ty} {ty eff op}
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) (DialectSignature.sig op)) (regArgs) :
    (Expr.mk op ty_eq eff_le args regArgs).args = args

content: @[simp]
theorem Expr.regArgs_mk {Γ : Ctxt d.Ty} {ty eff op}
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) (DialectSignature.sig op)) (regArgs) :
    (Expr.mk op ty_eq eff_le args regArgs).regArgs = regArgs

content: theorem Expr.denote_unfold (e : Expr d Γ eff ty) :
    e.denote = fun V => (· ++ V) <$> (e.denoteOp V)

-- Target theorem and proof:
@[simp] theorem Expr.denote_castPureToEff {e : Expr d Γ .pure t} :
    denote (e.castPureToEff eff) = fun V => pure (e.denote V)
:= by
  rcases e with ⟨op, rfl, eff_le, _, _⟩
  cases eff
  case pure => rfl
  case impure =>
    funext V
    simp only [castPureToEff, changeEffect, denote_unfold, denoteOp, op_mk, args_mk, regArgs_mk,
      EffectKind.pure_map, EffectKind.pure_liftEffect]