import Mathlib.Logic.Function.Iterate
import Mathlib.Tactic.NormNum

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

BitStream

def BitStream : Type := Nat → Bool

BitStream.zero

private theorem two_le_add_iff_odd_and_odd (n m : Nat) :
    2 ≤ n % 2 + m % 2 ↔ n % 2 = 1 ∧ m % 2 = 1

private theorem add_odd_iff_neq (n m : Nat) :
    (n + m) % 2 = 1 ↔ (n % 2 = 1) ≠ (m % 2 = 1)

import Mathlib.Tactic.NormNum

import Mathlib.Logic.Function.Iterate

section UpStream

namespace Int

end Int

end UpStream

def BitStream : Type := Nat → Bool

namespace BitStream

section Basic

section Lemmas

end Lemmas

end Basic

section OfNat

end OfNat

section ToBitVec

abbrev ofBitVecSext {w} (x : BitVec w) : BitStream :=
  fun i => if i < w then x.getLsbD i else x.msb

section Lemmas

end Lemmas

end ToBitVec

section BitwiseOps

section Lemmas

variable {w : Nat}

variable (x y : BitStream) (i : Nat)

variable (x y : BitVec (w+1))

end Lemmas

end BitwiseOps

section Scan

end Scan

section FindIndex

section Arith

def addAux (x y : BitStream) (i : Nat) :  Bool × Bool :=
  let carryIn : Bool := match i with
    | 0 => false
    | i + 1 => (addAux x y i).2
  Prod.swap (BitVec.adcb (x i) (y i) carryIn)

def add (x y : BitStream) : BitStream :=
  fun n => (addAux x y n).1

abbrev zero   : BitStream := fun _ => false

section Lemmas

variable {w : Nat} {x y : BitVec w} {a b a' b' : BitStream}

local infix:20 " ≈ʷ " => EqualUpTo w

theorem ofBitVec_add : ofBitVecSext (x + y) ≈ʷ (ofBitVecSext x) + (ofBitVecSext y) :=

:= by
  intros n a
  have add_lemma : ⟨(x + y).getLsbD n, BitVec.carry (n + 1) x y false ⟩ = (ofBitVecSext x).addAux (ofBitVecSext y) n := by
    induction n
    case zero =>
      simp [addAux, BitVec.adcb, BitVec.carry, BitVec.getLsbD, a,
        two_le_add_iff_odd_and_odd, add_odd_iff_neq]
      bv_decide
    case succ i ih =>
      simp [addAux, ← ih (by omega), BitVec.adcb, a, BitVec.carry_succ, BitVec.getElem_add];
  simp [HAdd.hAdd, Add.add, BitStream.add, ← add_lemma, a, -BitVec.add_eq, -Nat.add_eq]


import Blase.SingleWidth.Defs
import Mathlib.Data.Fin.Basic
import Mathlib.Data.Bool.Basic
import Blase.Fast.BitStream
import Blase.Blase.Fast.BitStream

syntax "slt" : MLIR.Pretty.uniform_op

syntax "slt" : MLIR.Pretty.uniform_op

macro_rules
  | `(mlir_op| $res:mlir_op_operand = const ($x)
      $[: $outer_type]? ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _)
      `(mlir_op| $res:mlir_op_operand = "const"()
          {val = $x:num : $outer_type} : ($outer_type) -> ($outer_type) )

macro_rules
  | `(mlir_op| $res:mlir_op_operand = li ($x)
     $[: $outer_type]?  ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _ )
      `(mlir_op| $res:mlir_op_operand = "li"()
          {imm = $x:num : $outer_type } : ($outer_type) -> ($outer_type))

macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithImmediate $reg1 , $x : $t)  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {imm = $x:num : $t} : ($t) -> ($t) )

macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithShamt $reg1 , $x  : $t )  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {shamt = $x:num : $t}  : ($t) -> ($t) )

BitStream

def BitStream : Type := Nat → Bool

borrow

def borrow (x y : BitStream) : BitStream :=
  fun n => (subAux x y n).2

subAux

def subAux (x y : BitStream) : Nat → Bool × Bool
  | 0 => (xor (x 0) (y 0), !(x 0) && y 0)
  | n+1 =>
    let borrow := (subAux x y n).2
    let a := x (n + 1)
    let b := y (n + 1)
    (xor a (xor b borrow), !a && b || ((!(xor a b)) && borrow))

nxor

def nxor (a b : BitStream) : BitStream := fun i => a i == b i

scanAnd

def scanAnd (s : BitStream) : BitStream := scanl true Bool.and s

scanl

abbrev scanl (init : Bool) (f : Bool → Bool → Bool) (s : BitStream) : BitStream :=
  fun n => match n with
    | 0 => f init (s 0)
    | n+1 => f (scanl init f s n) (s (n + 1))

scanOr

def scanOr (s : BitStream) : BitStream := scanl false Bool.or s

falseIffNeq

abbrev falseIffNeq (n : Nat) : BitStream := fun i => decide (i == n)

falseIffLt

abbrev falseIffLt (n : Nat) : BitStream := fun i => decide (i ≥ n)

falseIffGe

abbrev falseIffGe (n : Nat) : BitStream := fun i => decide (i < n)

falseIffEq

abbrev falseIffEq (n : Nat) : BitStream := fun i => decide (i != n)

falseIffGt

abbrev falseIffGt (n : Nat) : BitStream := fun i => decide (i ≤ n)

Term.width

def Term.width (t : Term) : WidthExpr :=
  match t with

  | .ofNat w _n => w
  | .var _v w => w
  | .add w _a _b => w
  | .zext _a wnew => wnew
  | .setWidth _a wnew => wnew
  | .sext _a wnew => wnew
  | .bor w _a _b => w
  | .band w _a _b => w
  | .bxor w _a _b => w
  | .bnot w _a => w
  | .boolVar _v => WidthExpr.const 1
  | .boolConst _b => WidthExpr.const 1
  | .shiftl w _a _k => w
  | .bvOfBool _b => WidthExpr.const 1
  | binWidthRel _k wa wb => WidthExpr.const 0
  | binRel _k w _a _b => w
  | or _p1 _p2 => WidthExpr.const 0
  | and _p1 _p2 => WidthExpr.const 0
  | pvar _v => WidthExpr.const 0
  | boolBinRel _k _a _b => WidthExpr.const 0

falseIffLe

abbrev falseIffLe (n : Nat) : BitStream := fun i => decide (i > n)

negOne

abbrev negOne : BitStream := fun _ => true

shiftLeft

def shiftLeft (x : BitStream) (k : Nat) : BitStream :=
  fun i => if i < k then false else x (i - k)

ofNat

def ofNat (x : Nat) : BitStream :=
  Nat.testBit x

one

abbrev one    : BitStream := (· == 0)

zero

abbrev zero   : BitStream := fun _ => false

Term.denote

def Term.denote (w : Nat) (t : Term) (vars : List (BitVec w)) : BitVec w :=
  match t with
  | ofNat n => BitVec.ofNat w n
  | var n => vars.getD n default
  | zero => 0#w
  | negOne => -1#w
  | one  => 1#w
  | and a b => (a.denote w vars) &&& (b.denote w vars)
  | or a b => (a.denote w vars) ||| (b.denote w vars)
  | xor a b => (a.denote w vars) ^^^ (b.denote w vars)
  | not a => ~~~ (a.denote w vars)
  | add a b => (a.denote w vars) + (b.denote w vars)
  | sub a b => (a.denote w vars) - (b.denote w vars)
  | neg a => - (a.denote w vars)

  | shiftL a n => (a.denote w vars) <<< n

Predicate.denote

def Predicate.denote (p : Predicate) (w : Nat) (vars : List (BitVec w)) : Prop :=
  match p with
  | .width .ge k => k ≤ w
  | .width .gt k => k < w
  | .width .le k => w ≤ k
  | .width .lt k => w < k
  | .width .neq k => w ≠ k
  | .width .eq k => w = k
  | .binary .eq t₁ t₂ => t₁.denote w vars = t₂.denote w vars
  | .binary .neq t₁ t₂ => t₁.denote w vars ≠ t₂.denote w vars
  | .binary .sle  t₁ t₂ => ((t₁.denote w vars).sle (t₂.denote w vars)) = true
  | .binary .slt  t₁ t₂ => ((t₁.denote w vars).slt (t₂.denote w vars)) = true
  | .binary .ule  t₁ t₂ => ((t₁.denote w vars).ule (t₂.denote w vars)) = true
  | .binary .ult  t₁ t₂ => (t₁.denote w vars).ult (t₂.denote w vars) = true
  | .land  p q => p.denote w vars ∧ q.denote w vars
  | .lor  p q => p.denote w vars ∨ q.denote w vars

Predicate

inductive Predicate : Type where

| width (wp : WidthPredicate) (n : Nat) : Predicate
| binary (p : BinaryPredicate) (t₁ t₂ : Term)
| land  (p q : Predicate) : Predicate
| lor (p q : Predicate) : Predicate
deriving Repr, Inhabited, Lean.ToExpr

WidthPredicate

inductive WidthPredicate
| eq
| neq
| lt
| le
| gt
| ge
deriving Repr, Inhabited, Lean.ToExpr

ofBitVecSext

abbrev ofBitVecSext {w} (x : BitVec w) : BitStream :=
  fun i => if i < w then x.getLsbD i else x.msb

toBitVec

def toBitVec (w : Nat) (x : BitStream) : BitVec w :=
  match w with
  | 0   => 0#0
  | w+1 => (x.toBitVec w).cons (x w)

@[simp] theorem subAux_eq_BitVec_carry (a b : BitStream) (w i : Nat) (hi : i < w) :
    (a.subAux b i).2 = !(BitVec.carry (i + 1) (a.toBitVec w) ((~~~b).toBitVec w) true)

Predicate.eval

private theorem BitVec.lt_eq_decide_ult {x y : BitVec w} : (x < y) = decide (x.ult y)

theorem Predicate.evalUlt_denote_false_iff {w : Nat} (a b : Term) (vars : List (BitVec w)) :
    evalUlt (a.eval (List.map .ofBitVecSext vars)) (b.eval (List.map .ofBitVecSext vars)) w = false ↔
    (Term.denote w a vars < Term.denote w b vars)

import Mathlib.Data.Bool.Basic

import Mathlib.Data.Fin.Basic

import Blase.Fast.BitStream

import Blase.SingleWidth.Defs

open Term

open BitStream in

def Term.eval (t : Term) (vars : List BitStream) : BitStream :=
  match t with
  | var n       => vars.getD n default
  | zero        => BitStream.zero
  | one         => BitStream.one
  | negOne      => BitStream.negOne
  | ofNat n     => BitStream.ofNat n
  | and t₁ t₂   => (t₁.eval vars) &&& (t₂.eval vars)
  | or t₁ t₂    => (t₁.eval vars) ||| (t₂.eval vars)
  | xor t₁ t₂   => (t₁.eval vars) ^^^ (t₂.eval vars)
  | not t       => ~~~(t.eval vars)
  | add t₁ t₂   => (Term.eval t₁ vars) + (Term.eval t₂ vars)
  | sub t₁ t₂   => (Term.eval t₁ vars) - (Term.eval t₂ vars)
  | neg t       => -(Term.eval t vars)

  | shiftL t n  => BitStream.shiftLeft (Term.eval t vars) n

def Predicate.evalEq (t₁ t₂ : BitStream) : BitStream := (t₁ ^^^ t₂).concat false |>.scanOr

def Predicate.evalNeq (t₁ t₂ : BitStream) : BitStream := (t₁.nxor t₂).concat true |>.scanAnd

def Predicate.evalLor (t₁ t₂ : BitStream) : BitStream := (t₁ &&& t₂)

def Predicate.evalLand (t₁ t₂ : BitStream) : BitStream := (t₁ ||| t₂)

def Predicate.evalUlt (t₁ t₂ : BitStream) : BitStream := (~~~ (t₁.borrow t₂)).concat true

def Predicate.evalMsbEq (t₁ t₂ : BitStream) : BitStream :=
  (t₁ ^^^ t₂).concat false

def Predicate.evalSlt (t₁ t₂ : BitStream) : BitStream :=
    (((Predicate.evalUlt t₁ t₂)) ^^^ (Predicate.evalMsbEq t₁ t₂))

open BitStream in

def Predicate.eval (p : Predicate) (vars : List BitStream) : BitStream :=
  match p with
  | .width .eq n => BitStream.falseIffEq n
  | .width .neq n => BitStream.falseIffNeq n
  | .width .lt n => BitStream.falseIffLt n
  | .width .le n => BitStream.falseIffLe n
  | .width .gt n => BitStream.falseIffGt n
  | .width .ge n => BitStream.falseIffGe n
  | lor p q => Predicate.evalLor (p.eval vars) (q.eval vars)
  | land p q => Predicate.evalLand (p.eval vars) (q.eval vars)
  | binary .eq t₁ t₂ => Predicate.evalEq (t₁.eval vars) (t₂.eval vars)

  | binary .neq t1 t2 => Predicate.evalNeq (t1.eval vars) (t2.eval vars)
  | binary .ult t₁ t₂ => Predicate.evalUlt (t₁.eval vars) (t₂.eval vars)
  | binary .ule t₁ t₂ =>
     Predicate.evalLor
       (Predicate.evalEq (t₁.eval vars) (t₂.eval vars))
       (Predicate.evalUlt (t₁.eval vars) (t₂.eval vars))
  | binary .slt t₁ t₂ => Predicate.evalSlt (t₁.eval vars) (t₂.eval vars)
  | binary .sle t₁ t₂ => Predicate.evalLor
       (Predicate.evalEq (t₁.eval vars) (t₂.eval vars))
       (Predicate.evalSlt (t₁.eval vars) (t₂.eval vars))

section Predicate

end Predicate

theorem Predicate.evalUlt_denote_true_iff {w : Nat} (a b : Term) (vars : List (BitVec w)) :
    evalUlt (a.eval (List.map .ofBitVecSext vars)) (b.eval (List.map .ofBitVecSext vars)) w = true ↔
      (Term.denote w b vars) ≤ (Term.denote w a vars) :=

:= by
  obtain ⟨h₁, h₂⟩ := evalUlt_denote_false_iff a b vars
  constructor
  · intros h
    by_contra h'
    simp at h'
    specialize (h₂ h')
    simp [h₂] at h
  · intros h
    by_contra h'
    simp at h'
    specialize (h₁ h')
    bv_omega


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

HVector.decidableEqReg

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

idxOf?

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

map

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

HVectorLiteral

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

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

R

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

Ty

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

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

ExplicitISL

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

ExpTy

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

ISL

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

ISLTy

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

ISLOp

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

RegIndex

def RegIndex := Fin RegFile.numRegisters

RegFile.numRegisters

def RegFile.numRegisters := 32

Dialect

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

Op

inductive Op (q : Nat) (n : Nat)
  | add : Op q n
  | sub : Op q n
  | mul : Op q n
  | mul_constant : Op q n

  | leading_term : Op q n
  | monomial : Op q n
  | monomial_mul : Op q n
  | from_tensor : Op q n
  | to_tensor : Op q n
  | const (c : R q n) : Op q n
  | const_int (c : Int) : Op q n
  | const_idx (i : Nat) : Op q n

RegFile

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

RegFile.registerWidth

def RegFile.registerWidth := 64

read

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

write

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

instDCTyDenote

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

DCPlus

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

def_signature for DCPlus where

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

sndVal

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)

Ty

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

Var

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

Context

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

VarName

def VarName : Type := String

hasType

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

Expr.decidableEq

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

AffineExpr

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

mux

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)

branch

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)

merge

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)

cMerge

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)

fork

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)

fstVal

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)

forkVal

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)

join

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)

Op

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

fstValPure

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)

sink

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)

supp

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)

muxVal

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)

sndValPure

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)

sourceOnes

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

snd

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)

source

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

fst

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)

TokenStream

def TokenStream := Stream Unit

Stream

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

ValueStream

def ValueStream := Stream

VariadicValueStream

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

branch

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

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

cMerge

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

fork

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

forkVal

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)

join

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

corec

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

merge

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

mux

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)

muxVal

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)

not

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)

sink

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

source

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

sourceOnes

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

supp

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

mapOpt

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

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

BabyBear

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

Ty

inductive Ty
  | felt

instDCTyDenote

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

branch

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

Term

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

pack

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

select

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)

unpack

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)

DC

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

def_signature for DC where

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

Op

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

branch

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)

fork

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)

fst

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)

sink

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)

merge

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)

pack

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)

select

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)

source

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

join

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)

fstVal

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)

snd

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)

unpack

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)

sndVal

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)

pair

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

ExTy

inductive ExTy
  | nat
  | bool
  deriving DecidableEq

instance : TyDenote Ty where toType := Ty.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

Index

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

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

def_denote for TestDialect

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

Op

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

TestDialect

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

Ty

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

IntW

def IntW w := PoisonOr <| BitVec w

PoisonOr

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

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

SimpleReg.int

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

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

ExTy

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

instHandshakeTyDenote

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

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

Handshake

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

def_signature for Handshake where

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)

Ty

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

sync

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)

snd

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)

Op

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

merge

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)

supp

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)

branch

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)

fst

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)

altMerge

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)

branch

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

controlMerge

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

fork

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

join

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

merge

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

mux

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

not

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

sink

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

supp

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

sync

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

Op

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

ExTy

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

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

Op

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

VariadicDialect

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

def_signature for VariadicDialect where

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

Ty

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)

TyDenote

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

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

Ty

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

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

SLLVMOp

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

Ty.mem

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

Ty.ptr

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

MemorySSAState

def MemorySSAState := PoisonOr MemoryState

PoisonOr

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

MemoryState

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

GlobalState

structure GlobalState where
  alloc : AllocState
  mem : MemoryState

BlockId

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

Block

inductive Block where

  | dead

  | live (b : LiveBlock)

LiveBlock

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

Ptr

def Ptr : Type := PoisonOr Pointer

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

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

Ctxt

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

  deriving Repr, Lean.ToExpr, DecidableEq

Ctxt.cons

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

Ctxt.Var

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

Ctxt.Valuation

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

Ctxt.Valuation.nil

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

Ctxt.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

Ctxt.Valuation.reassignVars

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

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


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

syntax "max" : MLIR.Pretty.uniform_op

syntax "max" : MLIR.Pretty.uniform_op

syntax "slt" : MLIR.Pretty.uniform_op

syntax "xor" : MLIR.Pretty.uniform_op

macro_rules
  | `(mlir_op| $res:mlir_op_operand = const ($x)
      $[: $outer_type]? ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _)
      `(mlir_op| $res:mlir_op_operand = "const"()
          {val = $x:num : $outer_type} : ($outer_type) -> ($outer_type) )

macro_rules
  | `(mlir_op| $res:mlir_op_operand = li ($x)
     $[: $outer_type]?  ) => do
      let outer_type ← outer_type.getDM `(mlir_type| _ )
      `(mlir_op| $res:mlir_op_operand = "li"()
          {imm = $x:num : $outer_type } : ($outer_type) -> ($outer_type))

macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithImmediate $reg1 , $x : $t)  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {imm = $x:num : $t} : ($t) -> ($t) )

macro_rules
| `(mlir_op| $res:mlir_op_operand = $op1:MLIR.Pretty.RV.opWithShamt $reg1 , $x  : $t )  => do
    let some opName := MLIR.EDSL.Pretty.extractOpName op1.raw
      | Macro.throwUnsupported
    `(mlir_op| $res:mlir_op_operand = $opName ($reg1) {shamt = $x:num : $t}  : ($t) -> ($t) )

carry

def carry (initCarry : Bool) (x y : BitStream) : BitStream :=
  fun n => (addAux' initCarry x y n).2

BitStream

def BitStream : Type := Nat → Bool

addAux'

def addAux' (carryIn : Bool) (x y : BitStream) (i : Nat) :  Bool × Bool :=
  let carryIn : Bool := match i with
    | 0 => carryIn
    | i + 1 => (addAux' carryIn x y i).2
  Prod.swap (BitVec.adcb (x i) (y i) carryIn)

BitVec.ofFn

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

f

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

worklistRun

def worklistRun (final : S → Bool) (inits : Array S)
    (hinits : inits.toList.Nodup) (f : S → Array (BitVec n × S)) : CNFA n :=
  ⟨worklistRun' _ S final inits hinits f, worklistRun'_wf (BitVec n) S⟩

worklistRun'

def worklistRun' (final : S → Bool) (inits : Array S) (hinits : inits.toList.Nodup) (f : S → Array (A × S)) : RawCNFA A :=
  let st0 := worklist.initState _ _ inits hinits final
  go st0
where go (st0 : worklist.St A S) : RawCNFA A :=
  if hemp : st0.worklist.isEmpty then st0.m else
  let sa? := st0.worklist.back?
  match heq : sa? with
  | some sa =>
    let wl := st0.worklist.pop
    let st1 := { st0 with worklist := wl,
                          worklist_nodup := by admit /- proof elided -/

worklist.St

structure worklist.St where
  m : RawCNFA A
  map : Std.HashMap S State := ∅
  worklist : Array S := ∅
  worklist_nodup : worklist.toList.Nodup
  worklist_incl : ∀ sa ∈ worklist, sa ∈ map

worklist.initState

def worklist.initState (inits : Array S) (hinits : inits.toList.Nodup) (final? : S → Bool) : worklist.St A S :=
  let m := RawCNFA.empty (A := A)
  let mapm := inits.foldl (init := (Std.HashMap.emptyWithCapacity, m)) fun (map, m) sa =>
    let (s, m) := m.newState
    let m := m.addInitial s
    let m := if final? sa then m.addFinal s else m
    (map.insert sa s, m)
  let map := mapm.1
  let m := mapm.2
  let worklist_incl : ∀ sa ∈ inits, sa ∈ map :=

RawCNFA.statesFinset

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

State

abbrev State := Nat

RawCNFA.empty

def RawCNFA.empty : RawCNFA A := {
  stateMax := 0
  initials := ∅
  finals := ∅
  trans := ∅
}

processOneElem

def processOneElem (final : S → Bool) (s : State) (st : worklist.St A S) : A × S → worklist.St A S :=
  fun (a', sa') =>
    let (s', st') := st.addOrCreateState _ _ (final sa') sa'
    let m := st'.m.addTrans a' s s'
    { st' with m }

worklist.St.addOrCreateState

def worklist.St.addOrCreateState (st : worklist.St A S) (final? : Bool) (sa : S) : State × worklist.St A S :=
  match heq : st.map[sa]? with
  | some s => (s, st)
  | none =>
    let (s, m) := st.m.newState
    let m := if final? then m.addFinal s else m
    let map := st.map.insert sa s
    let worklist := st.worklist.push sa
    have worklist_nodup : worklist.toList.Nodup := by admit /- proof elided -/

CNFA

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

FSM

structure FSM (arity : Type) : Type 1 where

  ( α  : Type )
  [ i : FinEnum α ]
  [ h : Hashable α ]
  [ dec_eq : DecidableEq α ]

  ( initCarry : α → Bool )

  outputCirc : Circuit (α ⊕ arity)
  nextStateCirc : α → Circuit (α ⊕ arity)

Circuit

inductive Circuit (α : Type u) : Type u
  | tru : Circuit α
  | fals : Circuit α

  | var : (positive: Bool) → α → Circuit α
  | and : Circuit α → Circuit α → Circuit α
  | or : Circuit α → Circuit α → Circuit α
  | xor : Circuit α → Circuit α → Circuit α
deriving Repr, DecidableEq

Var

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

sub

def sub (x y : BitStream) : BitStream :=
  fun n => (subAux x y n).1

subAux

def subAux (x y : BitStream) : Nat → Bool × Bool
  | 0 => (xor (x 0) (y 0), !(x 0) && y 0)
  | n+1 =>
    let borrow := (subAux x y n).2
    let a := x (n + 1)
    let b := y (n + 1)
    (xor a (xor b borrow), !a && b || ((!(xor a b)) && borrow))

ofTerm

abbrev ofTerm (t : Term) : FSM (Fin t.arity) := termEvalEqFSM t |>.toFSM

Term

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

termEvalEqFSM

def termEvalEqFSM : ∀ (t : Term), FSMTermSolution t
  | ofNat n =>
    { toFSM := FSM.ofNat n,
      good := by admit /- proof elided -/

or

def or : FSM Bool :=
  { α := Empty,
    initCarry := Empty.elim,
    outputCirc := Circuit.var true (inr true) ||| Circuit.var true (inr false),
    nextStateCirc := fun a => a.elim
  }

shiftLeft

def shiftLeft (n : Nat) : FSM Unit :=
  match n with
  | 0 => FSM.id
  | n + 1 => composeUnaryAux (FSM.ls false) (shiftLeft n)

id

def id : FSM Unit := {
 α := Empty,
 initCarry := Empty.elim,
 outputCirc := Circuit.var true (inr ()),
 nextStateCirc := Empty.elim
}

ls

def ls (b : Bool) : FSM Unit :=
  { α := Unit,
    initCarry := fun _ => b,
    nextStateCirc := fun () => Circuit.var true (inr ()),
    outputCirc := Circuit.var true (inl ())
  }

composeUnaryAux

def composeUnaryAux
    (p : FSM Unit)
    (q : FSM arity) :
    FSM arity :=
  p.compose
    arity
    _
    (λ _ => id)
    (λ _ => q)

compose

def compose [FinEnum arity] [DecidableEq arity] [Hashable arity]
    (new_arity : Type)
    (q_arity : arity → Type)
    (vars : ∀ (a : arity), q_arity a → new_arity)

    (q : ∀ (a : arity), FSM (q_arity a)) :
    FSM new_arity :=
  { α := p.α ⊕ (Σ a, (q a).α),
    i := by admit /- proof elided -/

and

def and : FSM Bool :=
  { α := Empty,
    initCarry := Empty.elim,
    nextStateCirc := fun a => a.elim,
    outputCirc := Circuit.var true (inr true) &&& Circuit.var true (inr false),
  }

xor

def xor : FSM Bool :=
  { α := Empty,
    initCarry := Empty.elim,
    outputCirc := Circuit.var true (inr true) ^^^ Circuit.var true (inr false),
    nextStateCirc := Empty.elim
  }

neg

def neg : FSM Unit :=
  { α := Unit,
    i := by admit /- proof elided -/

composeBinary

def composeBinary
    (p : FSM Bool)
    {t₁ t₂ : Term}
    (q₁ : FSMTermSolution t₁)
    (q₂ : FSMTermSolution t₂) :
    FSM (Fin (max t₁.arity t₂.arity)) := composeBinaryAux p q₁.toFSM q₂.toFSM

composeBinaryAux

def composeBinaryAux
    (p : FSM Bool)
    (q₁ : FSM (Fin a₁))
    (q₂ : FSM (Fin a₂)) :
    FSM (Fin (max a₁ a₂)) :=
  p.compose (Fin (max a₁ a₂))
    (λ b => Fin (cond b a₁ a₂))
    (λ b i => Fin.castLE (by admit /- proof elided -/
    ) i)
    (λ b => match b with
      | true => q₁
      | false => q₂)

FSMTermSolution

structure FSMTermSolution (t : Term) extends FSM (Fin t.arity) where
  ( good : t.evalFin = toFSM.eval )

Term.evalFin

@[simp] def Term.evalFin (t : Term) (vars : Fin (arity t) → BitStream) : BitStream :=
  match t with
  | var n => vars (Fin.last n)
  | zero    => BitStream.zero
  | one     => BitStream.one
  | negOne  => BitStream.negOne
  | ofNat n => BitStream.ofNat n
  | and t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ &&& x₂
  | or t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ||| x₂
  | xor t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ^^^ x₂
  | not t     => ~~~(t.evalFin vars)
  | add t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ + x₂
  | sub t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ - x₂
  | neg t       => -(Term.evalFin t vars)

  | shiftL t n  => BitStream.shiftLeft (Term.evalFin t vars) n

Predicate.evalFin

@[simp] def Predicate.evalFin (p : Predicate) (vars : Fin (arity p) → BitStream) : BitStream :=
match p with
| .width .eq n => BitStream.falseIffEq n
| .width .neq n => BitStream.falseIffNeq n
| .width .lt n => BitStream.falseIffLt n
| .width .le n => BitStream.falseIffLe n
| .width .gt n => BitStream.falseIffGt n
| .width .ge n => BitStream.falseIffGe n
| .binary .eq t₁ t₂ =>
    let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    Predicate.evalEq x₁ x₂
| .binary .neq t₁ t₂  =>
    let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    Predicate.evalNeq x₁ x₂
| .land p q =>

  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLand x₁ x₂
| .lor p q =>

  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLor x₁ x₂
| .binary .slt p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalSlt x₁ x₂
| .binary .sle p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLor (Predicate.evalSlt x₁ x₂) (Predicate.evalEq x₁ x₂)
| .binary .ult p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  (Predicate.evalUlt x₁ x₂)
| .binary .ule p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLor (Predicate.evalUlt x₁ x₂) (Predicate.evalEq x₁ x₂)

Predicate.evalUlt

def Predicate.evalUlt (t₁ t₂ : BitStream) : BitStream := (~~~ (t₁.borrow t₂)).concat true

borrow

def borrow (x y : BitStream) : BitStream :=
  fun n => (subAux x y n).2

Predicate.evalLor

def Predicate.evalLor (t₁ t₂ : BitStream) : BitStream := (t₁ &&& t₂)

Predicate.evalSlt

def Predicate.evalSlt (t₁ t₂ : BitStream) : BitStream :=
    (((Predicate.evalUlt t₁ t₂)) ^^^ (Predicate.evalMsbEq t₁ t₂))

Predicate.evalMsbEq

def Predicate.evalMsbEq (t₁ t₂ : BitStream) : BitStream :=
  (t₁ ^^^ t₂).concat false

Predicate.evalLand

def Predicate.evalLand (t₁ t₂ : BitStream) : BitStream := (t₁ ||| t₂)

Predicate.evalNeq

def Predicate.evalNeq (t₁ t₂ : BitStream) : BitStream := (t₁.nxor t₂).concat true |>.scanAnd

nxor

def nxor (a b : BitStream) : BitStream := fun i => a i == b i

scanAnd

def scanAnd (s : BitStream) : BitStream := scanl true Bool.and s

scanl

abbrev scanl (init : Bool) (f : Bool → Bool → Bool) (s : BitStream) : BitStream :=
  fun n => match n with
    | 0 => f init (s 0)
    | n+1 => f (scanl init f s n) (s (n + 1))

Predicate.evalEq

def Predicate.evalEq (t₁ t₂ : BitStream) : BitStream := (t₁ ^^^ t₂).concat false |>.scanOr

scanOr

def scanOr (s : BitStream) : BitStream := scanl false Bool.or s

Predicate

inductive Predicate : Type where

| width (wp : WidthPredicate) (n : Nat) : Predicate
| binary (p : BinaryPredicate) (t₁ t₂ : Term)
| land  (p q : Predicate) : Predicate
| lor (p q : Predicate) : Predicate
deriving Repr, Inhabited, Lean.ToExpr

falseIffNeq

abbrev falseIffNeq (n : Nat) : BitStream := fun i => decide (i == n)

falseIffLt

abbrev falseIffLt (n : Nat) : BitStream := fun i => decide (i ≥ n)

falseIffLe

abbrev falseIffLe (n : Nat) : BitStream := fun i => decide (i > n)

falseIffGe

abbrev falseIffGe (n : Nat) : BitStream := fun i => decide (i < n)

falseIffEq

abbrev falseIffEq (n : Nat) : BitStream := fun i => decide (i != n)

falseIffGt

abbrev falseIffGt (n : Nat) : BitStream := fun i => decide (i ≤ n)

Term.width

def Term.width (t : Term) : WidthExpr :=
  match t with

  | .ofNat w _n => w
  | .var _v w => w
  | .add w _a _b => w
  | .zext _a wnew => wnew
  | .setWidth _a wnew => wnew
  | .sext _a wnew => wnew
  | .bor w _a _b => w
  | .band w _a _b => w
  | .bxor w _a _b => w
  | .bnot w _a => w
  | .boolVar _v => WidthExpr.const 1
  | .boolConst _b => WidthExpr.const 1
  | .shiftl w _a _k => w
  | .bvOfBool _b => WidthExpr.const 1
  | binWidthRel _k wa wb => WidthExpr.const 0
  | binRel _k w _a _b => w
  | or _p1 _p2 => WidthExpr.const 0
  | and _p1 _p2 => WidthExpr.const 0
  | pvar _v => WidthExpr.const 0
  | boolBinRel _k _a _b => WidthExpr.const 0

Term.arity

@[simp] def Term.arity : Term → Nat
| (var n) => n+1
| zero => 0
| one => 0
| negOne => 0
| ofNat _ => 0
| Term.and t₁ t₂ => max (arity t₁) (arity t₂)
| Term.or t₁ t₂ => max (arity t₁) (arity t₂)
| Term.xor t₁ t₂ => max (arity t₁) (arity t₂)
| Term.not t => arity t
| add t₁ t₂ => max (arity t₁) (arity t₂)
| sub t₁ t₂ => max (arity t₁) (arity t₂)
| neg t => arity t

| shiftL t .. => arity t

negOne

abbrev negOne : BitStream := fun _ => true

shiftLeft

def shiftLeft (x : BitStream) (k : Nat) : BitStream :=
  fun i => if i < k then false else x (i - k)

ofNat

def ofNat (x : Nat) : BitStream :=
  Nat.testBit x

one

abbrev one    : BitStream := (· == 0)

zero

abbrev zero   : BitStream := fun _ => false

ofNat

def ofNat (n : Nat)  : FSM (Fin 0) :=
  match hn : n with
  | 0 => FSM.zero

  | n' + 1 =>
    let bit := n.testBit 0
    let m := n / 2
    have h : m < n := by admit /- proof elided -/

zero

def zero : FSM (Fin 0) :=
  { α := Empty,
    initCarry := Empty.elim,
    nextStateCirc := Empty.elim,
    outputCirc := Circuit.fals
  }

composeUnary

def composeUnary
    (p : FSM Unit)
    {t : Term}
    (q : FSMTermSolution t) :
    FSM (Fin t.arity) := composeUnaryAux p q.toFSM

one

def one : FSM (Fin 0) :=
  { α := Unit,
    i := by admit /- proof elided -/

var

def var (n : ℕ) : FSM (Fin (n+1)) :=
  { α := Empty,
    i := by admit /- proof elided -/

add

def add : FSM Bool :=
  { α := Unit,
    initCarry := λ _ => false,
    nextStateCirc := fun () =>
      Circuit.var true (inr true) &&& Circuit.var true (inr false) |||
      Circuit.var true (inr true) &&& Circuit.var true (inl ()) |||
      Circuit.var true (inr false) &&& Circuit.var true (inl ()),
    outputCirc := Circuit.var true (inr true) ^^^
                  Circuit.var true (inr false) ^^^
                  Circuit.var true (inl ()),
  }

negOne

def negOne : FSM (Fin 0) :=
  { α := Empty,
    i := by admit /- proof elided -/

sub

def sub : FSM Bool :=
  { α := Unit,
    initCarry := fun _ => false,
    outputCirc := Circuit.var true (inr true) ^^^
                  Circuit.var true (inr false) ^^^
                  Circuit.var true (inl ()),
    nextStateCirc := fun _ =>
      (Circuit.var false (inr true) &&& Circuit.var true (inr false)) |||
      (Circuit.var false (inr true) ^^^ Circuit.var true (inr false)) &&&
      (Circuit.var true (inl ()))
  }

not

def not : FSM Unit :=
  { α := Empty,
    initCarry := Empty.elim,
    nextStateCirc := Empty.elim,
    outputCirc := Circuit.var false (inr ())
  }

add

def add (x y : BitStream) : BitStream :=
  fun n => (addAux x y n).1

addAux

def addAux (x y : BitStream) (i : Nat) :  Bool × Bool :=
  let carryIn : Bool := match i with
    | 0 => false
    | i + 1 => (addAux x y i).2
  Prod.swap (BitVec.adcb (x i) (y i) carryIn)

neg

def neg (x : BitStream) : BitStream :=
  fun n => (negAux x n).1

negAux

def negAux (x : BitStream) : Nat → Bool × Bool
  | 0 => (x 0, !(x 0))
  | n+1 =>
    let borrow := (negAux x n).2
    let a := x (n + 1)
    (xor (!a) borrow, !a && borrow)

def liftMaxSuccSucc2 (n m : Nat) : Fin (m + 1) → Fin (max n m + 3) :=
  fun k => if _ : k = Fin.last m then max n m + 1 else k.castLE (by admit /- proof elided -/
  )

def liftMaxSuccSucc1 (n m : Nat) : Fin (n + 1) → Fin (max n m + 3) :=
  fun k => if _ : k = Fin.last n then (max n m).cast else k.castLE (by admit /- proof elided -/
  )

def liftLast3 n : Fin 3 → Fin (n + 3)
| 0 => n
| 1 => n + 1
| 2 => Fin.last (n + 2)

CNFA.inter

def CNFA.inter (m1 m2 : CNFA n) : CNFA n := product (fun b1 b2 => b1 && b2) m1 m2

product

def product (final? : Bool → Bool → Bool) (m₁ m₂ : CNFA n) : CNFA n :=
  worklistRun (m₁.m.states × m₂.m.states) final (product.inits m₁ m₂)
    (by admit /- proof elided -/
    ) f
where final (ss : m₁.m.states × m₂.m.states) := final? (ss.1 ∈ m₁.m.finals) (ss.2 ∈ m₂.m.finals)
      f (ss : m₁.m.states × m₂.m.states) :=
        let (s1, s2) := ss
        (FinEnum.toList (α := BitVec n)).foldl (init := Array.empty) fun as a =>
          product.prodArray' (λ s₁ s₂ ↦ (a, (s₁, s₂)))
            (fun s' => m₁.wf.trans_tgt_lt (s := s1) (a := a)) (fun s' => m₂.wf.trans_tgt_lt (s := s2) (a := a)) as

product.prodArray'

@[inline]
def product.prodArray' (a : Array γ) :=
  m₁.attachWith _ hm₁ |>.fold (init := a) fun is s1 =>
    m₂.attachWith _ hm₂ |>.fold (init := is) fun is s2 =>
      is.push (f s1 s2)

def product.inits_nodup : inits m₁ m₂ |>.toList.Nodup :=

product.inits

def product.inits (m₁ m₂ : CNFA n) :=
  product.prodArray Prod.mk @m₁.wf.initials_lt @m₂.wf.initials_lt

product.prodArray

@[inline]
def product.prodArray := prodArray' f hm₁ hm₂ (Array.emptyWithCapacity <| m₁.size * m₂.size)

def CNFA.inter_bv_language (m₁ m₂ : CNFA n) :
    m₁.bv_recognizes L₁ →
    m₂.bv_recognizes L₂ →
    (m₁.inter m₂).bv_recognizes (L₁ ∩ L₂) :=

HashSet.inter

def HashSet.inter [BEq A] [Hashable A] (m1 m2 : Std.HashSet A) : Std.HashSet A :=
  m1.fold (init := ∅) fun mi x => if m2.contains x then mi.insert x else mi

NFA'

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

BitVecs

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

eval

def eval (x : arity → BitStream) : BitStream :=
  fun n => (p.nextBit (p.carry x n) (fun i => x i n)).2

nextBit

def nextBit : p.State → (arity → Bool) → p.State × Bool :=
  fun carry inputBits =>
    let input := Sum.elim carry inputBits
    let newState : p.State  := fun (a : p.α) => (p.nextStateCirc a).eval input
    let outBit : Bool       := (p.outputCirc).eval input
    (newState, outBit)

State

abbrev State : Type := p.α → Bool

carry

def carry (x : arity → BitStream) : ℕ → p.State
  | 0 => p.initCarry
  | n+1 => (p.nextBit (carry x n) (fun i => x i n)).1

carryBV

def carryBV (x : ar → BitVec w) : p.State :=
  p.carry (fun ar => .ofBitVecSext (x ar)) w

evalBV

def evalBV {w} (x : ar → BitVec w) : BitVec w :=
  BitVec.ofFn fun k => p.eval (fun ar => .ofBitVecSext (x ar)) k

ofBitVecSext

abbrev ofBitVecSext {w} (x : BitVec w) : BitStream :=
  fun i => if i < w then x.getLsbD i else x.msb

Term.language

def Term.language (t : Term) : Set (BitVecs (t.arity + 1)) :=
  { bvs : BitVecs (t.arity + 1) | t.evalFinBV (fun n => bvs.bvs.get n) = bvs.bvs.get t.arity }

Formula.arity

@[simp]
def Formula.arity : Formula → Nat
| width _ _ => 0
| atom _ t1 t2 => max t1.arity t2.arity
| msbSet t => t.arity
| unop _ φ => φ.arity
| binop _ φ1 φ2 => max φ1.arity φ2.arity

Term.evalFinBV

@[simp] def Term.evalFinBV (t : Term) (vars : Fin (arity t) → BitVec w) : BitVec w :=
  match t with
  | .var n => vars (Fin.last n)
  | .zero    => BitVec.zero w
  | .one     => 1
  | .negOne  => -1
  | .ofNat n => BitVec.ofNat _ n
  | .and t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ &&& x₂
  | .or t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ||| x₂
  | .xor t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ^^^ x₂
  | .not t     => ~~~(t.evalFinBV vars)

  | .add t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ + x₂
  | .sub t₁ t₂ =>
      let x₁ := t₁.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFinBV (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ - x₂
  | .neg t       => -(t.evalFinBV vars)
  | .shiftL a n => (a.evalFinBV vars) <<< n

enc

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

BitVecs'

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

dec

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

accepts

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

accepts'

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

def worklistRun_spec : (worklistRun S final inits hinits f |>.Sim $ nfa' inits final f) :=
  worklistRun'_spec inits final f

nfa'

def nfa' : NFA' n :=
  { σ := _, M := nfa inits final f }

nfa

def nfa : NFA A S where
  start := { sa | sa ∈ inits }
  accept := { sa | final sa }
  step sa a := { sa' | (a, sa') ∈ f sa }

def worklistRun'_spec :
    (worklistRun' A S final inits hinits f |>.Sim $ nfa inits final f) :=

StInv

structure StInv (m : RawCNFA A) (map : Std.HashMap S State) where
  wf : m.WF
  map_states : ∀ (sa : S) s, map[sa]? = some s → s ∈ m.states
  map_surj : ∀ s : m.states, ∃ (sa : S), map[sa]? = some s.val
  map_inj : ∀ {s} {sa sa' : S}, map[sa]? = some s → map[sa']? = some s → sa = sa'

worklist.St.D

def worklist.St.D (st : worklist.St A S) : Set S := st.visited

worklist.St.visited

def worklist.St.visited (st : worklist.St A S) : Set S := { s : S | s ∈ st.map ∧ s ∉ st.worklist }

def worklistGo_spec {st : worklist.St A S} (inv : StInv A S st.m st.map) :
    st.sim inits final f ∅ →
    (worklistRun'.go A S final f st |>.Sim $ nfa inits final f) :=

worklist.St.rel

def worklist.St.rel (st : worklist.St A S) : SetRel State S := {(s, sa) | st.map[sa]? = some s }

def processOneElem_mot (s : State) (sa : S) (n : ℕ) (st : worklist.St A S) : Prop :=
  st.map[sa]? = some s ∧
  sa ∈ st.visited ∧
  StInv A S st.m st.map ∧
  st.sim inits final f  {(sa1, a, sa') | sa1 = sa ∧ ∃ k ≥ n, (f sa)[k]? = some (a, sa') }

worklist.St.sim

abbrev worklist.St.sim {st : worklist.St A S} (T : Set (S × A × S)) :=
  st.m.Simul (nfa inits final f) st.rel st.D T

RawCNFA.Sim

def RawCNFA.Sim (m : RawCNFA A) (A : NFA A S) := ∃ R, RawCNFA.Simul m A R ⊤ ∅

RawCNFA.Simul

structure RawCNFA.Simul (m : RawCNFA A) (M : NFA A Q) (R : SetRel State Q) (D : Set Q) (T : Set (Q × A × Q)) where
  accept {s q} : s ~[R] q → (s ∈ m.finals ↔ q ∈ M.accept)
  initial₁ {s} : s ∈ m.initials → ∃ q ∈ M.start, s ~[R] q
  initial₂ {q} : q ∈ M.start → ∃ s ∈ m.initials, s ~[R] q
  trans_match₁ {s s' a q} : s ~[R] q → s' ∈ m.tr s a → ∃ q', q' ∈ M.step q a ∧ s' ~[R] q'
  trans_match₂ {s a q q'} : s ~[R] q → q' ∈ M.step q a → q ∈ D → (q, a, q') ∉ T → ∃ s', s' ∈ m.tr s a ∧ s' ~[R] q'

RawCNFA.SimulFun

structure RawCNFA.SimulFun (m : RawCNFA A) (M : NFA A Q) (f : m.states ≃ Q)  where
  accept {q} : ((f.invFun q).val ∈ m.finals ↔ q ∈ M.accept)
  initial {q} : q ∈ M.start ↔ (f.invFun q).val ∈ m.initials
  trans_match {a q q'} : q' ∈ M.step q a ↔ (f.invFun q').val ∈ m.tr (f.invFun q) a

RawCNFA.tr

@[inline]
def RawCNFA.tr (m : RawCNFA A) s a := m.trans.getD (s, a) ∅

R

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

CNFA.Sim

def CNFA.Sim (m : CNFA n) (M : NFA' n) :=
  m.m.Sim M.M

def CNFA.bv_recognizes (m : CNFA n) (L : Set (BitVecs n)) :=
  ∃ L', m.recognizes L' ∧ L = dec '' L'

RawCNFA.recognizes

def RawCNFA.recognizes (m : RawCNFA A) (L : Language A) :=
  ∃ (σ : Type) (M : NFA A σ), m.Sim M ∧ M.accepts = L

CNFA.recognizes

def CNFA.recognizes (m : CNFA n) (L : Language (BitVec n)) :=
  ∃ (M : NFA' n), m.Sim M ∧ M.M.accepts = L

BitVecs.cast

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

Valuation.cons

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

CNFA.minimize

def CNFA.minimize (m : CNFA n) : CNFA n :=
  let mᵣ := m.reverse.determinize
  mᵣ.reverse.determinize

CNFA.determinize

def CNFA.determinize (m : CNFA n) : CNFA n :=
  worklistRun (BitVec m.m.stateMax)
    (fun ss => ss.any fun n b => b == true && n ∈ m.m.finals)
    (determinize.inits m)
    (by admit /- proof elided -/
    )
    f
where
  f := fun (ss : BitVec m.m.stateMax) =>
        (FinEnum.toList (BitVec n)).foldl (init := Array.empty) fun ts a =>
          let ss' := m.m.transSetBV ss a
          ts.push (a, ss')

CNFA.determinize.inits

def CNFA.determinize.inits (m : CNFA n) : Array (BitVec m.m.stateMax) :=
  #[BitVec.ofFn (fun n => n ∈ m.m.initials)]

CNFA.reverse

def CNFA.reverse (m : CNFA n) : CNFA n :=
  ⟨m.m.reverse, RawCNFA.reverse_spec m.wf |>.1⟩

RawCNFA.reverse

def RawCNFA.reverse (m : RawCNFA A) : RawCNFA A :=
  let m' := { stateMax := m.stateMax, trans := Std.HashMap.emptyWithCapacity m.trans.size, initials := m.finals, finals := m.initials}
  m.trans.fold (init := m') processState
where
  processState := fun m' (s, a) ss' =>
    ss'.fold (init := m') fun m' s' => m'.addTrans a s' s

CNFA.toNFA'

def CNFA.toNFA' (m : CNFA n) : NFA' n := ⟨_, m.toNFA⟩

CNFA.toNFA

def CNFA.toNFA (m : CNFA n) : NFA (BitVec n) m.m.states where
  start := { s | s.val ∈ m.m.initials }
  accept := { s | s.val ∈ m.m.finals }
  step s₁ a := { s₂ | s₂.val ∈ m.m.tr s₁.val a }

RawCNFA.states

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

reverse

def reverse (M : NFA' n) : NFA' n where
  σ := _
  M := M.M.reverse

def CNFA.determinize_spec (m : CNFA n)
  {M : NFA' n} (hsim : m.Sim M) :
    m.determinize.Sim M.determinize :=

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

structure _root_.SetRel.set_eq (R : SetRel α β) (A : Set α) (B : Set β) where
  fwd : a ∈ A → ∃ b ∈ B, a ~[R] b
  bwd : b ∈ B → ∃ a ∈ A, a ~[R] b

RawCNFA.lift

@[inline]
def RawCNFA.lift (m₁: RawCNFA (BitVec n1)) (f : Fin n1 → Fin n2) : RawCNFA (BitVec n2) :=
  let trans := (List.range m₁.stateMax).foldl (init := ∅) fun m2 s => processState m2 s
  { m₁ with trans }
where

CNFA.lift

@[inline]
def CNFA.lift (m: CNFA n1) (f : Fin n1 → Fin n2) : CNFA n2 :=
  ⟨m.m.lift f, m.m.lift_wf m.wf⟩

BitVecs.transport

def BitVecs.transport (f : Fin n → Fin m) (bvs : BitVecs m) : BitVecs n :=
  { w := bvs.w, bvs := bvs.bvs.transport f }

BitVec.transport

def BitVec.transport (f : Fin n2 → Fin n1) (bv : BitVec n1) : BitVec n2 :=
  BitVec.ofFn fun i => bv.getLsbD (f i)

List.Vector.transport

def List.Vector.transport (v : Vector α m) (f : Fin n → Fin m) : Vector α n :=
  Vector.ofFn fun i => v.get (f i)

BitVecs'.transport

def BitVecs'.transport (f : Fin n → Fin m) (bvs' : BitVecs' m): BitVecs' n :=
  bvs'.map fun bv => bv.transport f

RawCNFA.proj

@[inline]
def RawCNFA.proj (m1: RawCNFA (BitVec n1)) (f : Fin n2 → Fin n1) : RawCNFA (BitVec n2) :=
  let trans := m1.trans.keysArray.foldl (init := Std.HashMap.emptyWithCapacity) process
  { m1 with trans }
where

def CNFA.proj_spec (m : CNFA n2) (f : Fin n1 → Fin n2) {M : NFA' n2} :
    m.Sim M → (m.proj f |>.Sim (M.proj f)) :=

CNFA.proj

@[inline]
def CNFA.proj (m: CNFA n2) (f : Fin n1 → Fin n2) : CNFA n1 :=
  ⟨m.m.proj f, m.m.proj_wf m.wf⟩

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

ext

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

lemma bisim_comp (m : RawCNFA A) :
    m.Sim M₁ → M₁.Bisim M₂ → m.Sim M₂

lemma bisimul_comp {m : RawCNFA A} :
    m.Simul M₁ R₁ ⊤ ∅ → M₁.Bisimul R₂ M₂ →
    m.Simul M₂ (R₁.comp R₂) ⊤ ∅

lemma CNFA.bv_recognizes_equiv {m : CNFA n} :
    m.bv_recognizes L ↔ ∃ (M : NFA' n), m.Sim M ∧ M.accepts = L

@[simp]
lemma List.Vector.append_get_ge {x : List.Vector α n} {y : List.Vector α m} {i : Fin (n+m)} (hlt: n ≤ i) :
    (x ++ y).get i = y.get ((i.cast (Nat.add_comm n m) |>.subNat n hlt))

@[simp]
lemma List.Vector.append_get_lt {x : List.Vector α n} {y : List.Vector α m} {i : Fin (n+m)} (hlt: i < n) :
    (x ++ y).get i = x.get (i.castLT hlt)

lemma CNFA.minimize_bv_language {m : CNFA n} :
    m.bv_recognizes L → m.minimize.bv_recognizes L

lemma CNFA.minimize_language {m : CNFA n} :
    m.recognizes L → m.minimize.recognizes L

lemma CNFA.reverse_language {m : CNFA n} (hl : m.recognizes L) : m.reverse.recognizes L.reverse

lemma CNFA.reverse_spec {m : CNFA n} : m.reverse.Sim m.toNFA'.reverse

lemma RawCNFA.reverse_spec {m : RawCNFA A} (hwf : m.WF) :
    let m'

lemma RawCNFA.reverse_spec_procesState {m : RawCNFA A} (hwf : m.WF) s₀ a₀ ss' (hs₀ : s₀ ∈ m.states) :
    let motive m' ss'

lemma CNFA.determinize_language {m : CNFA n} :
    m.recognizes L → m.determinize.recognizes L

@[simp]
lemma CNFA.lift_bv_language {m : CNFA n1} {f : Fin n1 → Fin n2} :
    m.bv_recognizes L → (m.lift f |>.bv_recognizes (BitVecs.transport f ⁻¹' L))

lemma CNFA.lift_spec (m : CNFA n1) (f : Fin n1 → Fin n2) {M : NFA' n1} :
    m.Sim M → (m.lift f |>.Sim (M.lift f))

lemma CNFA.proj_bv_language {m : CNFA n2} {f : Fin n1 → Fin n2} :
    m.bv_recognizes L → (m.proj f |>.bv_recognizes (BitVecs.transport f '' L))

NFA.sa

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

NFA.correct

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

BVNRel

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

NFA'.sa

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

langRel

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

Alphabet

abbrev Alphabet (arity: Type) [FinEnum arity] := BitVec (FinEnum.card arity + 1)

finFunToBitVec

def finFunToBitVec [fe : FinEnum carry] (c : carry → Bool) : BitVec (FinEnum.card carry) :=
  BitVec.ofFn fun i => c (fe.equiv.invFun i)

bitVecToFinFun

def bitVecToFinFun [FinEnum ar] (bv : BitVec $ FinEnum.card ar) : ar → Bool :=
  fun c => bv[FinEnum.equiv.toFun c]

inFSMRel

def NFA'.ofFSM_sa (p : FSM arity) : (NFA'.ofFSM' p).sa := fun q _ bvn =>
    inFSMRel p bvn ∧ q = p.carryBV (fun ar => bvn.get (FinEnum.equiv.toFun ar))

def NFA'.ofFSM_correct (p : FSM arity) :
    (NFA'.ofFSM' p).correct (ofFSM_sa p) (fun _ bvn => inFSMRel p bvn) :=

NFA.msbState

inductive NFA.msbState : Type where
| i | f
deriving DecidableEq, Fintype

liftOp

def liftOp n : Fin (n + 1) → Fin (n + 3) :=
  fun k =>
    if k = n then Fin.last (n+2) else k.castLE (by admit /- proof elided -/
    )

@[simp]
def liftOp_unchanged (k : Fin n) : liftOp n k.castSucc = k.castLE (by simp) :=

TermBinop

inductive TermBinop where
| and | or | xor | add | sub

TermBinop.openTerm

def TermBinop.openTerm (op : TermBinop) : Term := op.subst (.var 0) (.var 1)

@[simp]
def TermBinop.openTerm_arity (op : TermBinop) : op.openTerm.arity + 1 = 3 :=

swapLastTwoBlock

def swapLastTwoBlock (x : Fin (n + 3)) : Fin (n + 3) :=
  if x = Fin.last (n+2) then n
  else if x = n+1 then Fin.last (n + 2)
  else if x = n then n + 1
  else x

TermUnop

inductive TermUnop where
| neg | not | shiftL (k : Nat)

TermUnop.openTerm

def TermUnop.openTerm (op : TermUnop) : Term :=
  match op with
  | .neg => .neg (.var 0)
  | .not => .not (.var 0)
  | .shiftL k => .shiftL (.var 0) k

def TermUnop.openTerm_arity (op : TermUnop) : op.openTerm.arity = 1 :=

TermUnop.subst

def TermUnop.subst (op : TermUnop) (t : Term) : Term :=
  match op with
  | .neg => .neg t
  | .not => .not t
  | .shiftL k => .shiftL t k

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

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

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

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

@[simp]
lemma NFA'.ofFSM_spec (t : Term) :
    (ofFSM (FSM.ofTerm t)).accepts = t.language

lemma CNFA.ofFSM_spec (p : FSM arity) :
    (CNFA.ofFSM p).Sim (NFA'.ofFSM p)

lemma CNFA.ofFSM_bv_language :
    (CNFA.ofFSM (FSM.ofTerm t)).bv_recognizes t.language

lemma TermBinop.subst_arity' {op : TermBinop} : (op.subst t₁ t₂).arity + 1= t₁.arity ⊔ t₂.arity + 1

@[simp]
lemma BitVecs.cast_eq  (x : BitVecs n) (h : n = n') : h ▸ x = x.cast h

lemma TermBinop.alt_lang {t₁ t₂ : Term} (op : TermBinop) :
  (op.subst_arity' ▸ (op.subst t₁ t₂).language) =
    let lop : Set (BitVecs 3) := op.openTerm_arity ▸ op.openTerm.language
    let lop' : Set (BitVecs ((t₁.arity ⊔ t₂.arity) + 3)) := lop.lift (liftLast3 (max t₁.arity t₂.arity))
    let l₁ := t₁.language.lift (liftMaxSuccSucc1 t₁.arity t₂.arity)
    let l₂ := t₂.language.lift (liftMaxSuccSucc2 t₁.arity t₂.arity)
    let l := l₁ ∩ l₂ ∩ lop'
    l.proj (liftOp _)

@[simp]
lemma TermUnop.subst_arity' {op : TermUnop} : (op.subst t).arity + 1 = t.arity + 1

import Batteries.Data.Fin.Basic

import Batteries.Data.Fin.Lemmas

import Blase.SingleWidth.Defs

import Blase.AutoStructs.Constructions

import Blase.AutoStructs.Defs

import Blase.AutoStructs.FiniteStateMachine

import Mathlib.Tactic.Ring

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

import Mathlib.Data.BitVec

import Mathlib.Tactic.FinCases

open Fin.NatCast

open Mathlib

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

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

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

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

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

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

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

section fsm

abbrev Alphabet (arity: Type) [FinEnum arity] := BitVec (FinEnum.card arity + 1)

variable {arity : Type} [FinEnum arity]

def finFunToBitVec [fe : FinEnum carry] (c : carry → Bool) : BitVec (FinEnum.card carry) :=
  BitVec.ofFn fun i => c (fe.equiv.invFun i)

def bitVecToFinFun [FinEnum ar] (bv : BitVec $ FinEnum.card ar) : ar → Bool :=
  fun c => bv[FinEnum.equiv.toFun c]

def NFA.ofFSM (p : FSM arity) : NFA (Alphabet arity) (p.α → Bool) where
  start := { q | q = p.initCarry }
  accept := ⊤
  step s a := {s' |
    let (s'', b) := p.nextBit s (bitVecToFinFun (a.truncate $ FinEnum.card arity))
    s' = s'' ∧ a.msb = b }

@[simp]
abbrev inFSMRel (p : FSM arity) {w} (bvn : List.Vector (BitVec w) _) :=
  bvn.get (Fin.last (FinEnum.card arity)) = p.evalBV (fun ar => bvn.get (FinEnum.equiv.toFun ar))

def NFA'.ofFSM_sa (p : FSM arity) : (NFA'.ofFSM' p).sa := fun q _ bvn =>
    inFSMRel p bvn ∧ q = p.carryBV (fun ar => bvn.get (FinEnum.equiv.toFun ar))

def NFA'.ofFSM_correct (p : FSM arity) :
    (NFA'.ofFSM' p).correct (ofFSM_sa p) (fun _ bvn => inFSMRel p bvn) :=

open BitStream in

def CNFA.ofFSM (p : FSM arity) : CNFA (FinEnum.card arity + 1) :=
  worklistRun (BitVec (FinEnum.card p.α))
    (fun _ => true)
    #[finFunToBitVec p.initCarry]
    (by admit /- proof elided -/
    )
    f
  where

end fsm

section nfas_relations

inductive NFA.msbState : Type where
| i | f
deriving DecidableEq, Fintype

end nfas_relations

def liftOp n : Fin (n + 1) → Fin (n + 3) :=
  fun k =>
    if k = n then Fin.last (n+2) else k.castLE (by admit /- proof elided -/
    )

@[simp]
def liftOp_unchanged (k : Fin n) : liftOp n k.castSucc = k.castLE (by simp) :=

inductive TermBinop where
| and | or | xor | add | sub

def TermBinop.subst (op : TermBinop) (t₁ t₂ : Term) : Term :=
  match op with
  | .and => .and t₁ t₂
  | .or => .or t₁ t₂
  | .xor => .xor t₁ t₂
  | .add => .add t₁ t₂
  | .sub => .sub t₁ t₂

def TermBinop.openTerm (op : TermBinop) : Term := op.subst (.var 0) (.var 1)

@[simp]
def TermBinop.openTerm_arity (op : TermBinop) : op.openTerm.arity + 1 = 3 :=

def TermBinop.termGadget (t : TermBinop) : CNFA 3 :=
  match t with
  | .and => FSM.ofTerm (.and (.var 0) (.var 1)) |> CNFA.ofFSM
  | .or => FSM.ofTerm (.or (.var 0) (.var 1)) |> CNFA.ofFSM
  | .xor => FSM.ofTerm (.xor (.var 0) (.var 1)) |> CNFA.ofFSM
  | .add => FSM.ofTerm (.add (.var 0) (.var 1)) |> CNFA.ofFSM
  | .sub => FSM.ofTerm (.sub (.var 0) (.var 1)) |> CNFA.ofFSM

def autOfTermBinop (op : TermBinop) (m₁ : CNFA (n + 1)) (m₂ : CNFA (m + 1)) : CNFA ((n ⊔ m) + 1 ) :=
  let mop : CNFA 3 := op.termGadget
  let f₁ := liftMaxSuccSucc1 n m
  let m1' := m₁.lift f₁
  let f₂ := liftMaxSuccSucc2 n m
  let m2' := m₂.lift f₂
  let mop := mop.lift $ liftLast3 (max (FinEnum.card (Fin n)) (FinEnum.card (Fin m)))
  let m := CNFA.inter m1' m2' |> CNFA.inter mop
  let mfinal := m.proj (liftOp _)
  mfinal.minimize

def swapLastTwoBlock (x : Fin (n + 3)) : Fin (n + 3) :=
  if x = Fin.last (n+2) then n
  else if x = n+1 then Fin.last (n + 2)
  else if x = n then n + 1
  else x

inductive TermUnop where
| neg | not | shiftL (k : Nat)

def TermUnop.openTerm (op : TermUnop) : Term :=
  match op with
  | .neg => .neg (.var 0)
  | .not => .not (.var 0)
  | .shiftL k => .shiftL (.var 0) k

def TermUnop.openTerm_arity (op : TermUnop) : op.openTerm.arity = 1 :=

def TermUnop.subst (op : TermUnop) (t : Term) : Term :=
  match op with
  | .neg => .neg t
  | .not => .not t
  | .shiftL k => .shiftL t k

lemma autOfTermBinop_bv_language op {t₁ t₂ : Term} (m₁ : CNFA (t₁.arity + 1)) (m₂ : CNFA (t₂.arity + 1)) :
    m₁.bv_recognizes t₁.language →
    m₂.bv_recognizes t₂.language →
    (autOfTermBinop op m₁ m₂ |>.bv_recognizes (op.subst_arity' ▸ (op.subst t₁ t₂).language)) :=

:= by
  rintro hrec₁ hrec₂
  simp [autOfTermBinop]
  rw [TermBinop.alt_lang]
  simp
  apply CNFA.minimize_bv_language
  apply CNFA.proj_bv_language
  ac_nf
  apply CNFA.inter_bv_language
  · apply CNFA.lift_bv_language
    rcases op <;> simp [TermBinop.termGadget, TermBinop.openTerm] <;> apply CNFA.ofFSM_bv_language
  · apply CNFA.inter_bv_language
    · apply CNFA.lift_bv_language; assumption
    · apply CNFA.lift_bv_language; assumption