-- LeanMLIR/LeanMLIR/ErasedContext.lean

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

-- Repository definitions used:
content: protected instance HVector.decidableEqReg [DecidableEq d.Op] [DecidableEq d.Ty] :
    ∀ {l : RegionSignature d.Ty}, DecidableEq (HVector (fun t => Com d t.1 .impure t.2) l)
  | _, .nil, .nil => isTrue rfl
  | _, .cons x₁ v₁, .cons x₂ v₂ =>
    letI := HVector.decidableEqReg v₁ v₂
    letI := Com.decidableEq x₁ x₂
    decidable_of_iff (x₁ = x₂ ∧ v₁ = v₂) (by admit /- proof elided -/
    )

content: def 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 -/
          ⟩

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

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

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

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

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

content: 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]ₕ

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

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

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

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

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

content: def RegIndex := Fin RegFile.numRegisters

content: def RegFile.numRegisters := 32

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

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

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

content: def RegFile.registerWidth := 64

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

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

content: 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₁]ₕ

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

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

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

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

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

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

content: def VarName : Type := String

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

content: protected instance Expr.decidableEq [DecidableEq d.Op] [DecidableEq d.Ty] :
    {Γ : Ctxt d.Ty} → {ty : List d.Ty} → DecidableEq (Expr d Γ eff ty)
  | Γ, _, .mk op₁ rfl eff_le₁ arg₁ regArgs₁, .mk op₂ eq eff_le₂ arg₂ regArgs₂ =>
    if ho : op₁ = op₂ then by
      subst ho
      letI := HVector.decidableEq arg₁ arg₂
      letI := HVector.decidableEqReg regArgs₁ regArgs₂
      exact decidable_of_iff (arg₁ = arg₂ ∧ regArgs₁ = regArgs₂) (by admit /- proof elided -/
      )
    else isFalse (by admit /- proof elided -/
    )

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: def TokenStream := Stream Unit

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

content: def ValueStream := Stream

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

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

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

content: 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,
  )

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: inductive Ty
  | felt

content: 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₂]ₕ

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: inductive ExTy
  | nat
  | bool
  deriving DecidableEq

content: instance : TyDenote Ty where toType := Ty.toType

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

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

content: abbrev Index := ℕ

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: 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₁]ₕ

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: inductive ExTy
  | nat
  deriving DecidableEq, Repr

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

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

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

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

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

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

content: inductive Ty
| int
deriving DecidableEq, Repr

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

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

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

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

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

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

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

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

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

content: def MemorySSAState := PoisonOr MemoryState

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

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

content: structure GlobalState where
  alloc : AllocState
  mem : MemoryState

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

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

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

content: def Ptr : Type := PoisonOr Pointer

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

-- Repository lemmas used:
content: @[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)

-- Local definitions used:
content: class TyDenote (β : Type) : Type 1 where
  toType : β → Type

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

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

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

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

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

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

content: def 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 lemmas used:

-- Target theorem and proof:
@[simp] theorem Valuation.reassignVars_eq [DecidableEq Ty] (V : Γ.Valuation) :
    V.reassignVars vs (vs.map V) = V
:= by
  funext t w
  unfold reassignVars
  induction vs
  case nil => rfl
  case cons v vs ih =>
    by_cases h_eq : w.eq v
    · have := h_eq.ty_eq
      subst this
      have := h_eq.to_eq
      subst this
      simp
    · unfold Var.eq at h_eq
      simp only [HVector.idxOf?, h_eq, ↓reduceDIte, List.get_eq_getElem, List.length_cons,
        HVector.map_cons]
      split at ih <;> simp_all