Ctxt

structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty

  deriving Repr, Lean.ToExpr, DecidableEq

EffectKind

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

Ty

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

Valuation.nil

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

Valuation.cons

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

Var

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

map

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

Hom

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

dropUntil

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

appendCases

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

Hom.castCodomain

def Hom.castCodomain (h : Δ = Δ') (f : Γ.Hom Δ) : Γ.Hom Δ' :=
  fun _t v => (f v).castCtxt h

appendInl

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

Valuation.cast

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

dropUntilHom

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

dropUntilDiff

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

emptyElim

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

cons

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

Diff

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

Diff.Valid

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

Hom.id

@[simp] abbrev Hom.id {Γ : Ctxt Ty} : Γ.Hom Γ :=
  fun _ v => v

appendInr

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

length

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

map

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)

HVectorLiteral

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
    )

(q

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

Op.signature

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

Op.sig

@[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 _ => []

Op

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

Op.outTy

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

Op.regSig

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

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

ExOp

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

ExTy

inductive ExTy
  | nat
  | bool
  deriving DecidableEq

add

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)

Ex

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

cst

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)

Dialect

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

TyDenote

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

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

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

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

Op.sig

@[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]

Op

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

Ty

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

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

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

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

    | .hs o => MLIR2Handshake.instDialectSignatureHandshake.signature o

Op

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

liftSig

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

liftTy

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

Ty

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

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

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

map

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

Stream

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

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

    | .dc o => MLIR2DC.instDialectSignatureDC.signature o

Op

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

liftSig

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

liftTy

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

Ty

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

[SIG

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

ExTy

inductive ExTy
  | nat
  | bool
  deriving DecidableEq, Repr

LLVMPlusRiscVSignature

@[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 := []}

Op

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

Ty

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

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

ExOp

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

add

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)

ExTy

inductive ExTy
  | nat
  deriving DecidableEq, Repr

Expr

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

cst

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

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

ExOp

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

add

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)

[TyDenote

@[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⟩

HasTy

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

Scf.Op

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

iterate

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

ScfArith

abbrev ScfArith := Scf Arith

Scf

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

Op

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

run

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

Ty

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

Valuation.instAppendHVector

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

neg

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

axpy

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

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

Ty

inductive Ty
  | int

  | int2
  deriving DecidableEq, Lean.ToExpr

Op

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

Arith

abbrev Arith : Dialect := {Op, Ty}

@[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⟩

add

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

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

TyDenote.toType

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

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

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

RegionSignature

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

Signature.mk

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

DialectSignature

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

DialectSignature.sig

def sig          := Signature.sig ∘ s.signature

DialectSignature.regSig

def regSig       := Signature.regSig ∘ s.signature

DialectSignature.returnTypes

def returnTypes  := Signature.returnTypes ∘ s.signature

DialectSignature.effectKind

def effectKind   := Signature.effectKind ∘ s.signature

Expr.outContext

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

Regions

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

Expr.args

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

Com.outContext

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

Expr.contextHom

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

Lets.getPureExpr

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

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

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

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⟩

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


theorem mem_matchArg {Δ_out}
    {lets : Lets d Γ_in eff Γ_out}
    {matchLets : Lets d Δ_in .pure Δ_out}
    {l : List d.Ty} {argsₗ : HVector (Var Γ_out) l}
    {argsᵣ : HVector (Var Δ_out) l} {ma : Mapping Δ_in Γ_out}
    {varMap : Mapping Δ_in Γ_out}
    (hvarMap : ((), varMap) ∈ matchArg lets matchLets argsₗ argsᵣ ma)
    {t' v'} : ⟨t', v'⟩ ∈ matchLets.varsOfVec argsᵣ → ⟨t', v'⟩ ∈ varMap :=
  match l, argsₗ, argsᵣ/- , ma, varMap, hvarMap -/ with
  | .nil, .nil, .nil /- , _, varMap, _ -/ => by simp [Lets.varsOfVec]
  | .cons t ts, .cons vₗ argsₗ, .cons vᵣ args /-, ma, varMap, h -/ => by
    simp only [matchArg, bind, Option.mem_def, StateT.bind, Option.bind_eq_some_iff] at hvarMap
    rcases hvarMap with ⟨ma', h₁, h₂⟩
    simp only [HVector.vars_cons, Finset.biUnion_insert, Finset.mem_union,
      Finset.mem_biUnion, Sigma.exists, Lets.varsOfVec]
    rintro (h | ⟨a, b, hab⟩)
    · exact AList.keys_subset_keys_of_entries_subset_entries
        (isMonotone_matchArg _ _ h₂)
        (mem_matchVar (matchLets := matchLets) h₁ h)
    · apply mem_matchArg h₂
      unfold Lets.varsOfVec
      apply Finset.mem_biUnion.mpr ⟨_, hab.1, hab.2⟩

Mapping

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

Ctxt

structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty

  deriving Repr, Lean.ToExpr, DecidableEq

Var

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

Lets

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

Expr.decidableEq

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

Var

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

Valuation.nil

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

EffectKind

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

appendCases

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

last

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

Valuation.cons

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

HVector.decidableEqReg

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

Lets.varsOfVec

def Lets.varsOfVec (lets : Lets d Γ_in eff Γ_out) (vs : HVector Γ_out.Var ts) :
    VarSet Γ_in :=
  (vs.vars).biUnion (fun v => lets.vars v.2)

Lets.vars

def Lets.vars : Lets d Γ_in eff Γ_out → Var Γ_out t → VarSet Γ_in
  | .nil, v => VarSet.ofVar v
  | .var lets e, v => by admit /- proof elided -/
      | right v => exact lets.vars v
      | left _ => exact lets.varsOfVec e.args

Com.vars

def Com.vars (com : Com d Γ eff ts) : VarSet Γ :=
  com.toLets.varsOfVec com.returnVars

Expr.returnVars

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

Com.outContext

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

Com.rec'

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

Com

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 β

Valuation.cast

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

Expr.outContext

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

ofFin

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

Com.returnVars

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

Com.toLets

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

VarSet

abbrev VarSet (Γ : Ctxt Ty) : Type :=
  Finset (Σ t, Γ.Var t)

Expr.args

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

sig

def sig          := Signature.sig ∘ s.signature

Signature

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

DialectSignature

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

HVector.vars

def HVector.vars {l : List d.Ty} (T : HVector (Var Γ) l) : VarSet Γ :=
  T.foldl (fun _ s a => insert ⟨_, a⟩ s) ∅

ofVar

@[simp]
def ofVar {Γ : Ctxt Ty} (v : Γ.Var t) : VarSet Γ :=
  {⟨_, v⟩}

foldl

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

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

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

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

@[simp] theorem HVector.vars_cons {t  : d.Ty} {l : List d.Ty}
    (v : Var Γ t) (T : HVector (Var Γ) l) :
    (HVector.cons v T).vars = insert ⟨_, v⟩ T.vars

@[simp] theorem appendCases_appendInr (v : Γ.Var t) :
    appendCases (motive := motive) left right v.appendInr = (right v)

MatchVar.IsMonotone

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

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

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

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

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

theorem MatchVar.isMonotone_unifyVars  : IsMonotone (unifyVars w v)

theorem isMonotone_matchVarArg_aux (lets : Lets d Γ_in eff Γ_out) :
    (
     ∀  (Δ_out : Ctxt d.Ty)
        (matchLets : Lets d Δ_in EffectKind.pure Δ_out) (l : List d.Ty)
        (argsl : HVector Γ_out.Var l) (argsr : HVector Δ_out.Var l),
        (matchArg lets matchLets argsl argsr).IsMonotone
    )
    ∧ (
      ∀ (Δ_out : Ctxt d.Ty) (t : d.Ty) (v : Γ_out.Var t)
        (matchLets : Lets d Δ_in EffectKind.pure Δ_out)
        (w : Var Δ_out t),
        (matchVar lets v matchLets w).IsMonotone
    )

theorem isMonotone_matchArg [DecidableEq d.Op]
    {Γ_out Δ_in Δ_out : Ctxt d.Ty}
    {lets : Lets d Γ_in eff Γ_out}
    {matchLets : Lets d Δ_in .pure Δ_out}
    {l : List d.Ty}
    {argsl : HVector (Var Γ_out) l}
    {argsr : HVector (Var Δ_out) l} :
    (matchArg lets matchLets argsl argsr).IsMonotone

import LeanMLIR.Framework

import LeanMLIR.Transforms.Rewrite.Mapping

open Ctxt (Var VarSet Valuation Hom)

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

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

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

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

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

open MatchVarM

variable [DecidableEq d.Op]

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

def matchVar {Γ_in Γ_out Δ_in Δ_out : Ctxt d.Ty} {t : d.Ty} [DecidableEq d.Op]
    (lets : Lets d Γ_in eff Γ_out) (v : Var Γ_out t) :
    (matchLets : Lets d Δ_in .pure Δ_out) →
    (w : Var Δ_out t) →
    MatchVar Δ_in Γ_out

  | @Lets.var _ _ _ _ Δ_out ts matchLets matchExpr, w => by admit /- proof elided -/
      | right w =>
        exact matchVar lets v matchLets w
      | left w => exact do
        let ⟨ts', w', ie⟩ ← lets.getPureExpr v
        if hs : ∃ h : ie.op = matchExpr.op, ie.regArgs = (h ▸ matchExpr.regArgs) then
          have hts : Ctxt.ofList ts' = ts := by admit /- proof elided -/

end

section MatchVar

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

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

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

namespace MatchVarResult

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

section Left

variable {w : Δ_out.Var t}

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

end Left

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

end MatchVarResult

end MatchVar

section SubsetEntries

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

open MatchVar

section UnifyVars

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

end UnifyVars

variable [DecidableEq d.Op]

end SubsetEntries

namespace MatchArgResult

variable [DecidableEq d.Op] {Γ_in Γ_out Δ_in Δ_out te}
          {lets : Lets d Γ_in eff Γ_out}
          {matchLets : Lets d Δ_in .pure Δ_out}
          {matchExpr : Expr d Δ_out .pure te}
          {u us}
          {v : Γ_out.Var u} {vs : HVector Γ_out.Var us}
          {w : Δ_out.Var u} {ws : HVector Δ_out.Var us}
          {mapIn : Mapping _ _}
          (mapOut : MatchArgResult lets matchLets (v ::ₕ vs) (w ::ₕ ws) mapIn)

end MatchArgResult

section DenoteLemmas

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

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

variable [Monad d.m] [LawfulMonad d.m] [DialectDenote d]

section DenoteIntoSubtype

end DenoteIntoSubtype

variable {Γ_in Γ_out Δ_in Δ_out : Ctxt d.Ty}
    {lets : Lets d Γ_in eff Γ_out}
    {matchLets : Lets d Δ_in .pure Δ_out}


State

abbrev State := Nat

RawCNFA.states

def RawCNFA.states (m : RawCNFA A) : Finset State := Finset.range m.stateMax

RawCNFA.statesFinset

instance RawCNFA.statesFinset (m : RawCNFA A) : Fintype m.states := (Finset.range m.stateMax).fintypeCoeSort

CNFA

structure CNFA (n : Nat) where
  m : RawCNFA (BitVec n)
  wf : m.WF

Std.HashSet.toSet

def Std.HashSet.toSet [BEq α] [Hashable α] (m : HashSet α) : Set α := { x | x ∈ m }

theorem Std.HashSet.fold_induction [BEq α] [LawfulBEq α] [Hashable α]
  {f : β → α → β} {m : HashSet α} {motive : β → Set α → Prop} :
    motive b ∅ →
    (∀ b x s, x ∉ s → motive b s → motive (f b x) (s ∪ {x})) →
    motive (m.fold f b) m.toSet

theorem Std.HashSet.toSet_toList[BEq α] [LawfulBEq α] [Hashable α] (m : HashSet α) : m.toSet = { x | x ∈ m.toList }

@[simp]
lemma bv_to_set_or {m : CNFA n} (x y : BitVec m.m.stateMax) :
    (s ∈ bv_to_set (x ||| y)) ↔ (s ∈ bv_to_set x ∨ s ∈ bv_to_set y)

@[simp]
lemma bv_to_set_shift (x s : Nat) :
    (s ∈ bv_to_set (1#w <<< x)) ↔ (s = x ∧ x < w)

import Mathlib.Data.Fintype.Prod

import Blase.AutoStructs.Worklist

import Mathlib.Tactic.ApplyFun

open SetRel

section sink

variable {A : Type} [BEq A] [Hashable A] [DecidableEq A] [FinEnum A]

end sink

section generic_prod

variable {α} [BEq α] [Hashable α] [LawfulBEq α]

variable {β} [BEq β] [Hashable β] [LawfulBEq β]

variable {S₁ : Finset α} {S₂ : Finset β}

variable {γ} (f : S₁ → S₂ → γ) (hinj : Function.Injective2 f)

variable {m₁ : Std.HashSet α} (hm₁ : ∀ s₁ ∈ m₁, s₁ ∈ S₁)

variable {m₂ : Std.HashSet β} (hm₂ : ∀ s₂ ∈ m₂, s₂ ∈ S₂)

end generic_prod

section product

variable {A : Type} [BEq A] [LawfulBEq A] [Hashable A] [DecidableEq A] [FinEnum A]

end product

section determinization

variable {A : Type} [BEq A] [LawfulBEq A] [Hashable A] [DecidableEq A] [FinEnum A]

private def bv_to_set (bv : BitVec w) : Set State :=
  { s | bv.getLsbD s }

@[simp]
lemma transBV_spec {m : CNFA n} {res} {s : m.m.states} :
    s' ∈ bv_to_set (m.m.transBV' res s a) ↔
      (s' ∈ bv_to_set res ∨ s' ∈ m.m.tr s a) :=

:= by
  let motive (bv : BitVec m.m.stateMax) (X : Set State) :=
    ∀ s' (hlt : s' < m.m.stateMax),
      s' ∈ bv_to_set bv ↔ (s' ∈ bv_to_set res ∨ s' ∈ X)
  suffices h : motive (m.m.transBV' res s a) (m.m.tr s a).toSet by
    by_cases hlt : s' < m.m.stateMax
    · simp_all [motive]
    · constructor
      · simp_all [bv_to_set]
      · rintro (h | h); simp_all [bv_to_set]
        apply m.wf.trans_tgt_lt at h; simp_all only [RawCNFA.states, not_lt, Finset.mem_range]
        suffices _ : m.m.stateMax < m.m.stateMax by simp_all
        exact Nat.lt_of_le_of_lt hlt h
  apply Std.HashSet.fold_induction
  · simp [motive]
  rintro bv s S hnin ih s' hlt
  simp only [bv_to_set_or, bv_to_set_shift, Set.union_singleton, Set.mem_insert_iff]
  constructor
  · rintro (hold | ⟨rfl, hin⟩)
    · apply ih _ hlt |>.mp at hold; tauto
    · tauto
  · rintro (hold | rfl | hS)
    · left; apply (ih _ hlt).mpr; tauto
    · tauto
    · left; apply (ih _ hlt).mpr; tauto


BitVecs

structure BitVecs (n : Nat) where
  w : Nat
  bvs : List.Vector (BitVec w) n

RelationOrdering

inductive RelationOrdering
| lt | le | gt | ge
deriving Repr, Fintype

NFA'

structure NFA' (n : Nat) where
  σ : Type
  M : NFA (BitVec n) σ

Relation

inductive Relation
| eq
| signed (ord : RelationOrdering)
| unsigned (ord : RelationOrdering)
deriving Repr

evalRelation

def evalRelation (rel : Relation) {w} (bv1 bv2 : BitVec w) : Prop :=
  match rel with
  | .eq => bv1 = bv2
  | .signed .lt => bv1.slt bv2
  | .signed .le => bv1.sle  bv2
  | .signed .gt => bv2.slt bv1
  | .signed .ge => bv2.sle bv1
  | .unsigned .lt => bv1.ult bv2
  | .unsigned .le => bv1.ule bv2
  | .unsigned .gt => bv2.ult bv1
  | .unsigned .ge => bv2.ule bv1

def bv2 : BitVec 4 := BitVec.ofNat 4 1

def bv1 : BitVec 4 := BitVec.ofNat 4 5

accepts

def accepts (M : NFA' n) : Set (BitVecs n) := dec '' M.accepts'

accepts'

def accepts' (M : NFA' n) : Set (BitVecs' n) := M.M.accepts

BitVecs'

abbrev BitVecs' (n : Nat) := List (BitVec n)

dec

@[simps]
def dec (bvs' : BitVecs' n) : BitVecs n where
  w := bvs'.length
  bvs := List.Vector.ofFn fun k => BitVec.ofFn fun i => bvs'[i].getLsbD k

BitVec.ofFn

def BitVec.ofFn {w : Nat} (f : Fin w → Bool) : BitVec w :=
  BitVec.iunfoldr (fun i _ => ((), f i)) () |>.2

BitStream

def BitStream : Type := Nat → Bool

enc

def enc (bvs : BitVecs n) : BitVecs' n :=
  (List.finRange bvs.w).map (fun i =>
    BitVec.ofFn (fun (k : Fin n) => (bvs.bvs.get k)[i]))

instFinEnumBV

instance instFinEnumBV : FinEnum (BitVec w) where
  card := 2^w
  equiv := {
    toFun := fun x => x.toFin
    invFun := fun x => BitVec.ofFin x
    left_inv := by admit /- proof elided -/

ext

@[ext]
theorem ext {x y : BitStream} (h : ∀ i, x i = y i) : x = y

Relation.absAutOfRelation

lemma NFA.correct_spec {M : NFA α σ} {ζ : M.sa} {L : Language α} :
    M.correct ζ L → M.accepts = L

@[simp]
lemma in_enc : x ∈ enc '' S ↔ dec x ∈ S

@[simp]
lemma dec_snoc_in_langRel {n} {R : BVNRel n} {w : BitVecs' n} {a : BitVec n} :
    dec (w ++ [a]) ∈ langRel R ↔
      R (List.Vector.ofFn fun k => .cons (a.getLsbD k) ((dec w).bvs.get k))

lemma NFA'.correct_spec {M : NFA' n} {ζ : M.sa} {L : BVNRel n} :
    M.correct ζ L → M.accepts = langRel L

lemma NFA'.correct2_spec {M : NFA' 2} {ζ : M.sa2} {L : BVRel} :
    M.correct2 ζ L → M.accepts = langRel2 L

lemma NFA'.autEq_correct : autEq.correct2 (fun _ => eqRel) eqRel

lemma BitVec.ule_iff_ult_or_eq {w : ℕ} (bv1 bv2 : BitVec w):
    (bv1.ule bv2) = true ↔ (bv1.ult bv2) = true ∨ bv1 = bv2

@[simp]
lemma ucmp_tricho {bv1 bv2 : BitVec w} : (bv2.ult bv1) = false → (bv1.ult bv2) = false → bv1 = bv2

lemma NFA'.autUnsignedCmp_correct cmp : autUnsignedCmp cmp |>.correct2 autUnsignedCmpSA cmp.urel

private lemma BitVec.sle_iff_slt_or_eq {w : ℕ} (bv1 bv2 : BitVec w):
    (bv1.sle bv2) = true ↔ (bv1.slt bv2) = true ∨ bv1 = bv2

lemma NFA'.autSignedCmp_correct cmp : autSignedCmp cmp |>.correct2 autSignedCmpSA cmp.srel

import Batteries.Data.Fin.Basic

import Batteries.Data.Fin.Lemmas

import Blase.SingleWidth.Defs

import Blase.AutoStructs.Constructions

import Blase.AutoStructs.Defs

import Blase.AutoStructs.FiniteStateMachine

import Mathlib.Tactic.Ring

import Mathlib.Data.Nat.Size -- TODO: remove and get rid of shiftLeft_eq_mul_pow use

import Mathlib.Data.BitVec

import Mathlib.Tactic.FinCases

open Fin.NatCast

open Mathlib

def NFA.sa (_ : NFA α σ) := σ → Language α

structure NFA.correct (M : NFA α σ) (ζ : M.sa) (L : Language α) where
  cond1 : ∀ w, (w ∈ L ↔ ∃ q ∈ M.accept, w ∈ ζ q)
  cond2 : ∀ w q, q ∈ M.eval w ↔ w ∈ ζ q

abbrev BVRel := ∀ ⦃w⦄, BitVec w → BitVec w → Prop

abbrev BVNRel n := ∀ ⦃w⦄, List.Vector (BitVec w) n → Prop

def NFA'.sa (M : NFA' n) := M.σ → BVNRel n

def NFA'.sa2 (M : NFA' 2) := M.σ → BVRel

def langRel (R : BVNRel n) : Set (BitVecs n) :=
  { bvs | R bvs.bvs }

def langRel2 (R : BVRel) : Set (BitVecs 2) :=
  { bvs | R (bvs.bvs.get 0) (bvs.bvs.get 1) }

structure NFA'.correct (M : NFA' n) (ζ : M.sa) (L : BVNRel n) where
  cond1 : ∀ ⦃w⦄ (bvn : List.Vector (BitVec w) n), (L bvn ↔ ∃ q ∈ M.M.accept, ζ q bvn)
  cond2 q : q ∈ M.M.start ↔ ζ q (List.Vector.replicate n .nil)
  cond3 q a {w} (bvn : List.Vector (BitVec w) n) : q ∈ M.M.stepSet { q | ζ q bvn } a ↔
              ζ q (List.Vector.ofFn fun k => BitVec.cons (a.getLsbD k) (bvn.get k))

structure NFA'.correct2 (M : NFA' 2) (ζ : M.sa2) (L : BVRel) where
  cond1 : ∀ (bv1 bv2 : BitVec w), (L bv1 bv2 ↔ ∃ q ∈ M.M.accept, ζ q bv1 bv2)
  cond2 q : q ∈ M.M.start ↔ ζ q .nil .nil
  cond3 q a w (bv1 bv2 : BitVec w) : q ∈ M.M.stepSet { q | ζ q bv1 bv2 } a ↔
              ζ q (BitVec.cons (a.getLsbD 0) bv1) (BitVec.cons (a.getLsbD 1) bv2)

section fsm

variable {arity : Type} [FinEnum arity]

open BitStream in

end fsm

section nfas_relations

def NFA.autEq : NFA (BitVec 2) Unit :=
  { start := ⊤, accept := ⊤, step _ a := { _s' | if a = 0 ∨ a = 3 then true else false }}

def NFA'.autEq : NFA' 2 :=
  ⟨Unit, NFA.autEq⟩

def NFA'.eqRel : BVRel := fun _ x y => x = y

inductive NFA.unsignedCmpState : Type where
| eq | gt | lt
deriving Fintype, DecidableEq

def NFA.unsignedCmpStep (q : NFA.unsignedCmpState) (a : BitVec 2) : List NFA.unsignedCmpState :=
  match q, a with
  | .eq, 0 => [ .eq ] | .eq, 3 => [ .eq ] | .eq, 1 => [ .gt ] | .eq, 2 => [ .lt ]
  | .gt, 0 => [ .gt ] | .gt, 1 => [ .gt ] | .gt, 3 => [ .gt ] | .gt, 2 => [ .lt ]
  | .lt, 0 => [ .lt ] | .lt, 1 => [ .gt ] | .lt, 2 => [ .lt ] | .lt, 3 => [ .lt ]

def NFA.autUnsignedCmp (cmp: RelationOrdering) : NFA (BitVec 2) unsignedCmpState where
  step s a := { s' | s' ∈ unsignedCmpStep s a }
  start := {s | s = .eq }
  accept := { s | s ∈ match cmp with | .lt => [unsignedCmpState.lt] | .le => [.lt, .eq] | .gt => [.gt] | .ge => [.gt, .eq] }

def NFA'.autUnsignedCmp (cmp: RelationOrdering) : NFA' 2 :=
  ⟨_, NFA.autUnsignedCmp cmp⟩

def RelationOrdering.urel (cmp : RelationOrdering) : BVRel :=
  match cmp with
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .le => fun _ bv1 bv2 => bv1.ule bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1
  | .ge => fun _ bv1 bv2 => bv2.ule bv1

def NFA'.autUnsignedCmpSA (q : NFA.unsignedCmpState) : BVRel :=
  match q with
  | .eq => fun _ bv1 bv2 => bv1 = bv2
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1

inductive NFA.signedCmpState : Type where
| eq | gt | lt | ltfin | gtfin
deriving DecidableEq, Fintype

def NFA.signedCmpStep (q : NFA.signedCmpState) (a : BitVec 2) : List NFA.signedCmpState :=
  match q, a with
  | .eq, 0 => [ .eq ] | .eq, 3 => [.eq] | .eq, 1 => [.gt, .ltfin] | .eq, 2 => [ .lt, .gtfin ]
  | .gt, 0 => [ .gt, .gtfin ] | .gt, 1 => [ .gt, .ltfin ] | .gt, 3 => [ .gt, .gtfin ] | .gt, 2 => [ .lt, .gtfin ]
  | .lt, 0 => [ .lt, .ltfin ] | .lt, 1 => [ .gt, .ltfin ] | .lt, 2 => [ .lt, .gtfin ] | .lt, 3 => [ .lt, .ltfin ]
  | .gtfin, _ => ∅
  | .ltfin, _ => ∅

def NFA.autSignedCmp (cmp: RelationOrdering) : NFA (BitVec 2) signedCmpState where
  step s a := { s' | s' ∈ signedCmpStep s a }
  start := { s | s = signedCmpState.eq }
  accept := { s | s ∈ match cmp with | .lt => [NFA.signedCmpState.ltfin] | .le => [.ltfin, .eq] | .gt => [.gtfin] | .ge => [.gtfin, .eq] }

def NFA'.autSignedCmp (cmp: RelationOrdering) : NFA' 2 :=
  ⟨_, NFA.autSignedCmp cmp⟩

def RelationOrdering.srel (cmp : RelationOrdering) : BVRel :=
  match cmp with
  | .lt => fun _ bv1 bv2 => bv1.slt bv2
  | .le => fun _ bv1 bv2 => bv1.sle bv2
  | .gt => fun _ bv1 bv2 => bv2.slt bv1
  | .ge => fun _ bv1 bv2 => bv2.sle bv1

def NFA'.autSignedCmpSA (q : NFA.signedCmpState) : BVRel :=
  match q with
  | .eq => fun _ bv1 bv2 => bv1 = bv2
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1
  | .ltfin => fun _ bv1 bv2 => bv1.slt bv2
  | .gtfin => fun _ bv1 bv2 => bv2.slt bv1

end nfas_relations

def Relation.absAutOfRelation (rel : Relation) : NFA' 2 :=
  match rel with
  | .eq => NFA'.autEq
  | .unsigned cmp => NFA'.autUnsignedCmp cmp
  | .signed cmp => NFA'.autSignedCmp cmp

@[simp]
lemma autOfRelation_accepts (r : Relation) :
    r.absAutOfRelation.accepts = r.language :=

:= by
  simp [Relation.absAutOfRelation]
  rcases r with ⟨⟩ | ⟨cmp⟩ | ⟨cmp⟩ <;> simp
  · rw [NFA'.correct2_spec NFA'.autEq_correct]
    simp [langRel2, NFA'.eqRel, evalRelation]
  · rw [NFA'.correct2_spec (NFA'.autSignedCmp_correct cmp)]
    simp [langRel2, evalRelation, RelationOrdering.srel]
    cases cmp <;> simp
  · rw [NFA'.correct2_spec (NFA'.autUnsignedCmp_correct cmp)]
    simp [langRel2, evalRelation, RelationOrdering.urel]
    cases cmp <;> simp


theorem correct : Com.denote (lhs v0) = Com.denote (rhs v0)

import LeanMLIR.Util
import LeanMLIR.Framework
import LeanMLIR.ErasedContext
import Mathlib.Tactic.Linarith
import Mathlib.Logic.Function.Iterate
import LeanMLIR.Tactic

macro "simp_peephole" loc:(location)? : tactic =>

macro "simp_peephole" loc:(location)? : tactic =>
  `(tactic|(

      first
      | rw [funext_iff (α := Ctxt.Valuation _)] $[$loc]?
      | change ∀ (_ : Ctxt.Valuation _), _ $[$loc]?
      | skip

      simp (config := {failIfUnchanged := false}) only
        [Expr.denote_castPureToEff, simp_denote] $[$loc]?

  ))

cons

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

Ctxt

structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty

  deriving Repr, Lean.ToExpr, DecidableEq

Var

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

Expr.decidableEq

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

Var

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

HVector.decidableEqReg

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

Com

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 β

FlatCom

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

Lets

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

EffectKind

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

DialectDenote

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

sig

def sig          := Signature.sig ∘ s.signature

Signature

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

DialectSignature

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

regSig

def regSig       := Signature.regSig ∘ s.signature

RegionSignature

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

Ty

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

effectKind

def effectKind   := Signature.effectKind ∘ s.signature

returnTypes

def returnTypes  := Signature.returnTypes ∘ s.signature

Dialect

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

Op

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

TyDenote

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

Valuation.instAppendHVector

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

Ty

inductive Ty
  | int

  | int2
  deriving DecidableEq, Lean.ToExpr

Op

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

Com.ret

def Com.ret {Γ : Ctxt d.Ty} {ty : d.Ty} {eff : EffectKind} : Γ.Var ty → Com d Γ eff [ty] :=
  (Com.rets [·]ₕ)

Com.letPure

def Com.letPure (e : Expr d Γ .pure t) (body : Com d (e.outContext) eff u) : Com d Γ eff u :=
  body.var (e.castPureToEff eff)

Expr.outContext

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

Expr.castPureToEff

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

Expr.changeEffect

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

Expr.args

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

Expr.regArgs

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

Regions

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

Expr.op

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

Com.denote

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

Lets.denote

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

HVector.denote

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)

FlatCom.denote

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

RegionSignature.map

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

Signature.map

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
  }

map

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

Expr.denote

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

Com.outContext

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

Com.rec'

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

Valuation.cast

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

liftEffect

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

toMonad

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

Valuation.nil

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

Valuation.cons

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

Com.ty

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

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

TyDenote.toType

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

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

theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

ScfFunctor.Scf.LoopBody

abbrev LoopBody (t : Type) : Type := Int → t → t

ScfFunctor.Scf.LoopBody.IndexInvariant

def IndexInvariant (f : LoopBody t) : Prop :=
  ∀ (i j : Int) (v : t), f i v = f j v

ScfFunctor.Scf.LoopBody.atZero

def atZero (f : LoopBody t) : t → t := fun v => f 0 v

ScfFunctor.Arith.Ty

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

ScfFunctor.Arith.Arith

abbrev Arith : Dialect := {Op, Ty}

ScfFunctor.ScfArith

abbrev ScfArith := Scf Arith

ScfFunctor.ForAddToMul.rhs

theorem eq_invariant_fn
    (f : LoopBody t) (g : t → t) (hf : ∀ (i : Int) (v : t), f i v = g v) :
    LoopBody.IndexInvariant f ∧ atZero f = g

@[simp]
theorem eval' {f : LoopBody t} (hf : LoopBody.IndexInvariant f) (i : Int) (v : t) :
    f i v = f.atZero v

import Mathlib.Logic.Function.Iterate

import Mathlib.Tactic.Linarith

import LeanMLIR.Framework

import LeanMLIR.Tactic

import LeanMLIR.ErasedContext

import LeanMLIR.Util

open LeanMLIR

open Ctxt(Var)

namespace ScfFunctor

open TyDenote

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

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

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

namespace Scf

section InheritedInstances

variable {d : Dialect} [TyDenote d.Ty] [DialectSignature d] [DialectDenote d]

end InheritedInstances

abbrev LoopBody (t : Type) : Type := Int → t → t

namespace LoopBody

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

def atZero (f : LoopBody t) : t → t := fun v => f 0 v

end LoopBody

namespace LoopBody.IndexInvariant

end LoopBody.IndexInvariant

namespace LoopBody.counterDecorator

end LoopBody.counterDecorator

variable [TyDenote d.Ty] [DialectSignature d] [DialectDenote d]
  [B : HasBool d] [N : HasNat d] [Z : HasInt d]

open Ctxt (Valuation) in

end Scf

namespace Arith

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

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

abbrev Arith : Dialect := {Op, Ty}

end Arith

abbrev ScfArith := Scf Arith

open LeanMLIR.SingleReturnCompat (Com Expr)

@[simp_denote] def cst (n : ℤ) : Expr ScfArith Γ .pure .int  :=
  Expr.mk
    (op := .coe <| .const n)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

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

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

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

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

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

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

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

namespace ForAddToMul

def lhs (vincrement : ℤ) : Com ScfArith ⟨[  .nat,   .int]⟩ .impure .int :=
    Com.letPure (cst 0) <|
    Com.letPure  (cst 1) <|
    Com.var (for_ (t := .int)
                        ⟨  1, rfl⟩
                        ⟨  0, rfl⟩
                        ⟨  2, rfl⟩
                        ⟨  3, rfl⟩ (
      Com.letPure (cst vincrement) <|
      Com.letPure (add ⟨0, rfl⟩ ⟨2, rfl⟩)
      <| Com.ret ⟨0, rfl⟩)) <|
  Com.ret ⟨0, rfl⟩

def rhs (vincrement : ℤ) : Com ScfArith ⟨[  .nat,   .int]⟩ .pure .int :=
  Com.var (cst vincrement) <|
  Com.var (axpy ⟨0, rfl⟩ ⟨1, rfl⟩ ⟨2, rfl⟩) <|
  Com.ret ⟨0, rfl⟩

open Scf in

open Arith in

theorem correct : Com.denote (lhs v0) = Com.denote (rhs v0) :=

:= by
  unfold lhs rhs
  simp_peephole
  intros A B
  rw [Scf.LoopBody.counterDecorator.const_index_fn_iterate (f' := fun v => v0 + v)] <;> try rfl
  simp only [add_left_iterate, nsmul_eq_mul, Int.mul_comm]


syntax "llvm.and"     : 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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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}"

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

Lets

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

Expr.decidableEq

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

Var

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

Ctxt

structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty

  deriving Repr, Lean.ToExpr, DecidableEq

Valuation.nil

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

EffectKind

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

Com

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 β

FlatCom

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

HVector.decidableEqReg

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

Var

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

Hom.with

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

Hom

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

TyDenote

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

Valuation.instAppendHVector

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

map

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

Com.outContextHom

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

Com.outContext

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

Com.rec'

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

Valuation.cast

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

Com.outContextDiff

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

Expr.outContext

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

Expr.bvars

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

returnTypes

def returnTypes  := Signature.returnTypes ∘ s.signature

Signature

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

DialectSignature

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

Com.bvars

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

Diff

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

Diff.Valid

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

toHom

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

castCtxt

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

Com.denoteLets

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

DialectDenote

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

Lets.denote

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

sig

def sig          := Signature.sig ∘ s.signature

regSig

def regSig       := Signature.regSig ∘ s.signature

RegionSignature

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

Ty

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

effectKind

def effectKind   := Signature.effectKind ∘ s.signature

Dialect

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

Op

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

HVector.denote

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)

FlatCom.denote

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

RegionSignature.map

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

Signature.map

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
  }

Expr.denote

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

Expr.op

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

liftEffect

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

toMonad

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

Com.denote

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

Com.ty

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

Valuation.cons

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

cast

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

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

Com.castPureToEff

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

Com.changeEffect

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

Expr.changeEffect

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

Expr.args

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

Expr.regArgs

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

Regions

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

com

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

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
  }]

Region

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

MetaLLVM

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

Ty

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

Op

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

MOp

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

binary

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

MOp.BinaryOp

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

LLVM

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

MTy

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

Width

abbrev Width φ := ConcreteOrMVar Nat φ

ConcreteOrMVar

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

const

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

MOp.UnaryOp

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

select

@[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?

IntW

def IntW w := PoisonOr <| BitVec w

PoisonOr

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

icmp

@[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'

icmp?

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

IntPred

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

icmp'

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

mkCom

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⟩

FlatCom.denoteLets

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

Com.toLets

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

Lets.castPureToEff

def Lets.castPureToEff (eff : EffectKind) : Lets d Γ_in .pure Γ_out → Lets d Γ_in eff Γ_out
  | .nil => .nil
  | .var body e => .var (body.castPureToEff eff) (e.castPureToEff eff)

Expr.castPureToEff

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

Expr.returnVars

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

ofFin

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

Com.returnVars

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

Valuation.comap

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

map

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)

HVectorLiteral

structure HVectorLiteral where
  u : Level
  v : Level
  α : Q(Type $u)
  A : Q($α → Type $v)
  elems : Array ((a : Q($α)) × Q($A $a))

Lets.ValidDenotation

def Lets.ValidDenotation (lets : Lets d Γ_in eff Γ_out) :=
  { V // ∀ {t ts} {v : Var _ t} {w : Var ⟨ts⟩ t} {e} ,
          lets.getPureExpr v = some ⟨ts, w, e⟩
          → (e.pdenoteOp V)[w] = V v }

Lets.denoteIntoSubtype

def Lets.denoteIntoSubtype (lets : Lets d Γ_in eff Γ_out) (Γv : Valuation Γ_in) :
    eff.toMonad d.m lets.ValidDenotation :=
  match lets with
    | .nil => return ⟨Γv, by admit /- proof elided -/
    ⟩
    | @Lets.var _ _ _ _ Γ_out eTy body e => do
        let ⟨Vout, h⟩ ← body.denoteIntoSubtype Γv
        let Ve ← e.denoteOpIntoSubtype Vout
        return ⟨Ve.val ++ Vout, by admit /- proof elided -/
        ⟩

Expr.denoteOpIntoSubtype

def Expr.denoteOpIntoSubtype (e : Expr d Γ_in eff tys) (Γv : Valuation Γ_in) :
    eff.toMonad d.m {x // e.IsDenotationForPureE Γv x} :=
  match h_pure : e.toPure? with
    | some ePure => pure ⟨ePure.denoteOp Γv, by admit /- proof elided -/
    ⟩
    | none => (Subtype.mk · (by admit /- proof elided -/
    )) <$> (e.denoteOp Γv)

Expr.IsDenotationForPureE

abbrev Expr.IsDenotationForPureE (e : Expr d Γ eff tys) (Γv : Valuation Γ)
    (x : HVector toType tys) : Prop :=
  ∀ (ePure : Expr d Γ .pure tys), e.toPure? = some ePure → ePure.denoteOp Γv = x

appendCases

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

Expr.toPure?

def Expr.toPure? (e : Expr d Γ eff ty) : Option (Expr d Γ .pure ty) :=
  if h : e.HasPureOp then
    some <| e.toPure h
  else
    none

Expr.toPure

def Expr.toPure (e : Expr d Γ eff ty) (h : e.HasPureOp) : Expr d Γ .pure ty :=
  ⟨e.op, e.ty_eq, EffectKind.le_of_eq h, e.args, e.regArgs⟩

Expr.ty

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

Expr.HasPureOp

def Expr.HasPureOp (e : Expr d Γ eff ty) : Prop :=
  DialectSignature.effectKind e.op = .pure

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

TyDenote.toType

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

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

theorem Com.denoteLets_eq {com : Com d Γ eff t} : com.denoteLets = com.toLets.denote

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

@[simp, grind=] theorem castCtxt_rfl (h : Γ = Γ) : v.castCtxt h = v

@[simp] theorem Com.returnVars_castPureToEff (eff : _) (com : Com d Γ .pure tys) :
    (com.castPureToEff eff).returnVars = com.returnVars.map (fun _ v => v.castCtxt (by simp))

@[simp] theorem Valuation.comap_apply {Γi Γo : Ctxt Ty}
    (V : Γi.Valuation) (f : Ctxt.Hom Γo Γi) (v : Γo.Var t) :
    V.comap f v = V (f v)

@[simp] theorem Com.denoteLets_castPureToEff {com : Com d Γ .pure ty} :
    denoteLets (com.castPureToEff eff)
    = fun V => pure (com.denoteLets V |>.comap fun _ v => v.castCtxt (by simp))

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

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

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

@[simp] theorem Ctxt.Valuation.comap_outContextHom_denoteLets {com : Com d Γ .pure ty} {V} :
    Valuation.comap (com.denoteLets V) com.outContextHom = V

@[simp] theorem Com.bvars_castPureToEff {com : Com d Γ .pure ty} :
    (com.castPureToEff eff).bvars = com.bvars

@[simp] theorem Valuation.comap_with [DecidableEq Ty] {Γ Δ : Ctxt Ty}
    {V : Valuation Γ} {map : Δ.Hom Γ} {vs : HVector Δ.Var ty} {ws : HVector Γ.Var ty} :
    V.comap (map.with vs ws) = (V.comap map).reassignVars vs (ws.map V)

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

@[simp, grind=] theorem castCtxt_castCtxt (h₁ : Γ = Δ) (h₂ : Δ = Ξ) :
    (v.castCtxt h₁).castCtxt h₂ = v.castCtxt (by simp [*])

theorem Lets.denote_eq_denoteIntoSubtype (lets : Lets d Γ_in eff Γ_out) (Γv : Valuation Γ_in) :
    lets.denote Γv = Subtype.val <$> (lets.denoteIntoSubtype Γv)

theorem Expr.denoteOp_eq_denoteOpIntoSubtype (e : Expr d Γ eff tys) (V : Valuation Γ) :
    e.denoteOp V = Subtype.val <$> e.denoteOpIntoSubtype V

Zipper.insertPureCom

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

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

import LeanMLIR.Framework.Basic

import LeanMLIR.Transforms.Rewrite.Match

open Ctxt (Valuation Var Hom)

variable (d : Dialect) [DialectSignature d]

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

namespace Zipper

variable {d}

section Denote

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

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

end Denote

section ToCom

variable {Γ_mid}

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

end ToCom

section InsertCom

variable [DecidableEq d.Ty]

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 }

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)

section Lemmas

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

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 :=

:= by
  simp only [denote_insertPureCom, Valuation.comap_with,
  Valuation.comap_outContextHom_denoteLets, Com.denoteLets_returnVars]
  unfold Zipper.denote
  simp [Lets.denote_eq_denoteIntoSubtype, h]


theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

R.fromTensorFinsupp

theorem f_deg_eq : (f q n).degree = 2^n

theorem f_monic : Monic (f q n)

theorem R.fromPoly_rep'_eq_ideal :
    forall a : (ZMod q)[X],
      ∃ i ∈ Ideal.span {f q n}, (R.fromPoly (n:=n) a).representative' = a + i

theorem R.representative_fromPoly :
    forall a : (ZMod q)[X], (R.fromPoly (n:=n) a).representative = a %ₘ (f q n)

@[simp]
theorem R.representative_fromPoly_eq (x : (ZMod q)[X]) (DEGREE: x.degree < (f q n).degree) :
   R.representative q n (R.fromPoly (n:=n) x) = x

theorem Polynomial.degree_toFinsupp [Semiring M] [DecidableEq M]
  (xs : List M) :
  degree { toFinsupp := List.toFinsupp (l := xs) } ≤ List.length xs

theorem R.fromTensorFinsupp_degree (q : Nat) (coeffs : List Int):
  (R.fromTensorFinsupp q coeffs).degree ≤ coeffs.length

theorem R.fromTensorFinsupp_coeffs (coeffs : List Int) :
  Polynomial.coeff (fromTensorFinsupp q coeffs) i = ↑(List.getD coeffs i 0)

theorem R.fromTensorFinsupp_concat_monomial (c : Int) (cs : List Int) :
    (R.fromTensorFinsupp q (cs ++ [c])) =
      (R.fromTensorFinsupp q cs) +
        (Polynomial.monomial cs.length (Int.cast c : (ZMod q)))

theorem R.fromTensor_eq_fromTensorFinsupp_fromPoly {coeffs : List Int} :
    R.fromTensor (q := q) (n := n) coeffs =
  R.fromPoly (q := q) (n := n) (R.fromTensorFinsupp q coeffs)

import Mathlib.RingTheory.Polynomial.Quotient

import Mathlib.RingTheory.Ideal.Defs

import Mathlib.RingTheory.Ideal.Basic

import Mathlib.Data.ZMod.Defs

import Mathlib.Data.ZMod.Basic

import Mathlib.Algebra.MonoidAlgebra.Basic

import Mathlib.Algebra.Polynomial.RingDivision

import Mathlib.Data.Finset.Sort

import Mathlib.Data.List.ToFinsupp

import Mathlib.Data.List.Basic

import Mathlib.Tactic.Cases

import LeanMLIR.Framework

open Polynomial -- for R[X] notation

section CommRing

variable (q t : Nat) [Fact (q > 1)] (n : Nat)

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

abbrev R := (ZMod q)[X] ⧸ (Ideal.span {f q n})

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

end CommRing

section Representative

variable (q t n : Nat)

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

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

variable [Fact (q > 1)]

end Representative

section Coeff

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

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

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

end Coeff

section FinnSupp

variable {q n : Nat}

noncomputable def R.fromTensorFinsupp (q : Nat) (coeffs : List Int) : (ZMod q)[X] :=
  Polynomial.ofFinsupp (List.toFinsupp (coeffs.map Int.cast))

end FinnSupp

section Tensor

variable {q n : Nat} [Fact (q > 1)]

theorem R.coeff_fromTensor (tensor : List Int)
    (htensorlen : tensor.length < 2^n) :
    (R.fromTensor (q := q) (n := n) tensor).coeff i = (tensor.getD i 0) :=

:= by
  rw [fromTensor_eq_fromTensorFinsupp_fromPoly]
  have hfromTensorFinsuppDegree := fromTensorFinsupp_degree q tensor
  rw [coeff, representative_fromPoly_eq]
  apply fromTensorFinsupp_coeffs
  case DEGREE =>
    generalize htensor_degree : degree (fromTensorFinsupp q tensor) = tensor_degree
    rw [f_deg_eq]
    cases tensor_degree
    case bot =>
      rw [WithBot.lt_def]
      simp
      exists 2^n
    norm_cast
    case coe tensor_degree =>
      /- I hate this coercion stuff -/
      apply WithBot.coe_strictMono
      norm_cast
      have htrans : tensor_degree  ≤ List.length tensor := by
        rw [htensor_degree] at hfromTensorFinsuppDegree
        rw [← WithBot.coe_le_coe]
        assumption
      apply Nat.lt_of_le_of_lt htrans htensorlen