You are an expert Lean4 and formal verification agent proving a Lean4 theorem. Your output **must** strictly adhere to the specified format.

### Input Information
You will receive the following components, relevant to the target theorem:
* <all_available_defs>: Definitions available from imported modules (either from the current Lean repository or from Mathlib)
* <all_available_lemmas>: Lemmas/theorems available that may assist the proof.
* <local_ctx>...</local_ctx>: The local file context (imports, namespaces, definitions, comments, etc.) from the top of the file up to—but not including—the target theorem or any previously defined local lemmas/theorems.
* <target_theorem>...</target_theorem>: The target theorem statement, ending with `:=`.

**Notes**
1. Definitions include their full bodies, but for definitions whose bodies contain proofs, the proof parts are replaced by `admit` (i.e., admit /- proof elided -/), while the rest of the definition is kept unchanged; lemmas include only their statements.
2. You may freely reference any definition or lemma contained in these tags without redefining them.
3. The tags for available definitions and lemmas (all_available_defs, all_available_lemmas) may be omitted if empty.

### Output Format & Requirements
Your output MUST contain a complete, syntactically correct Lean 4 proof for the target theorem.

**Strict Output Structure:**
<reasoning> 
...
</reasoning> 

<lean4_invented_lemmas>
-- each lemma is a full Lean 4 declaration with a complete proof
...
</lean4_invented_lemmas>

<lean4_proof>
-- Lean 4 proof body of the target theorem
...
</lean4_proof>

**Requirements:**
1. <reasoning> and <lean4_invented_lemmas> are OPTIONAL. Include them only if they aid proof construction.
2. <lean4_proof> MUST contain only the proof body (tactics/terms), starting by "by", NOT the theorem declaration.
3. <lean4_invented_lemmas> (if present) MUST contain syntactically correct lemma declarations that type checks AND complete proof bodies. Defining a lemma using `axiom` is NOT allowed.
4. <reasoning> contains high-level proof planning.
5. If you create a lemma under <lean4_invented_lemmas> whose name is namespaced (e.g., X.Y.Z), you MUST always refer to it using its fully qualified name. Do NOT rely on shortened names. For example, write rw [X.Y.Z], not rw [Z] for invented lemmas.
6. All Lean 4 code MUST be syntactically correct. No `sorry`, `admit`, or non-Lean text. Comments allowed using `--` or `/- ... -/`.
7. Follow the tags and structure EXACTLY. No extra text before, between, or after the sections.
8. Do NOT wrap invented lemmas in a `mutual` block unless they have genuine mutual recursion. Declare each lemma independently.

### Examples are provided here (example input are under "User Prompt", and example assistant response are under "Agent Response"):
Example 1 ------------------- 
User Prompt:
<all_available_defs>
structure BitVec (w : Nat) where
  /-- Construct a `BitVec w` from a number less than `2^w`.
  O(1), because we use `Fin` as the internal representation of a bitvector. -/
  ofFin ::

  /-- Interpret a bitvector as a number less than `2^w`.
  O(1), because we use `Fin` as the internal representation of a bitvector. -/
  toFin : Fin (hPow 2 w)
</all_available_defs>

<all_available_lemmas>
theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) : (l₁ ++ l₂)[i] = l₁[i]
theorem getElem_eq_testBit_toNat (x : BitVec w) (i : Nat) (h : i < w) :
  x[i] = x.toNat.testBit i := rfl
@[simp, grind =] theorem toNat_cons (b : Bool) (x : BitVec w) :
    (cons b x).toNat = (b.toNat <<< w) ||| x.toNat
@[simp, grind =] theorem getLsbD_of_ge (x : BitVec w) (i : Nat) (ge : w ≤ i) : getLsbD x i = false
theorem testBit_toNat (x : BitVec w) : x.toNat.testBit i = x.getLsbD i
</all_available_lemmas>

<local_ctx>
import Mathlib.Computability.NFA
import Mathlib.Data.FinEnum
import Mathlib.Data.Rel
import Mathlib.Data.Vector.Basic
import Blase.AutoStructs.ForLean

open Set
open Mathlib
open SetRel

-- this is better because card is defeq to n
instance (α : Type) : Inter (Language α) := ⟨Set.inter⟩
instance (α : Type) : Union (Language α) := ⟨Set.union⟩

/--
The set of `n`-tuples of bit vectors of an arbitrary width.
-/
structure BitVecs (n : Nat) where
  w : Nat
  bvs : List.Vector (BitVec w) n

def BitVecs.cast (bvs : BitVecs n) (h : n = n') : BitVecs n' :=
  { w := bvs.w, bvs := h ▸ bvs.bvs }

abbrev BitVecs.empty : BitVecs n := ⟨0, List.Vector.replicate n .nil⟩
abbrev BitVecs.singleton {w : Nat} (bv : BitVec w) : BitVecs 1 := ⟨w, bv ::ᵥ .nil⟩
abbrev BitVecs.pair {w : Nat} (bv1 bv2 : BitVec w) : BitVecs 2 := ⟨w, bv1 ::ᵥ bv2 ::ᵥ .nil⟩

/--
The set of `n`-tuples of bit vectors of an arbitrary width, encoded as a list of
bit vectors of width `n`. The width of the encoded bit vectors is the length of
the list.
-/
abbrev BitVecs' (n : Nat) := List (BitVec n)

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

@[simp]
</local_ctx>

<target_theorem>
lemma dec_snoc n (bvs' : BitVecs' n) (a : BitVec n) :
  dec (bvs' ++ [a]) =
    { w := bvs'.length + 1
      bvs := List.Vector.ofFn fun k =>
        BitVec.cons (a.getLsbD k) ((dec bvs').bvs.get k) } :=
</target_theorem>

Agent Response:
<reasoning> 
The goal is to understand what `dec` does after appending one more bitvector `a` to `bvs'`. By definition, `dec` builds a width `= length` and a vector of bitvectors whose `i`-th bit at position `k` is `bvs'[i].getLsbD k`. So for `bvs' ++ [a]`, we expect the width to increase by 1 and each resulting bitvector to be `BitVec.cons` of the new bit coming from `a`.

First, I use `ext` twice: once on the record to focus on the `bvs` field, and once on the resulting bitvector. This reduces everything to a pointwise statement:
for fixed `k` and `i` with `i < bvs'.length + 1`,
the bit `(bvs' ++ [a])[i].getLsbD k` should match the bit at index `i` of `BitVec.cons (a.getLsbD k) ((dec bvs').bvs.get k)`.

At this point I need a clean description of indexing into `BitVec.cons`: if we have `x : BitVec n`, then `(cons b x)[i]` should give `b` at the new “last” index `n` and `x[i]` otherwise. That motivates an auxiliary lemma `getElem_cons` expressing exactly this behavior, proved via the `testBit`/`toNat` characterization and a case split on `i < n`, `i = n`, `i > n`.

Using `getElem_cons`, the RHS becomes an `if` on whether `i = bvs'.length`. So I split into cases:
- If `i = bvs'.length`, then `(bvs' ++ [a])[i]` is the appended `a`, so both sides reduce to `a.getLsbD k`.
- If `i ≠ bvs'.length`, the bound `i < bvs'.length + 1` forces `i < bvs'.length`. Then `(bvs' ++ [a])[i] = bvs'[i]` by `List.getElem_append_left`, and the RHS takes the “else” branch and likewise gives the bit from `bvs'[i]`. In both subcases, `simp` (plus the definition of `dec`) finishes.

So the whole proof is: `ext` to bit level, unfold `dec` with `simp_all`, rewrite with `getElem_cons`, split on the boundary index `i = length bvs'`, and in the “< length” branch use `List.getElem_append_left` and bounds to line up both sides.
</reasoning> 

<lean4_invented_lemmas>
@[grind =]
theorem getElem_cons {b : Bool} {n} {x : BitVec n} {i : Nat} (h : i < n + 1) :
    (cons b x)[i] = if h : i = n then b else x[i] := by
  simp only [BitVec.getElem_eq_testBit_toNat, BitVec.toNat_cons, Nat.testBit_or]
  rw [Nat.testBit_shiftLeft]
  rcases Nat.lt_trichotomy i n with i_lt_n | i_eq_n | n_lt_i
  · have p1 : ¬(n ≤ i) := by omega
    have p2 : i ≠ n := by omega
    simp [p1, p2]
  · simp only [i_eq_n, ge_iff_le, Nat.le_refl, decide_true, Nat.sub_self, Nat.testBit_zero,
    Bool.true_and, BitVec.testBit_toNat, BitVec.getLsbD_of_ge, Bool.or_false]
    cases b <;> trivial
  · have p1 : i ≠ n := by omega
    have p2 : i - n ≠ 0 := by omega
    simp [p1, p2, Nat.testBit_bool_toNat]
</lean4_invented_lemmas>

<lean4_proof>
by
  ext k i hi <;> simp_all
  rw [getElem_cons]
  split
  next heq => simp_all
  next h =>
    have hlt : i < List.length bvs' := by simp at hi; omega
    rw [List.getElem_append_left hlt]
    simp
</lean4_proof>

### Your task:
<all_available_defs>
DecidableEq in Init.Prelude
Lean.ToExpr in Lean.ToExpr
List in Init.Prelude
Repr in Init.Data.Repr
Nat in Init.Prelude
Fin in Init.Prelude
Option in Init.Prelude
reduceDIte in Lean.Meta.Tactic.Simp.BuiltinSimprocs.Core
ZMod in Mathlib.Data.ZMod.Defs
Int in Init.Data.Int.Basic
BitVec in Init.Prelude
String in Init.Prelude
Id in Init.Control.Id
Vector in Init.Data.Vector.Basic
StateM in Init.Control.State
Inhabited in Init.Prelude
Lean.Expr in Lean.Expr
Stream' in Mathlib.Data.Stream.Defs
Unit in Init.Prelude
Stream'.tail in Mathlib.Data.Stream.Defs
Id.run in Init.Control.Id
Bool in Init.Prelude
bool in Init.Control.Basic
List.replicate in Init.Data.List.Basic
BitVec.ofInt in Init.Data.BitVec.Basic
inferInstanceAs in Init.Prelude
Hashable in Init.Prelude
id in Init.Prelude
Std.HashMap in Std.Data.HashMap.Basic

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

def idxOf? (x : A a) {as} [DecidableEq α] [∀ a, DecidableEq (A a)] :
    HVector A as → Option { i : Fin <| as.length // as.get i = a }
  | .nil => none
  | .cons (a:=b) y ys =>
      if h : ∃ h : a = b, x = h ▸ y then
        some ⟨(0 : Fin <| _ + 1), h.1 ▸ rfl⟩
      else
        (ys.idxOf? x).map fun ⟨i, h⟩ =>
          ⟨i.succ, by admit /- proof elided -/
          ⟩

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)

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

instance : TyDenote Unit where toType := fun _ => Unit

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

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

instance : TyDenote ISL.Ty where
  toType := fun
    | .regIndex => RegIndex
    | .bits w => BitVec w

def_denote for ISL
  | .regConst r   => [r]ₕ
  | .regRead      => fun r regFile => ([regFile.read r]ₕ, regFile)
  | .regWrite     => fun r v regFile => ([]ₕ, regFile.write r v)
  | .bitsConst x  => [x]ₕ
  | .bitsAdd _    => fun (x y : BitVec _) => [x + y]ₕ

def ExplicitISL : Dialect where
  Op := ISLOp
  Ty := ExpTy
  m := Id

inductive ExpTy where
  | isl : ISL.Ty → ExpTy
  | regFile

def ISL : Dialect where
  Ty := ISLTy
  Op := ISLOp
  m := StateM RegFile

inductive ISLTy
  | regIndex
  | bits (w : Nat)
  deriving DecidableEq, ToExpr, Repr

inductive ISLOp
  | regConst (r : RegIndex)
  | regRead
  | regWrite
  | bitsConst {w : Nat} (x : BitVec w)
  | bitsAdd (w : Nat)
  deriving DecidableEq, ToExpr, Repr

def RegIndex := Fin RegFile.numRegisters

def RegFile.numRegisters := 32

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

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

structure RegFile where
  regs : Vector (BitVec RegFile.registerWidth) RegFile.numRegisters

def RegFile.registerWidth := 64

def read (self : RegFile) (r : RegIndex) : BitVec registerWidth :=
  self.regs[r.val]

def write (self : RegFile) (r : RegIndex) (v : BitVec registerWidth) : RegFile :=
  RegFile.mk <| self.regs.set r.val v

instance instDCTyDenote : TyDenote Ty where
toType := fun
| Ty.tokenstream => CIRCTStream.DCPlusOp.TokenStream
| Ty.tokenstream2 => CIRCTStream.DCPlusOp.TokenStream × CIRCTStream.DCPlusOp.TokenStream
| Ty.valuestream w => CIRCTStream.DCPlusOp.ValueStream (BitVec w)
| Ty.valuestream2 w => CIRCTStream.DCPlusOp.ValueStream (BitVec w) × CIRCTStream.DCPlusOp.ValueStream (BitVec w)
| Ty.valuetokenstream w => CIRCTStream.DCPlusOp.ValueStream (BitVec w) × CIRCTStream.DCPlusOp.TokenStream
| Ty.variadicvaluetokenstream w => CIRCTStream.DCPlusOp.VariadicValueStream w × CIRCTStream.DCPlusOp.TokenStream

def_denote for DCPlus where
  | .fst => fun s => [s.fst]ₕ
  | .fstVal _ => fun s => [s.fst]ₕ
  | .fstValPure _ => fun s => [s.fst]ₕ
  | .fstVal' _ => fun s => [s.fst.mapOpt (·[0]?)]ₕ
  | .snd => fun s => [s.snd]ₕ
  | .sndValPure _ => fun s => [s.snd]ₕ
  | .pair _ => fun s₁ s₂ => [(s₁, s₂)]ₕ
  | .sndVal _ => fun s => [s.snd]ₕ
  | .sndVal' _ => fun s => [s.fst.mapOpt (·[0]?)]ₕ
  | .tokVal' _ => fun s => [s.snd]ₕ
  | .merge => fun s₁ s₂ => [CIRCTStream.DCPlusOp.merge s₁ s₂]ₕ
  | .branch => fun s₁ s₂ => [CIRCTStream.DCPlusOp.branch s₁ s₂]ₕ
  | .fork => fun s => [CIRCTStream.DCPlusOp.fork s]ₕ
  | .forkVal => fun s => [CIRCTStream.DCPlusOp.forkVal s]ₕ
  | .join => fun s₁ s₂ => [CIRCTStream.DCPlusOp.join s₁ s₂]ₕ
  | .mux => fun s₁ s₂ c => [CIRCTStream.DCPlusOp.mux s₁ s₂ c]ₕ
  | .muxVal => fun s₁ s₂ c => [CIRCTStream.DCPlusOp.muxVal s₁ s₂ c]ₕ
  | .sink => fun s => [CIRCTStream.DCPlusOp.sink s]ₕ
  | .source => [CIRCTStream.DCPlusOp.source]ₕ
  | .sourceOnes => [CIRCTStream.DCPlusOp.sourceOnes]ₕ
  | .cMerge => fun s₁ s₂ => [CIRCTStream.DCPlusOp.cMerge s₁ s₂]ₕ
  | .supp => fun s₁ s₂ => [CIRCTStream.DCPlusOp.supp s₁ s₂]ₕ
  | .not => fun s₁ => [CIRCTStream.DCPlusOp.not s₁]ₕ

abbrev DCPlus : Dialect where
  Op := Op
  Ty := Ty

def_signature for DCPlus where
  | .fst => (Ty.tokenstream2) → (Ty.tokenstream)
  | .fstVal t => (Ty.valuetokenstream t) → Ty.valuestream t
  | .fstVal' t => (Ty.variadicvaluetokenstream t) → Ty.valuestream t
  | .fstValPure t => (Ty.valuestream2 t) → Ty.valuestream t
  | .snd => (Ty.tokenstream2) → (Ty.tokenstream)
  | .sndValPure t => (Ty.valuestream2 t) → Ty.valuestream t
  | .pair w => (Ty.valuestream w, Ty.valuestream w) → Ty.valuestream2 w
  | .sndVal t => (Ty.valuetokenstream t) → Ty.tokenstream
  | .sndVal' t => (Ty.variadicvaluetokenstream t) → Ty.valuestream t
  | .tokVal' t => (Ty.variadicvaluetokenstream t) → Ty.tokenstream
  | .merge => (Ty.tokenstream, Ty.tokenstream) → Ty.valuestream 1
  | .branch => (Ty.valuestream 1, Ty.tokenstream) → Ty.tokenstream2
  | .fork => (Ty.tokenstream) → Ty.tokenstream2
  | .forkVal => (Ty.valuestream 1) → Ty.valuestream2 1
  | .join => (Ty.tokenstream, Ty.tokenstream) → Ty.tokenstream
  | .mux => (Ty.tokenstream, Ty.tokenstream, Ty.valuestream 1) → Ty.tokenstream
  | .muxVal => (Ty.valuestream 1, Ty.valuestream 1, Ty.valuestream 1) → Ty.valuestream 1
  | .sink => (Ty.tokenstream) → Ty.tokenstream
  | .source => () → Ty.tokenstream
  | .sourceOnes => () → Ty.valuestream 1
  | .cMerge => (Ty.tokenstream, Ty.tokenstream) → Ty.valuetokenstream 1
  | .supp => (Ty.valuestream 1, Ty.tokenstream) → Ty.tokenstream
  | .not => (Ty.valuestream 1) → Ty.valuestream 1

def sndVal {Γ} (a : Γ.Var (MLIR2DCPlus.Ty.valuetokenstream r)) : Expr (DCPlus) Γ .pure (.tokenstream) :=
    Expr.mk
    (op := .sndVal r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

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

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

def Context (Ty : Type) : Type :=
  List (VarName × Ty)

def VarName : Type := String

def hasType (Γ : Context Ty) (v : VarName) (ty : Ty) : Prop :=
  Γ.lookup v = some ty

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

inductive AffineExpr
  | Var: String -> AffineExpr
  deriving DecidableEq, Repr

def mux {Γ : Ctxt _}  (a b : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) (c : Γ.Var (MLIR2DCPlus.Ty.valuestream 1)) : Expr (DCPlus) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .mux)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .cons c <| .nil)
    (regArgs := .nil)

def branch {Γ : Ctxt _} (c : Γ.Var (MLIR2DCPlus.Ty.valuestream 1)) (a : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) : Expr (DCPlus) Γ .pure (.tokenstream2) :=
  Expr.mk
    (op := .branch)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons c <| .cons a <| .nil)
    (regArgs := .nil)

def merge {Γ : Ctxt _}  (a b : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) : Expr (DCPlus) Γ .pure (.valuestream 1) :=
  Expr.mk
    (op := .merge)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def cMerge {Γ : Ctxt _} (a b : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) : Expr (DCPlus) Γ .pure (.valuetokenstream 1) :=
  Expr.mk
    (op := .cMerge)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def fork {Γ : Ctxt _}  (a : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) : Expr (DCPlus) Γ .pure (.tokenstream2) :=
  Expr.mk
    (op := .fork)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def fstVal (a : Γ.Var (MLIR2DCPlus.Ty.valuetokenstream r)) : Expr (DCPlus) Γ .pure (.valuestream r) :=
    Expr.mk
    (op := .fstVal r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def forkVal {Γ : Ctxt _}  (a : Γ.Var (MLIR2DCPlus.Ty.valuestream 1)) : Expr (DCPlus) Γ .pure (.valuestream2 1) :=
  Expr.mk
    (op := .forkVal)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def join {Γ : Ctxt _}  (a b : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) : Expr (DCPlus) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .join)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

inductive Op
| fst
| snd
| pair (w : Nat)
| fstVal (w : Nat)
| fstValPure (w : Nat)
| sndVal (w : Nat)
| sndValPure (w : Nat)
| fstVal' (w : Nat)
| sndVal' (w : Nat)
| tokVal' (w : Nat)
| fork
| forkVal
| join
| merge
| mux
| muxVal
| cMerge
| branch
| source
| sourceOnes
| sink
| supp
| not
deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

def fstValPure (a : Γ.Var (MLIR2DCPlus.Ty.valuestream2 r)) : Expr (DCPlus) Γ .pure (.valuestream r) :=
    Expr.mk
    (op := .fstValPure r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def sink {Γ : Ctxt _} (a : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) : Expr (DCPlus) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .sink)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def supp {Γ : Ctxt _} (a : Γ.Var (MLIR2DCPlus.Ty.tokenstream)) (c : Γ.Var (MLIR2DCPlus.Ty.valuestream 1)) : Expr (DCPlus) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .supp)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons c <| .cons a <| .nil)
    (regArgs := .nil)

def muxVal {Γ : Ctxt _}  (a b c : Γ.Var (MLIR2DCPlus.Ty.valuestream 1)) : Expr (DCPlus) Γ .pure (.valuestream 1) :=
  Expr.mk
    (op := .muxVal)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .cons c <| .nil)
    (regArgs := .nil)

def sndValPure (a : Γ.Var (MLIR2DCPlus.Ty.valuestream2 r)) : Expr (DCPlus) Γ .pure (.valuestream r) :=
    Expr.mk
    (op := .sndValPure r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def sourceOnes : Expr (DCPlus) Γ .pure (.valuestream 1) :=
  Expr.mk
    (op := .sourceOnes)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

def snd {Γ} (a : Γ.Var (MLIR2DCPlus.Ty.tokenstream2)) : Expr (DCPlus) Γ .pure (.tokenstream) :=
    Expr.mk
    (op := .snd)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def source : Expr (DCPlus) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .source)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

def fst (a : Γ.Var (MLIR2DCPlus.Ty.tokenstream2)) : Expr (DCPlus) Γ .pure (.tokenstream) :=
    Expr.mk
    (op := .fst)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def TokenStream := Stream Unit

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

def ValueStream := Stream

def VariadicValueStream (w : Nat) := CIRCTStream.Stream (List (BitVec w))

def branch (c : ValueStream (BitVec 1)) (x : TokenStream) : TokenStream × TokenStream  :=
  Stream.corec₂ (β := ValueStream (BitVec 1) × TokenStream) (c, x) fun ⟨c, x⟩ =>
    Id.run <| do
      match c 0 with
        | none => (none, none, (c.tail, x))
        | some x₀ =>
          if x₀.msb then
            (some (), none, (c.tail, x.tail))
          else
            (none, some (), (c.tail, x.tail))

def tail : Stream α → Stream α := Stream'.tail

def corec₂ {β} (s0 : β) (f : β → (Option α × Option γ × β)) : Stream α × Stream γ :=
  let f' := fun b =>
    let x := f b
    (x.fst, x.snd.fst)
  let g := (f · |>.snd.snd)
  let x := Stream'.corec f' g s0
  (
    fun i => (x i).fst,
    fun i => (x i).snd,
  )

def cMerge (x y : TokenStream) : ValueStream (BitVec 1) × TokenStream :=
  Stream.corec₂ (β := TokenStream × TokenStream) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some x', some _ => (some 1, some x', (x.tail, y))
    | some x', none => (some 1, some x', (x.tail, y.tail))
    | none, some y' => (some 0, some y', (x.tail, y.tail))
    | none, none => (none, none, (x.tail, y.tail))

def fork (x : TokenStream) : TokenStream × TokenStream  :=
  Stream.corec₂ (β := TokenStream) x
    fun x => Id.run <| do
      (x 0, x 0, x.tail)

def forkVal (x : ValueStream (BitVec 1)) : ValueStream (BitVec 1) × ValueStream (BitVec 1)  :=
  Stream.corec₂ (β := ValueStream (BitVec 1)) x
    fun x => Id.run <| do
      (x 0, x 0, x.tail)

def join (x y : TokenStream) : TokenStream  :=
  Stream.corec (β := TokenStream × TokenStream) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some _, some _ => (some (), (x.tail, y.tail))
    | some _, none => (none, (x, y.tail))
    | none, some _ => (none, (x.tail, y))
    | none, none => (none, (x.tail, y.tail))

def corec {α} {β} (s0 : β) (f : β → (Option α × β)) : Stream α :=
  Stream'.corec (f · |>.fst) (f · |>.snd) s0

def merge (x y : TokenStream) : ValueStream (BitVec 1) :=
  Stream.corec (β := TokenStream × TokenStream) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some _, some _ => (some 1, (x.tail, y))
    | some _, none => (some 1, (x.tail, y.tail))
    | none, some _ => (some 0, (x.tail, y.tail))
    | none, none => (none, (x.tail, y.tail))

def mux (x y : TokenStream) (c : ValueStream (BitVec 1)): TokenStream :=
  Stream.corec (β := TokenStream × TokenStream × ValueStream (BitVec 1)) (x, y, c)
  fun ⟨x, y, c⟩ =>
    match (c 0) with
    | none => (none, x, y, c.tail) 
    | some 1#1 =>
      match (x 0) with
      | none => (none, x.tail, y, c) 
      | some _ => (some (), x.tail, y, c.tail) 
    | some 0#1 =>
      match (y 0) with
      | none => (none, x, y.tail, c) 
      | some _ => (some (), x, y.tail, c.tail)

def muxVal (x y c : ValueStream (BitVec 1)): ValueStream (BitVec 1) :=
  Stream.corec (β := ValueStream (BitVec 1) × ValueStream (BitVec 1) × ValueStream (BitVec 1)) (x, y, c)
  fun ⟨x, y, c⟩ =>
    match (c 0) with
    | none => (none, x, y, c.tail) 
    | some 1#1 =>
      match (x 0) with
      | none => (none, x.tail, y, c) 
      | some e => (some e, x.tail, y, c.tail) 
    | some 0#1 =>
      match (y 0) with
      | none => (none, x, y.tail, c) 
      | some e => (some e, x, y.tail, c.tail)

def not (c : ValueStream (BitVec 1)) : (ValueStream (BitVec 1)) :=
  Stream.corec (β := ValueStream (BitVec 1)) c fun c =>
  match c 0 with
  | some 1 => (some 0, c.tail)
  | some 0 => (some 1, c.tail)
  | _ => (none, c.tail)

def sink (x : TokenStream) : TokenStream :=
  Stream.corec (β := TokenStream) x fun x => (none, x.tail)

def source : TokenStream :=
  Stream.corec () fun () => (some (), ())

def sourceOnes : ValueStream (BitVec 1) :=
  Stream.corec () fun () => (1#1, ())

def supp (c : ValueStream (BitVec 1)) (x : TokenStream) : TokenStream := (branch c x).snd

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

instance : TyDenote Ty where toType
  | .felt => Fin BabyBear

def BabyBear := 2^31 - 2^27 + 1

inductive Ty
  | felt

instance instDCTyDenote : TyDenote Ty where
toType := fun
| Ty.tokenstream => CIRCTStream.DCOp.TokenStream
| Ty.tokenstream2 => CIRCTStream.DCOp.TokenStream × CIRCTStream.DCOp.TokenStream
| Ty.valuestream w => CIRCTStream.DCOp.ValueStream (BitVec w)
| Ty.valuestream2 w => CIRCTStream.DCOp.ValueStream (BitVec w) × CIRCTStream.DCOp.ValueStream (BitVec w)
| Ty.valuetokenstream w => CIRCTStream.DCOp.ValueStream (BitVec w) × CIRCTStream.DCOp.TokenStream
| Ty.variadicvaluetokenstream w => CIRCTStream.DCOp.VariadicValueStream w × CIRCTStream.DCOp.TokenStream

def_denote for DC where
  | .fst => fun s => [s.fst]ₕ
  | .fstVal _ => fun s => [s.fst]ₕ
  | .fstVal' _ => fun s => [s.fst.mapOpt (·[0]?)]ₕ
  | .snd => fun s => [s.snd]ₕ
  | .pair _ => fun s₁ s₂ => [(s₁, s₂)]ₕ
  | .sndVal _ => fun s => [s.snd]ₕ
  | .sndVal' _ => fun s => [s.fst.mapOpt (·[0]?)]ₕ
  | .tokVal' _ => fun s => [s.snd]ₕ
  | .merge => fun s₁ s₂ => [CIRCTStream.DCOp.merge s₁ s₂]ₕ
  | .branch => fun s => [CIRCTStream.DCOp.branch s]ₕ
  | .fork => fun s => [CIRCTStream.DCOp.fork s]ₕ
  | .join => fun s₁ s₂ => [CIRCTStream.DCOp.join s₁ s₂]ₕ
  | .select => fun s₁ s₂ c => [CIRCTStream.DCOp.select s₁ s₂ c]ₕ
  | .sink => fun s => [CIRCTStream.DCOp.sink s]ₕ
  | .source => [CIRCTStream.DCOp.source]ₕ
  | .pack _ => fun s₁ s₂ => [CIRCTStream.DCOp.pack s₁ s₂]ₕ
  | .pack2 _ => fun s₁ => [CIRCTStream.DCOp.pack2 s₁]ₕ
  | .unpack _ => fun s => [CIRCTStream.DCOp.unpack s]ₕ
  | .unpack2 _ => fun s₁ s₂ => [CIRCTStream.DCOp.unpack2 s₁ s₂]ₕ

def branch (x : ValueStream (BitVec 1)): TokenStream × TokenStream  :=
  Stream.corec₂ (β := ValueStream (BitVec 1)) x fun x =>
    Id.run <| do
      match x 0 with
        | none => (none, none, (x.tail))
        | some x₀ =>
          if x₀.msb then
            (some (), none, (x.tail))
          else
            (none, some (), (x.tail))

inductive Term : Type
| var : Nat → Term
 
| zero : Term
 
| negOne : Term
 
| one : Term
 
| ofNat (n : Nat) : Term
 
| and : Term → Term → Term
 
| or : Term → Term → Term
 
| xor : Term → Term → Term
 
| not : Term → Term
 
| add : Term → Term → Term
 
| sub : Term → Term → Term
 
| neg : Term → Term

| shiftL : Term → Nat → Term

deriving Repr, Inhabited, Lean.ToExpr

def pack (x : ValueStream α) (y : TokenStream) : ValueStream α :=
  Stream.corec (β := ValueStream α × TokenStream) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some x₀, some _ => (x₀, (x.tail, y.tail))
    | some _, none => (none, (x, y.tail)) 
    | none, some _ => (none, (x.tail, y)) 
    | none, none => (none, (x.tail, y.tail))

def select (x y : TokenStream) (c : ValueStream (BitVec 1)): TokenStream :=
  Stream.corec (β := TokenStream × TokenStream × Stream (BitVec 1)) (x, y, c)
  fun ⟨x, y, c⟩ =>
    match (c 0) with
    | none => (none, x, y, c.tail) 
    | some 1#1 =>
      match (x 0) with
      | none => (none, x.tail, y, c) 
      | some _ => (some (), x.tail, y, c.tail) 
    | some 0#1 =>
      match (y 0) with
      | none => (none, x, y.tail, c) 
      | some _ => (some (), x, y.tail, c.tail)

def unpack (x : ValueStream (BitVec w)) : ValueStream (BitVec w) × TokenStream :=
  Stream.corec₂ (β := Stream (BitVec w)) (x)
    fun x => Id.run <| do
      match x 0 with
      | some _ => return (x 0, some (), x.tail)
      | none => return (none, none, x.tail)

abbrev DC : Dialect where
  Op := Op
  Ty := Ty

def_signature for DC where
  | .fst => (Ty.tokenstream2) → (Ty.tokenstream)
  | .fstVal t => (Ty.valuetokenstream t) → Ty.valuestream t
  | .fstVal' t => (Ty.variadicvaluetokenstream t) → Ty.valuestream t
  | .snd => (Ty.tokenstream2) → (Ty.tokenstream)
  | .pair w => (Ty.valuestream w, Ty.valuestream w) → Ty.valuestream2 w
  | .sndVal t => (Ty.valuetokenstream t) → Ty.tokenstream
  | .sndVal' t => (Ty.variadicvaluetokenstream t) → Ty.valuestream t
  | .tokVal' t => (Ty.variadicvaluetokenstream t) → Ty.tokenstream
  | .merge => (Ty.tokenstream, Ty.tokenstream) → Ty.valuestream 1
  | .branch => (Ty.valuestream 1) → Ty.tokenstream2
  | .fork => (Ty.tokenstream) → Ty.tokenstream2
  | .join => (Ty.tokenstream, Ty.tokenstream) → Ty.tokenstream
  | .select => (Ty.tokenstream, Ty.tokenstream, Ty.valuestream 1) → Ty.tokenstream
  | .sink => (Ty.tokenstream) → Ty.tokenstream
  | .source => () → Ty.tokenstream
  | .pack t => (Ty.valuestream t, Ty.tokenstream) → Ty.valuestream t
  | .pack2 t => (Ty.variadicvaluetokenstream t) → Ty.valuestream2 t
  | .unpack t => (Ty.valuestream t) → Ty.valuetokenstream t
  | .unpack2 t => (Ty.valuestream t, Ty.valuestream t) → Ty.variadicvaluetokenstream t

inductive Op
| fst
| snd
| pair (w : Nat)
| fstVal (w : Nat)
| sndVal (w : Nat)
| fstVal' (w : Nat)
| sndVal' (w : Nat)
| tokVal' (w : Nat)
| merge
| branch
| fork
| join
| select
| sink
| source
| pack (w : Nat)
| pack2 (w : Nat)
| unpack (w : Nat)
| unpack2 (w : Nat)
deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

def branch {Γ : Ctxt _} (a : Γ.Var (.valuestream 1)) : Expr (DC) Γ .pure (.tokenstream2) :=
  Expr.mk
    (op := .branch)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def fork (a : Γ.Var (.tokenstream)) : Expr (DC) Γ .pure (.tokenstream2) :=
  Expr.mk
    (op := .fork)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def unpack2 {r} {Γ : Ctxt _} (a : Γ.Var (.valuestream r)) (b : Γ.Var (.valuestream r)) : Expr (DC) Γ .pure (.variadicvaluetokenstream r) :=
  Expr.mk
    (op := .unpack2 r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def fst {Γ : Ctxt _} (a : Γ.Var (.tokenstream2)) : Expr (DC) Γ .pure (.tokenstream)  :=
  Expr.mk
    (op := .fst)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def sink {Γ : Ctxt _} (a : Γ.Var (.tokenstream)) : Expr (DC) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .sink)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def merge {Γ : Ctxt _} (a b : Γ.Var (.tokenstream)) : Expr (DC) Γ .pure (.valuestream 1) :=
  Expr.mk
    (op := .merge)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def pack {r} {Γ : Ctxt _} (a : Γ.Var (.valuestream r)) (b : Γ.Var (.tokenstream)) : Expr (DC) Γ .pure (.valuestream r) :=
  Expr.mk
    (op := .pack r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def select {Γ : Ctxt _} (a b : Γ.Var (.tokenstream)) (c : Γ.Var (.valuestream 1)) : Expr (DC) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .select)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .cons c <| .nil)
    (regArgs := .nil)

def source : Expr (DC) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .source)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

def join {Γ : Ctxt _} (a b : Γ.Var (.tokenstream)) : Expr (DC) Γ .pure (.tokenstream) :=
  Expr.mk
    (op := .join)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def pack2 {r} {Γ : Ctxt _} (a : Γ.Var (.variadicvaluetokenstream r)) : Expr (DC) Γ .pure (.valuestream2 r) :=
  Expr.mk
    (op := .pack2 r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a  <| .nil)
    (regArgs := .nil)

def fstVal {r} {Γ : Ctxt _} (a : Γ.Var (.valuetokenstream r))  : Expr (DC) Γ .pure (.valuestream r)  :=
  Expr.mk
    (op := .fstVal r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def snd {Γ : Ctxt _} (a : Γ.Var (.tokenstream2)) : Expr (DC) Γ .pure (.tokenstream)  :=
  Expr.mk
    (op := .snd)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def unpack {r} {Γ : Ctxt _} (a : Γ.Var (.valuestream r)) : Expr (DC) Γ .pure (.valuetokenstream r) :=
  Expr.mk
    (op := .unpack r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def sndVal {r} {Γ : Ctxt _} (a : Γ.Var (.valuetokenstream r))  : Expr (DC) Γ .pure (.tokenstream)  :=
  Expr.mk
    (op := .sndVal r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def pair {r} {Γ : Ctxt _} (a b: Γ.Var (.valuestream r)) : Expr (DC) Γ .pure (.valuestream2 r)  :=
  Expr.mk
    (op := .pair r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def pack2 (x : VariadicValueStream α × TokenStream) : (ValueStream (BitVec α)) × (ValueStream (BitVec α))  :=
  Stream.corec₂ (β := VariadicValueStream α × TokenStream) (x) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some x', some _ => (x'[0]?, x'[1]?, (x.tail, y.tail))
    | some _, none => (none, none, (x, y.tail))
    | none, some _ => (none, none, (x.tail, y)) 
    | none, none => (none, none, (x.tail, y.tail))

def unpack2 (x : ValueStream (BitVec w)) (y : ValueStream (BitVec w)) : VariadicValueStream w × TokenStream :=
  Stream.corec₂ (β := CIRCTStream.Stream (BitVec w) × CIRCTStream.Stream (BitVec w)) (x, y)
    fun (x, y) => Id.run <| do
      match x 0, y 0 with
      | some x', some y' => return (some [x', y'], some .unit, (x.tail, y.tail))
      | some _, none => return (none, none, (x, y.tail))
      | none, some _ => return (none, none, (x.tail, y)) 
      | none, none => return (none, none, (x.tail, y.tail))

@[reducible]
instance : TyDenote ExTy where
  toType
    | .nat => Nat
    | .bool => Bool

inductive ExTy
  | nat
  | bool
  deriving DecidableEq

instance : TyDenote Ty where toType := Ty.toType

def Ty.toType : Ty → Type
| .int => Int
| .ix => Index
| .tensor2d => Tensor2d' Int

instance : TyDenote Ty where
  toType
  | .int => Int
  | .ix => Index
  | .tensor1d => Tensor1d Int

abbrev Index := ℕ

structure Tensor1d (α : Type) [Inhabited α] where
  size : Index
  val :  Index → α
  spec : ∀ (ix: Index), ix >= size -> val ix = default

structure Tensor2d (α : Type) where
  size0 : Nat
  size1 : Nat
  val :  Fin size0 → Fin size1 → α

instance : TyDenote TestDialect.Ty where toType
  | .int => Int
  | .int2 => Int × Int

def_signature for TestDialect
  | .noop => () -> []
  | .mkPair => (.int, .int) -> .int2
  | .unPair => (.int2) -> [.int, .int]

def_denote for TestDialect
  | .noop => []ₕ
  | .unPair => fun (x, y) => [x, y]ₕ
  | .mkPair => fun x y => [(x, y)]ₕ

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

def TestDialect : Dialect where
  Ty := Ty
  Op := Op

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

instance : TyDenote Ty where
  toType
  | Ty.bv => BitVec 64

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

instance : TyDenote LLVM.Ty where
  toType := fun
    | bitvec w => LLVM.IntW w

def IntW w := PoisonOr <| BitVec w

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

@[reducible]
instance : TyDenote Ty where
  toType
    | .int => BitVec 32

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

instance : TyDenote Ty where toType
  | .int => BitVec 32

inductive ExTy
  | nat
  | bool
  deriving DecidableEq, Repr

instance : TyDenote Ty2 where
toType := fun
|  Ty2.bitvec w => BitVec w

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

instance instHandshakeTyDenote : TyDenote Ty where
toType := fun
| Ty.stream ty2 => CIRCTStream.Stream (TyDenote.toType ty2)
| Ty.stream2 ty2 => CIRCTStream.Stream (TyDenote.toType ty2) × CIRCTStream.Stream (TyDenote.toType ty2)
| Ty.stream2token ty2 => CIRCTStream.Stream (TyDenote.toType ty2) × CIRCTStream.Stream (TyDenote.toType (Ty2.bitvec 1))

def_denote for Handshake where
| .fst _ => fun s => [s.fst]ₕ
| .snd _ => fun s => [s.snd]ₕ
| .branch _ => fun s c => [HandshakeOp.branch s c]ₕ
| .merge _ => fun s₁ s₂ => [HandshakeOp.merge s₁ s₂]ₕ
| .altMerge _ => fun s₁ s₂ => [HandshakeOp.altMerge s₁ s₂]ₕ
| .fork _ => fun s => [HandshakeOp.fork s]ₕ
| .controlMerge _ => fun s₁ s₂ => [HandshakeOp.controlMerge s₁ s₂]ₕ
| .join _ => fun s₁ s₂ => [HandshakeOp.join s₁ s₂]ₕ
| .mux _ => fun s₁ s₂ c => [HandshakeOp.mux s₁ s₂ c]ₕ
| .sink _ => fun s => [HandshakeOp.sink s]ₕ
| .sync _ => fun s₁ s₂ => [HandshakeOp.sync s₁ s₂]ₕ
| .supp _ => fun s₁ s₂ => [HandshakeOp.supp s₁ s₂]ₕ
| .not => fun s₁ => [HandshakeOp.not s₁]ₕ

abbrev Handshake : Dialect where
  Op := Op
  Ty := Ty

def_signature for Handshake where
| .fst t => (Ty.stream2 t) → Ty.stream t
| .snd t => (Ty.stream2 t) → Ty.stream t
| .branch t => (Ty.stream t, Ty.stream (Ty2.bitvec 1)) → Ty.stream2 t
| .merge t => (Ty.stream t, Ty.stream t) → Ty.stream t
| .altMerge t => (Ty.stream t, Ty.stream t) → Ty.stream t
| .fork t => (Ty.stream t) → Ty.stream2 t 
| .controlMerge t => (Ty.stream t, Ty.stream t) → (Ty.stream2token t)
| .join t => (Ty.stream t, Ty.stream t) → (Ty.stream (Ty2.bitvec 1))
| .mux t => (Ty.stream t, Ty.stream t, Ty.stream (Ty2.bitvec 1)) → Ty.stream t
| .sink t => (Ty.stream t) → (Ty.stream (Ty2.bitvec 1))
| .sync t => (Ty.stream t, Ty.stream t) → Ty.stream2 t
| .supp t => (Ty.stream t, Ty.stream (Ty2.bitvec 1)) → Ty.stream t
| .not => (Ty.stream (Ty2.bitvec 1)) → Ty.stream (Ty2.bitvec 1)

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

def sync {Γ : Ctxt _} (a b : Var Γ (.stream r)) : Expr (Handshake) Γ .pure (.stream2 r)  :=
  Expr.mk
    (op := .sync r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def snd {Γ : Ctxt _} (a : Var Γ (.stream2 r)) : Expr (Handshake) Γ .pure (.stream r)  :=
  Expr.mk
    (op := .snd r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

inductive Op
| fst (t : Ty2)
| snd (t : Ty2)
| branch (t : Ty2)
| merge (t : Ty2)
| altMerge (t : Ty2)
| fork (t : Ty2)
| controlMerge (t : Ty2)
| join (t : Ty2)
| mux (t : Ty2)
| sink (t : Ty2)
| sync (t : Ty2)
| supp (t : Ty2)
| not
deriving Inhabited, DecidableEq, Repr, Lean.ToExpr

def merge {Γ : Ctxt _} (a b : Var Γ (.stream r)) : Expr (Handshake) Γ .pure (.stream r)  :=
  Expr.mk
    (op := .merge r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def supp {Γ : Ctxt _} (a : Var Γ (.stream r)) (b : Var Γ (.stream (.bitvec 1))) : Expr (Handshake) Γ .pure (.stream r)  :=
  Expr.mk
    (op := .supp r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons b <| .nil)
    (regArgs := .nil)

def branch {r} {Γ : Ctxt _} (a : Var Γ (.stream r)) (c : Var Γ (.stream (.bitvec 1))) : Expr (Handshake) Γ .pure (.stream2 r) :=
  Expr.mk
    (op := .branch r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons c <| .nil)
    (regArgs := .nil)

def fst {Γ : Ctxt _} (a : Var Γ (.stream2 r)) : Expr (Handshake) Γ .pure (.stream r)  :=
  Expr.mk
    (op := .fst r)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

def altMerge (x y : Stream α) : Stream α :=
  Stream.corec (β := Stream α × Stream α × ConsumeFrom) (x, y, .left) fun ⟨x, y, consume⟩ =>
    match consume with
      | .left  =>
        let x0 := x.head
        let x := x.tail
        let nextConsume := match x0 with
          | some _ => .right
          | none   => .left
        (x0, x, y, nextConsume)
      | .right =>
        let y0 := y.head
        let y := y.tail
        let nextConsume := match y0 with
          | some _ => .left
          | none   => .right
        (y0, x, y, nextConsume)

def branch (x : Stream α) (c : Stream (BitVec 1)) : Stream α × Stream α :=
  Stream.corec₂ (β := Stream α × Stream (BitVec 1)) (x, c)
    fun ⟨x, c⟩ => Id.run <| do
      let c₀ := c 0
      let c' := c.tail
      let x₀ := x 0
      let x' := x.tail
      match c₀, x₀ with
        | none, _ => (none, none, (x, c'))
        | _, none => (none, none, (x', c))
        | some c₀, some x₀ =>
          if c₀ = 1 then
            (some x₀, none, (x', c'))
          else
            (none, some x₀, (x', c'))

def controlMerge (x y : Stream α) : Stream α × Stream (BitVec 1) :=
  Stream.corec₂ (β := Stream α × Stream α) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some x', some _ => (some x', some 1, (x.tail, y))
    | some x', none => (some x', some 1, (x.tail, y.tail))
    | none, some y' => (some y', some 0, (x.tail, y.tail))
    | none, none => (none, none, (x.tail, y.tail))

def fork (x : Stream α) : Stream α × Stream α :=
  Stream.corec₂ (β := Stream α) x
    fun x => Id.run <| do
      let x0 := x 0
      let x' := x.tail
      (x0, x0, x')

def join (x y : Stream α) : Stream (BitVec 1) :=
  Stream.corec (β := Stream α × Stream α) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some _, some _ => (some 1, (x.tail, y.tail))
    | some _, none   => (none, (x, y.tail))
    | none, some _   => (none, (x.tail, y))
    | none, none     => (none, (x.tail, y.tail))

def merge (x y : Stream α) : Stream α :=
  Stream.corec (β := Stream α × Stream α) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some x', some _ => (some x', (x.tail, y))
    | some x', none => (some x', (x.tail, y.tail))
    | none, some y' => (some y', (x.tail, y.tail))
    | none, none => (none, (x.tail, y.tail))

def mux (x y : Stream α) (c : Stream (BitVec 1)) : Stream α :=
  Stream.corec (β := Stream α × Stream α × Stream (BitVec 1)) (x, y, c) fun ⟨x, y, c⟩ => Id.run <| do
    match x 0, y 0, c 0 with
      | none, _, some 1 => (none, (x.tail, y, c)) 
      | some _, _, some 1 => (x 0, (x.tail, y, c.tail)) 
      | _, none, some 0 => (none, (x, y.tail, c)) 
      | _, some _, some 0 => (y 0, (x, y.tail, c.tail)) 
      | _, _, none => (none, (x, y, c.tail))

def not (x : Stream (BitVec 1)) : Stream (BitVec 1) :=
  Stream.corec (β := Stream (BitVec 1)) x fun x =>
    match x 0 with
    | some 1 => (some 0, (x.tail))
    | some 0 => (some 0,  (x.tail))
    | none => (none, (x.tail))

def sink (x : Stream α) : Stream (BitVec 1) :=
  Stream.corec (β := Stream α) (x) fun (x) => (none, x.tail)

def supp (x : Stream α) (c : Stream (BitVec 1)) : Stream α := (branch x c).snd

def sync (x y : Stream α) : Stream α × Stream α :=
  Stream.corec₂ (β := Stream α × Stream α) (x, y) fun ⟨x, y⟩ =>
    match x 0, y 0 with
    | some x', some y' => (some x', some y', (x.tail, y.tail))
    | some _, none => (none, none, (x, y.tail))
    | none, some _ => (none, none, (x.tail, y))
    | none, none => (none, none, (x.tail, y.tail))

@[simp]
instance : TyDenote LLVMPlusRiscV.Ty where
  toType := fun
    | .llvm llvmTy => TyDenote.toType llvmTy
    | .riscv riscvTy => TyDenote.toType riscvTy

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

inductive ExTy
  | nat
  deriving DecidableEq, Repr

@[reducible]
instance : TyDenote ExTy where
  toType
    | .nat => Nat

instance : TyDenote (Dialect.Ty VariadicDialect) where
  toType := fun | .int => BitVec 32

def_denote for VariadicDialect where
  | .const z => BitVec.ofInt _ z
  | .add _   => fun xs => xs.foldl (fun ⟨⟩ => (· + ·)) 0#32

inductive Op : Type
| add (n : Nat) : Op
| const : (val : ℤ) → Op
deriving DecidableEq, Repr

def VariadicDialect : Dialect where
  Op := Op
  Ty := Ty

def_signature for VariadicDialect where
  | .const _  => () -> .int
  | .add n    => ${List.replicate n .int} → .int

inductive Ty
| int
deriving DecidableEq, Repr

instance : TyDenote (Dialect.Ty Comb) where
  toType := fun
  | .bitvec w => BitVec w

instance : TyDenote (Scf d).Ty := inferInstanceAs (TyDenote d.Ty)

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

instance : TyDenote Ty where
  toType
    | .int => ℤ
    | .bool => Bool
    | .nat => Nat

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

instance : TyDenote SLLVM.Ty where
  toType
  | .arith t => ⟦t⟧
  | .ptr => SLLVM.Ptr
  | .mem => MemorySSAState

inductive SLLVMOp where
  | arith (o : LLVM.Op)
  | ptradd
  | load (w : Nat)
  | store (w : Nat)
  | alloca (w : Nat)
  | loadPure (w : Nat)
  | storePure (w : Nat)
  deriving DecidableEq, Lean.ToExpr

@[match_pattern] abbrev Ty.mem : SLLVM.Ty := .mem

@[match_pattern] nonrec abbrev Op.arith : LLVM.Op → SLLVM.Op := .arith

@[match_pattern] nonrec abbrev Op.neg (w : Nat) : SLLVM.Op := arith <| Op.neg w
@[match_pattern] nonrec abbrev Op.not (w : Nat) : SLLVM.Op := arith <| Op.not w
@[match_pattern] nonrec abbrev Op.copy (w : Nat) : SLLVM.Op := arith <| Op.copy w
@[match_pattern] nonrec abbrev Op.freeze (w : Nat) : SLLVM.Op := arith <| Op.freeze w
@[match_pattern] nonrec abbrev Op.sext (w w' : Nat) : SLLVM.Op := arith <| Op.sext w w'
@[match_pattern] nonrec abbrev Op.zext (w w' : Nat) (flag : LLVM.NonNegFlag := { }) : SLLVM.Op := arith <| Op.zext w w' flag
@[match_pattern] nonrec abbrev Op.trunc (w w' : Nat) (flags : LLVM.NoWrapFlags := { }) : SLLVM.Op := arith <| Op.trunc w w' flags

@[match_pattern] nonrec abbrev Op.and (w : Nat) : SLLVM.Op := arith <| Op.and w
@[match_pattern] nonrec abbrev Op.or (w : Nat) (flag : LLVM.DisjointFlag := { }) : SLLVM.Op := arith <| Op.or w flag
@[match_pattern] nonrec abbrev Op.xor (w : Nat) : SLLVM.Op := arith <| Op.xor w
@[match_pattern] nonrec abbrev Op.shl (w : Nat) (flags : LLVM.NoWrapFlags := { }) : SLLVM.Op := arith <| Op.shl w flags
@[match_pattern] nonrec abbrev Op.lshr (w : Nat) (flag : LLVM.ExactFlag := { }) : SLLVM.Op := arith <| Op.lshr w flag
@[match_pattern] nonrec abbrev Op.ashr (w : Nat) (flag : LLVM.ExactFlag := { }) : SLLVM.Op := arith <| Op.ashr w flag
@[match_pattern] nonrec abbrev Op.add (w : Nat) (flags : LLVM.NoWrapFlags := { }) : SLLVM.Op := arith <| Op.add w flags
@[match_pattern] nonrec abbrev Op.mul (w : Nat) (flags : LLVM.NoWrapFlags := { }) : SLLVM.Op := arith <| Op.mul w flags
@[match_pattern] nonrec abbrev Op.sub (w : Nat) (flags : LLVM.NoWrapFlags := { }) : SLLVM.Op := arith <| Op.sub w flags

@[match_pattern] nonrec abbrev Op.icmp (c : LLVM.IntPred) (w : Nat) : SLLVM.Op := arith <| Op.icmp c w
@[match_pattern] nonrec abbrev Op.const (w : Nat) (val : Int) : SLLVM.Op := arith <| Op.const w val
@[match_pattern] nonrec abbrev Op.select (w : Nat) : SLLVM.Op := arith <| Op.select w

@[match_pattern] nonrec abbrev Op.udiv (w : Nat) (flag : LLVM.ExactFlag := { }) : SLLVM.Op := arith <| Op.udiv w flag
@[match_pattern] nonrec abbrev Op.sdiv (w : Nat) (flag : LLVM.ExactFlag := { }) : SLLVM.Op := arith <| Op.sdiv w flag
@[match_pattern] nonrec abbrev Op.urem : Nat → SLLVM.Op := arith ∘ Op.urem
@[match_pattern] nonrec abbrev Op.srem : Nat → SLLVM.Op := arith ∘ Op.srem

@[match_pattern] abbrev Ty.ptr : SLLVM.Ty := .ptr

def MemorySSAState := PoisonOr MemoryState

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

structure MemoryState where
  mem : Std.HashMap BlockId Block
  deriving Inhabited

structure GlobalState where
  alloc : AllocState
  mem : MemoryState

structure BlockId where
  id : Nat
  deriving DecidableEq, Hashable, Inhabited

inductive Block where
   
  | dead
   
  | live (b : LiveBlock)

structure LiveBlock where
  (length : Nat)
  (bytes : BitVec (8 * length))

def Ptr : Type := PoisonOr Pointer

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

instance instHandshakeTyDenote : TyDenote Ty where
toType := fun
| Ty.stream ty2 => CIRCTStream.Stream (TyDenote.toType ty2)
| Ty.stream2 ty2 => CIRCTStream.Stream (TyDenote.toType ty2) × CIRCTStream.Stream (TyDenote.toType ty2)
| Ty.stream2token ty2 => CIRCTStream.Stream (TyDenote.toType ty2) × CIRCTStream.Stream (TyDenote.toType (Ty2.bitvec 1))

def_denote for Handshake where
| .fst _ => fun s => [s.fst]ₕ
| .snd _ => fun s => [s.snd]ₕ
| .branch _ => fun s c => [HandshakeOp.branch s c]ₕ
| .merge _ => fun s₁ s₂ => [HandshakeOp.merge s₁ s₂]ₕ
| .altMerge _ => fun s₁ s₂ => [HandshakeOp.altMerge s₁ s₂]ₕ
| .fork _ => fun s => [HandshakeOp.fork s]ₕ
| .controlMerge _ => fun s₁ s₂ => [HandshakeOp.controlMerge s₁ s₂]ₕ
| .join _ => fun s₁ s₂ => [HandshakeOp.join s₁ s₂]ₕ
| .mux _ => fun s₁ s₂ c => [HandshakeOp.mux s₁ s₂ c]ₕ
| .sink _ => fun s => [HandshakeOp.sink s]ₕ
| .sync _ => fun s₁ s₂ => [HandshakeOp.sync s₁ s₂]ₕ
| .supp _ => fun s₁ s₂ => [HandshakeOp.supp s₁ s₂]ₕ
| .not => fun s₁ => [HandshakeOp.not s₁]ₕ

def test_add__nsw_is_safe_after := [llvm|
{
^0(%arg36 : i32):
  %0 = llvm.mlir.constant(-2147483648 : i32) : i32
  %1 = llvm.mlir.constant(1 : i32) : i32
  %2 = llvm.or %arg36, %0 : i32
  %3 = llvm.add %2, %1 overflow<nsw> : i32
  "llvm.return"(%3) : (i32) -> ()
}
]

def zext_ashr_before := [llvm|
{
^0(%arg11 : i8):
  %0 = llvm.mlir.constant(2 : i16) : i16
  %1 = llvm.zext %arg11 : i8 to i16
  %2 = llvm.ashr %1, %0 : i16
  %3 = llvm.and %2, %1 : i16
  "llvm.return"(%3) : (i16) -> ()
}
]
</all_available_defs>

<all_available_lemmas>
List.get_eq_getElem in Init.Data.List.Lemmas
List.length_cons in Init.Data.List.Basic

@[simp] theorem map_cons {A B : α → Type u} {as : List α} {f : (a : α) → A a → B a}
    {x : A a} {xs : HVector A as} :
    map f (cons x xs) = cons (f _ x) (map f xs)

theorem icmp_ne_and_pow2_shl_pow2_proof : icmp_ne_and_pow2_shl_pow2_before ⊑ icmp_ne_and_pow2_shl_pow2_after

theorem basic_preserve_nsw_thm (e e_1 : IntW 8) :
  LLVM.xor (add (LLVM.xor e_1 (const? 8 (-1))) e { «nsw» := true, «nuw» := false }) (const? 8 (-1)) ⊑
    sub e_1 e { «nsw» := true, «nuw» := false }

theorem xor_common_op_commute0_thm.extracted_1._1 : ∀ (x x_1 : BitVec 8), x_1 ^^^ x ||

theorem fold_mask_cmps_to_true_logical_proof : fold_mask_cmps_to_true_logical_before ⊑ fold_mask_cmps_to_true_logical_after

theorem bools_multi_uses2_logical_thm.extracted_1._30 : ∀ (x x_1 x_2 : BitVec 1),
  ¬x_2 ^^^ 1#1 = 1#1 → x_2 = 1#1 → 0#1 = 1#1 → ¬x = 1#1 → True → 0#1 + x = 0#1

theorem e_1497 :
     - 1 * t - 3 * (z ^^^ t) + 11 * z - 2 * (z |||  ~~~t) - 11 *  ~~~z - 7 * (z ||| t) + 1 *  ~~~(z ^^^ t) + 13 *  ~~~(z ||| t) + 23 *  ~~~(z |||  ~~~t) + 2 * (z &&&  ~~~t) - 1 * (z &&& t) = 1 *  ~~~(z &&&  ~~~z)

theorem ashrsgt_03_14_exact_proof : ashrsgt_03_14_exact_before ⊑ ashrsgt_03_14_exact_after
</all_available_lemmas>

<local_ctx>
import Mathlib.Data.Fintype.Basic

import LeanMLIR.HVector

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

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

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

variable {Ty : Type} {Γ Δ : Ctxt Ty}

namespace Ctxt

section Instances

open Lean in

end Instances

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

section GetElemLemmas

end GetElemLemmas

section Lemmas

variable (Γ : Ctxt Ty) (ts us : List Ty)

variable {m} [Monad m] [LawfulMonad m] (t u : m _) in

section Lemmas

variable {Γ Δ : Ctxt Ty} {tys : List Ty}

end Lemmas

end Lemmas

section Rec

end Rec

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

namespace Var

section Lemmas

variable {Γ : Ctxt Ty} {t : Ty}

end Lemmas

section Lemmas

variable {t} (v : Var Γ t)

end Lemmas

section Lemmas

variable {t : Ty} {v : Γ.Var t}

end Lemmas

section Lemmas

end Lemmas

end Var

end Ctxt

open Ctxt

namespace HVector

variable {A : α → _} {as : List α} (xs : HVector A as) {Γ : Ctxt α}

end HVector

namespace Ctxt

section Comp

variable {Γ Δ Ξ : Ctxt Ty} (f : Hom Γ Δ) (g : Hom Δ Ξ)

end Comp

section Lemmas

end Lemmas

variable {Γ Δ Δ' : Ctxt Ty} in

section Valuation

variable [TyDenote Ty]

def Valuation (Γ : Ctxt Ty) : Type :=
  ⦃t : Ty⦄ → Γ.Var t → (toType t)

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

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

variable {m} [Monad m] [LawfulMonad m] in

variable {V : Γ.Valuation} {W : Δ.Valuation}

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

section Lemmas

variable {ts : List Ty}

end Lemmas

def Valuation.reassignVars [DecidableEq Ty] {ts : List Ty} {Γ : Ctxt Ty}
    (V : Γ.Valuation) (var : HVector Γ.Var ts) (val : HVector toType ts) : Γ.Valuation :=
  fun _ vneedle =>
    match var.idxOf? vneedle with
    | none => V vneedle
    | some ⟨i, h⟩ => h ▸ val.get i
</local_ctx>

<target_theorem>
@[simp] theorem Valuation.reassignVars_eq [DecidableEq Ty] (V : Γ.Valuation) :
    V.reassignVars vs (vs.map V) = V :=
</target_theorem>

