syntax "max" : MLIR.Pretty.uniform_op

syntax "max" : 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) )

WidthPredicate

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

BitVecs

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

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

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

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)

BitVec.ofFn

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

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

Valuation.cons

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

BitVecs.cast

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

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

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

@[simp]
lemma BitVecs.transport_getElem {bvs : BitVecs m} (f : Fin n → Fin m) (i : Fin n) :
    (bvs.transport f).bvs.get i = bvs.bvs.get (f i)

Unop

inductive Unop
| neg
deriving Repr

Formula.language

lemma evalFin_eq {t : Term} {vars1 : Fin t.arity → BitVec w1} {vars2 : Fin t.arity → BitVec w2} :
    ∀ (heq : w1 = w2),
    (∀ n, vars1 n = heq ▸ vars2 n) →
    t.evalFinBV vars1 = heq ▸ t.evalFinBV vars2

@[simp]
lemma evalRelation_coe (rel : Relation) (bv1 bv2 : BitVec w1) (heq : w1 = w2) :
    evalRelation rel (heq ▸ bv1) (heq ▸ bv2) = evalRelation rel bv1 bv2

lemma helper1 : (k = 0) → (x ::ᵥ vs).get k = x

lemma helper2 : (k = 1) → (x ::ᵥ y ::ᵥ vs).get k = y

lemma formula_language_case_atom :
    let φ := Formula.atom rel t1 t2
    φ.language = λ (bvs : BitVecs φ.arity) => φ.sat (fun k => bvs.bvs.get k)

import Blase.AutoStructs.ForMathlib

import Blase.SingleWidth.Defs

open Fin.NatCast

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

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

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

def liftExcept2 n : Fin n → Fin (n + 2) :=
  fun k => Fin.castLE (by admit /- proof elided -/
  ) k

def liftMax1 (n m : Nat) : Fin n → Fin (max n m) :=
  fun k => k.castLE (by admit /- proof elided -/
  )

def liftMax2 (n m : Nat) : Fin m → Fin (max n m) :=
  fun k => k.castLE (by admit /- proof elided -/
  )

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

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 }

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

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

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

@[simp]
def Relation.language (rel : Relation) : Set (BitVecs 2) :=
  { bvs | evalRelation rel (bvs.bvs.get 0) (bvs.bvs.get 1) }

inductive Binop
| and | or | impl | equiv
deriving Repr

def evalBinop (op : Binop) (b1 b2 : Prop) : Prop :=
  match op with
  | .and => b1 ∧ b2
  | .or => b1 ∨ b2
  | .impl => b1 → b2
  | .equiv => b1 ↔ b2

def langBinop (op : Binop) (l1 l2 : Set (BitVecs n)) : Set (BitVecs n) :=
  match op with
  | .and => l1 ∩ l2
  | .or => l1 ∪ l2
  | .impl => l1ᶜ ∪ l2
  | .equiv => (l1ᶜ ∪ l2) ∩ (l2ᶜ ∪ l1)

inductive Unop
| neg
deriving Repr

inductive Formula : Type
| width : WidthPredicate → Nat → Formula
| atom : Relation → Term → Term → Formula
| msbSet : Term → Formula
| unop : Unop → Formula → Formula
| binop : Binop → Formula → Formula → Formula
deriving Repr

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

@[simp]
def WidthPredicate.sat (wp : WidthPredicate) (w n : Nat) : Bool :=
  match wp with
  | .eq => w = n
  | .neq => w ≠ n
  | .lt => w < n
  | .le => w ≤ n
  | .gt => w > n
  | .ge => w ≥ n

@[simp]
def Formula.sat {w : Nat} (φ : Formula) (ρ : Fin φ.arity → BitVec w) : Prop :=
  match φ with
  | .width wp n => wp.sat w n
  | .atom rel t1 t2 =>
    let bv1 := t1.evalFinBV (fun n => ρ $ Fin.castLE (by admit /- proof elided -/
    ) n)
    let bv2 := t2.evalFinBV (fun n => ρ $ Fin.castLE (by admit /- proof elided -/
    ) n)
    evalRelation rel bv1 bv2
  | .unop .neg φ => ¬ φ.sat ρ
  | .binop op φ1 φ2 =>
    let b1 := φ1.sat (fun n => ρ $ Fin.castLE (by admit /- proof elided -/
    ) n)
    let b2 := φ2.sat (fun n => ρ $ Fin.castLE (by admit /- proof elided -/
    ) n)
    evalBinop op b1 b2
  | .msbSet t => (t.evalFinBV ρ).msb

@[simp]
def _root_.Set.lift (f : Fin n → Fin m) (bvs : Set (BitVecs n)) : Set (BitVecs m) :=
  BitVecs.transport f ⁻¹' bvs

@[simp]
def _root_.Set.proj (f : Fin n → Fin m) (bvs : Set (BitVecs m)) : Set (BitVecs n) :=
  BitVecs.transport f '' bvs

@[simp]
def langMsb : Set (BitVecs 1) := { bvs | bvs.bvs.get 0 |>.msb }

@[simp]
def Formula.language (φ : Formula) : Set (BitVecs φ.arity) :=
  match φ with
  | .width wp n => { bvs | wp.sat bvs.w n }
  | .atom rel t1 t2 =>
    let l1 := t1.language.lift (liftMaxSucc1 (FinEnum.card $ Fin t1.arity) (FinEnum.card $ Fin t2.arity))
    let l2 := t2.language.lift (liftMaxSucc2 (FinEnum.card $ Fin t1.arity) (FinEnum.card $ Fin t2.arity))
    let lrel := rel.language.lift $ liftLast2 (max (FinEnum.card (Fin t1.arity)) (FinEnum.card (Fin t2.arity)))
    let l := lrel ∩ l1 ∩ l2
    l.proj (liftExcept2 _)
  | .unop .neg φ => φ.languageᶜ
  | .binop op φ1 φ2 =>
    let l1 := φ1.language.lift $ liftMax1 φ1.arity φ2.arity
    let l2 := φ2.language.lift $ liftMax2 φ1.arity φ2.arity
    langBinop op l1 l2
  | .msbSet t =>
    let lmsb := langMsb.lift $ fun _ => Fin.last t.arity
    let l' := t.language ∩ lmsb
    l'.proj fun n => n.castLE (by admit /- proof elided -/
    )

theorem formula_language (φ : Formula) :
    φ.language = { (bvs : BitVecs φ.arity) | φ.sat (fun k => bvs.bvs.get k) } :=

:= by
  let n : Nat := φ.arity
  induction φ
  case width wp n =>
    simp
  case atom rel t1 t2 =>
    apply formula_language_case_atom
  case unop op φ ih =>
    rcases op; simp [ih, Set.compl_def]
  case binop op φ1 φ2 ih1 ih2 =>
    unfold Formula.language
    ext1 bvs
    simp [ih1, ih2]
    have heq1 : (φ1.sat fun k => bvs.bvs.get (liftMax1 φ1.arity φ2.arity k)) ↔
           (φ1.sat fun n => bvs.bvs.get (Fin.castLE (by simp) n)) := by
      congr!
    have heq2 : (φ2.sat fun k => bvs.bvs.get (liftMax2 φ1.arity φ2.arity k)) = true ↔
           (φ2.sat fun n => bvs.bvs.get (Fin.castLE (by simp) n)) = true := by
      congr!
    rcases op <;>
      simp [evalBinop, langBinop, Set.instMembership] <;> simp_all <;> tauto
  case msbSet t =>
    ext1 bvs; simp only [Formula.arity, Formula.language, Set.proj, Set.lift, langMsb, Fin.isValue,
      Set.preimage_setOf_eq, Set.mem_image, Set.mem_inter_iff,
      Set.mem_setOf_eq, Formula.sat]
    rcases bvs with ⟨w, bvs⟩
    constructor
    · rintro ⟨bvsb, ⟨ht, hmsb⟩, heq⟩
      simp only [Fin.isValue, Formula.arity] at ht hmsb ⊢
      unfold Term.language at ht
      simp only [BitVecs.transport, List.Vector.transport] at hmsb
      simp at ht; rw [←ht] at hmsb; rw [←hmsb]
      simp [BitVecs.transport] at heq
      obtain ⟨hw, hbvs⟩ := heq
      simp; congr 1; simp [hw]
      rcases hw; simp
      congr 1; ext1 k
      simp at hbvs; simp [←hbvs, List.Vector.transport]; congr
    · intros heq
      use ⟨w,
        bvs ++ ((t.evalFinBV fun k => bvs.get $ k.castLE (by simp)) ::ᵥ List.Vector.nil)⟩
      unfold Term.language
      simp [BitVecs.transport, List.Vector.transport] at heq ⊢
      constructor; assumption
      ext1 k; simp; congr 1


import Blase.Blase.AutoStructs.ForLean
import Blase.AutoStructs.BundledNfa
import Mathlib.Algebra.Group.Nat.Range
import Blase.FinEnum

@[aesop 50% unsafe]
theorem Std.HashMap.mem_of_getElem? [BEq K] [LawfulBEq K] [Hashable K] [LawfulHashable K] {m : Std.HashMap K V} {k : K} :
    m[k]? = some v → k ∈ m

theorem Std.HashMap.mem_iff_getElem? [BEq K] [LawfulBEq K] [Hashable K] [LawfulHashable K] [Inhabited V] {m : Std.HashMap K V} {k : K} :
     k ∈ m ↔ ∃ v, m[k]? = some v

State

abbrev State := Nat

RawCNFA.WF

@[grind =, simp]
lemma RawCNFA.addInitial_tr {m : RawCNFA A} : (m.addInitial s'').tr s a = m.tr s a

@[grind =, simp, aesop 50% unsafe]
lemma states_addInitial (m : RawCNFA A) (s' : State) :
    (m.addInitial s').states = m.states

@[grind =] -- TODO: should I?
lemma addTrans_tr (m : RawCNFA A) [LawfulBEq A] {a b : A} {s₁ s₁' s₂ : State} :
    (m.addTrans a s₁ s₂).tr s₁' b = if s₁ = s₁' ∧ a = b then (m.tr s₁ a).insert s₂ else m.tr s₁' b

@[grind =, simp]
lemma mem_addTrans_tr (m : RawCNFA A) [LawfulBEq A] (a : A) (s1 s2 : State) :
    s' ∈ (m.addTrans a s1 s2).tr s b ↔
      (s = s1 ∧ s' = s2 ∧ b = a) ∨ s' ∈ m.tr s b

@[grind =, simp]
lemma newState_tr {m : RawCNFA A} : m.newState.2.tr s a = m.tr s a

@[grind =, simp, aesop 50% unsafe]
lemma states_newState (m : RawCNFA A) :
    m.newState.2.states = m.states ∪ { m.stateMax }

@[grind ., simp, aesop 50% unsafe]
lemma RawCNFA.WF.trans_src_lt' {m : RawCNFA A} (hwf : m.WF) :
    ∀ {s a}, (s, a) ∈ m.trans → s ∈ m.states

@[grind ., simp, aesop 50% unsafe]
lemma RawCNFA.WF.trans_src_lt'' [LawfulBEq A] {m : RawCNFA A} (hwf : m.WF) :
    ∀ {s a s'}, s' ∈ m.tr s a → s ∈ m.states

import Mathlib.Algebra.Group.Nat.Range

import Blase.AutoStructs.BundledNfa

import Blase.FinEnum

open SetRel

abbrev State := Nat

structure RawCNFA (A : Type 0) [BEq A] [Hashable A] [DecidableEq A] [FinEnum A] where
  stateMax : State
  initials : Std.HashSet State
  finals : Std.HashSet State
  trans : Std.HashMap (State × A) (Std.HashSet State)
deriving Repr

section sim

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

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

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

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

end sim

section basics

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

def RawCNFA.newState (m : RawCNFA A) : State × RawCNFA A :=
  let old := m.stateMax
  let m := { m with stateMax := old + 1 }
  (old, m)

def RawCNFA.addTrans (m : RawCNFA A) (a : A) (s s' : State) : RawCNFA A :=
  let ns := m.trans.getD (s, a) ∅
  let ns := ns.insert s'
  { m with trans :=  m.trans.insert (s, a) ns }

def RawCNFA.addInitial (m : RawCNFA A) (s : State) : RawCNFA A :=
  { m with initials := m.initials.insert s }

def RawCNFA.createSink (m : RawCNFA A) : State × RawCNFA A :=
  let (s, m) := m.newState
  let m := m.addInitial s
  let m := FinEnum.toList (α := A).foldl (init := m) fun m a =>
    m.addTrans a s s
  (s, m)

structure RawCNFA.WF (m : RawCNFA A) where
  initials_lt : ∀ {s}, s ∈ m.initials → s ∈ m.states
  finals_lt : ∀ {s}, s ∈ m.finals → s ∈ m.states
  trans_src_lt : ∀ s_a ∈ m.trans, s_a.1 ∈ m.states
  trans_tgt_lt : s' ∈ m.tr s a → s' ∈ m.states

@[grind =, simp]
lemma createSink_trans [LawfulBEq A] {m : RawCNFA A} (hwf : m.WF) :
    s₂ ∈ m.createSink.2.tr s₁ a ↔
      (s₁ = m.stateMax ∧ s₂ = m.stateMax) ∨
      (s₁ ∈ m.states ∧ s₂ ∈ m.states ∧ s₂ ∈ m.tr s₁ a) :=

:= by
  unfold RawCNFA.createSink
  simp
  let motive (mᵢ : RawCNFA A) (hwf : mᵢ.WF) (as : List A) :=
    mᵢ.states = m.states ∪ {m.stateMax} →
    ∀ a,
      s₂ ∈ (as.foldl (λ (m' : RawCNFA A) a => m'.addTrans a m.stateMax m.stateMax)
              (init := mᵢ)).tr s₁ a ↔
        (s₁ = m.stateMax ∧ s₂ = m.stateMax ∧ a ∈ as) ∨
        (s₂ ∈ mᵢ.tr s₁ a)
  suffices h : motive (m.newState.2.addInitial m.stateMax) (by simp_all) (FinEnum.toList A) by
    simp only [FinEnum.mem_toList, and_true, motive] at h
    rw [h] <;> simp only [RawCNFA.addInitial_tr, newState_tr,
                          states_addInitial, states_newState]
    constructor; on_goal 2 => tauto
    rintro (h | h); (exact .inl h); right
    use (RawCNFA.WF.trans_src_lt'' hwf h),
        (RawCNFA.WF.trans_tgt_lt hwf h),
        h
  generalize_proofs hwf'; revert hwf'
  generalize (m.newState.2.addInitial m.stateMax) = mi
  induction FinEnum.toList A generalizing mi
  case nil =>
    rintro hwf'; simp [motive]
  case cons a as ih =>
    simp [motive]
    rintro hwf' hstates b
    simp [motive] at ih
    rw [ih]
    · simp only [mem_addTrans_tr]; tauto
    · simp [hwf', hstates]
    · simp [hstates]


Vars.inputN

def Vars.inputN (inp : ι) (k : Nat) {n : Nat} (hkn : k < n := by admit /- proof elided -/
) : Vars σ ι n :=
  .inputs (Inputs.mk ⟨k, by admit /- proof elided -/
  ⟩ inp)

Vars

inductive Vars (σ : Type) (ι : Type) (n : Nat)
| state (s : Inputs σ (n + 1))
| inputs (is : Inputs ι n)
| outputs (os : Fin n)
deriving DecidableEq, Hashable

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 }

FSM

structure FSM (arity : Type) : Type 1 where

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

  ( initCarry : α → Bool )

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

Vars.stateN

def Vars.stateN (s : σ) (i : Nat) {n : Nat} (hin : i ≤ n := by admit /- proof elided -/
) : Vars σ ι n :=
  .state (Inputs.mk ⟨i, by admit /- proof elided -/
  ⟩ s)

Inputs

structure Inputs (ι : Type) (n : Nat) : Type  where
  ix : Fin n
  input : ι
deriving DecidableEq, Hashable

def Vars.state0 (s : σ) {n : Nat} : Vars σ ι n :=
  .state (Inputs.mk ⟨0, by admit /- proof elided -/
  ⟩ s)

bigOr

def bigOr {α : Type _}
    (cs : List (Circuit α)) : Circuit α :=
  match cs with
  | [] => Circuit.fals
  | c :: cs =>
    c ||| (Circuit.bigOr cs)

bigAnd

def bigAnd {α : Type _}
    (cs : List (Circuit α)) : Circuit α :=
  match cs with
  | [] => Circuit.tru
  | c :: cs =>
    c &&& (Circuit.bigAnd cs)

eval

@[simp]
def eval : Circuit α → (α → Bool) → Bool
  | tru, _ => true
  | fals, _ => false
  | var b x, f => if b then f x else !(f x)
  | and c₁ c₂, f => (eval c₁ f) && (eval c₂ f)
  | or c₁ c₂, f => (eval c₁ f) || (eval c₂ f)
  | xor c₁ c₂, f => Bool.xor (eval c₁ f) (eval c₂ f)

map

def map : ∀ (_c : Circuit α) (_f : α → β), Circuit β
  | tru, _ => tru
  | fals, _ => fals
  | var b x, f => var b (f x)
  | and c₁ c₂, f => (map c₁ f) &&& (map c₂ f)
  | or c₁ c₂, f => (map c₁ f) ||| (map c₂ f)
  | xor c₁ c₂, f => (map c₁ f) ^^^ (map c₂ f)

Vars.castLe

def Vars.castLe {n m : Nat} (v : Vars σ ι n) (hnm : n ≤ m) : Vars σ ι m :=
  match v with
  | .state ss => .state (ss.castLe (by admit /- proof elided -/
  ))
  | .inputs is => .inputs (is.castLe hnm)
  | .outputs os =>
    .outputs (os.castLE (by admit /- proof elided -/
    ))

castLe

def castLe (i : Inputs ι n) (hn : n ≤ m) : Inputs ι m where
  ix := ⟨i.ix, by admit /- proof elided -/
  ⟩
  input := i.input

lemma eval_map {c : Circuit α} {f : α → β} {g : β → Bool} :
    eval (map c f) g = eval c (λ x => g (f x))

@[simp]
theorem eval_bigOr_eq_false_iff
    (cs : List (Circuit α)) (env : α → Bool):
    (Circuit.bigOr cs).eval env = false ↔
    (∀ (c : Circuit α), c ∈ cs → c.eval env = false)

@[simp]
theorem Vars.castLe_inputs_eq_inputs {n i m : Nat} (hi : i < n) (hnm : n ≤ m) :
  (Vars.inputN inp i hi : Vars σ ι n).castLe hnm =
  Vars.inputN inp i (by omega)

@[simp]
theorem Vars.castLe_stateN_eq_stateN  {n i m : Nat} (hi : i ≤ n) (hnm : n ≤ m) :
  (Vars.stateN s i : Vars σ ι n).castLe hnm =
  Vars.stateN s i (hin := by omega)

@[simp]
theorem Vars.castLe_castLe_eq_castLe_self {α : Type _} {p q r : Nat}
  (v : Vars α σ p) (h : p ≤ q) (h' : q ≤ r) :
    (v.castLe h).castLe h' = v.castLe (by omega)

ReflectVerif.BvDecide.mkStateNeqCircuit

@[simp]
theorem mkCarryAssignCircuitNAux_eval_eq {arity : Type _}
    [DecidableEq arity]
    [Fintype arity]
    [Hashable arity]
    (p : FSM arity) (s : p.α) (n : Nat)
    {env : Vars p.α arity (n + 1) → Bool} :
    ((mkCarryAssignCircuitNAux p s n).eval env) = ((p.nextStateCirc s).eval
      (fun x => match x with | .inl x => env (Vars.stateN x n) | .inr x => env (Vars.inputN x n)))

theorem mkStateUniqueCircuitN_eq_false_iff {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (n : Nat)
  {env : Vars p.α arity n → Bool} :
  ((mkStateUniqueCircuitN p n).eval env = false) ↔
  (∀ (i : Nat) (hi : i < n), ∃ (s : p.α), env (Vars.stateN s i) ≠ env (Vars.stateN s n))

import Mathlib.Data.Bool.Basic

import Mathlib.Data.Fin.Basic

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Finset.Defs

import Mathlib.Data.Multiset.FinsetOps

import Blase.Fast.BitStream

import Blase.Fast.Defs

import Blase.Fast.FiniteStateMachine

import Blase.Fast.Decide

import Blase.SingleWidth.Syntax

import Lean.Meta.ForEachExpr

import Lean.Meta.Tactic.Simp.BuiltinSimprocs.BitVec

import Blase.Fast.ForLean

import Blase.Vars

import Blase.EnvBitstream

import Lean

open Fin.NatCast

namespace ReflectVerif

open Lean Meta Elab Tactic

namespace BvDecide

open Std Sat AIG

def mkCarryAssignCircuitNAux {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (s : p.α) (n : Nat) : Circuit (Vars p.α arity (n + 1)) :=
    (p.nextStateCirc s).map fun v =>
      match v with
        | .inl t => Vars.stateN t n
        | .inr i => Vars.inputN i n

def mkCarryAssignCircuitN {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) :
  Circuit (Vars p.α arity (n + 1)) :=
  let carrys := FinEnum.toList p.α |>.map fun s =>

    Circuit.xor
      (mkCarryAssignCircuitNAux p s n)
      (Circuit.var true <| Vars.stateN s (n + 1))
  Circuit.bigOr carrys

def mkOutputAssignCircuitNAux {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) : Circuit (Vars p.α arity (n + 1)) :=
    (p.outputCirc).map fun v =>
      match v with
        | .inl s' => Vars.stateN s' n
        | .inr i => Vars.inputN i n

def mkOutputAssignCircuitN {arity : Type _}
  [DecidableEq arity]
  [Fintype arity]
  [Hashable arity]
  (p : FSM arity) (n : Nat) :
  Circuit (Vars p.α arity (n + 1)) :=
    Circuit.xor
      (mkOutputAssignCircuitNAux p n)
      (Circuit.var true <| Vars.outputs ⟨n, by admit /- proof elided -/
      ⟩)

def mkStateNeqCircuit
  {arity : Type _} {i : Nat}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (s t : p.α → Circuit (Vars p.α arity i)) : Circuit (Vars p.α arity i) :=
  Circuit.bigAnd <| FinEnum.toList p.α |>.map fun a => ~~~ (s a) ^^^ (t a)

def mkStateUniqueCircuitN {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity]
  (p : FSM arity) (n : Nat) : Circuit (Vars p.α arity n) :=
  let sn : p.α → Circuit (Vars p.α arity n) := fun s =>
    Circuit.var true (Vars.stateN s n)
  let circs := (List.range n).attach |>.map fun ⟨i, hi⟩ =>
    let si : p.α → Circuit (Vars p.α arity n) := fun s =>
      Circuit.var true (Vars.stateN s i (by admit /- proof elided -/
      ))
    (mkStateNeqCircuit p si sn)
  Circuit.bigOr circs

structure KInductionCircuits {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity] (fsm : FSM arity) (n : Nat) where

  cInitCarryAssignCirc : Circuit (Vars fsm.α arity 0)

  cSuccCarryAssignCirc : Circuit (Vars fsm.α arity (n+2))

  cOutAssignCirc : Circuit (Vars fsm.α arity (n + 2))

  cStatesUniqueCirc : Circuit (Vars fsm.α arity n)

structure KInductionCircuits.IsLawful {arity : Type _}
    [DecidableEq arity] [Fintype arity] [Hashable arity] {fsm : FSM arity} {n : Nat}
  (circs : KInductionCircuits fsm n) where
  hCInitCarryAssignCirc :
    ∀ {env : Vars fsm.α arity 0 → Bool},
      (circs.cInitCarryAssignCirc.eval env = false)
      ↔ (∀ (s : fsm.α), fsm.initCarry s = env (Vars.state0 s))

  hCSuccCarryAssignCirc :
    ∀ {env : Vars fsm.α arity (n + 2) → Bool},
      (circs.cSuccCarryAssignCirc.eval env = false)
      ↔ (∀ (s : fsm.α) (i : Nat) (hi : i < n + 2),
        env (Vars.stateN s (i + 1)) =
          ((mkCarryAssignCircuitNAux fsm s i).map
            (fun v => v.castLe (by admit /- proof elided -/
            ))).eval env)
  hCOutAssignCirc :
    ∀ {env : Vars fsm.α arity (n + 2) → Bool},
      (circs.cOutAssignCirc.eval env = false)
      ↔ (∀ (i : Nat) (hi : i < n + 2),
        (fsm.outputCirc).eval
          (fun x => match x with
            | .inl s => env (Vars.stateN s i)
            | .inr j => env (Vars.inputN j i)) =
        env (Vars.outputs ⟨i, by admit /- proof elided -/
        ⟩))
  hCStatesUniqueCirc :
    ∀ {env : Vars fsm.α arity (n) → Bool},
      (circs.cStatesUniqueCirc.eval env = false)
      ↔ (∀ (i j : Nat) (hij : i < j ∧ j ≤ n),
        ∃ (s : fsm.α), env (Vars.stateN s i) ≠ env (Vars.stateN s j))

namespace KInductionCircuits

variable {arity : Type _}
  {fsm : FSM arity}

def castCircLe {n m : Nat} (c : Circuit (Vars fsm.α arity n)) (hnm : n ≤ m := by admit /- proof elided -/
) :
    Circuit (Vars fsm.α arity m) :=
  c.map (fun v => v.castLe hnm)

variable [DecidableEq arity] [Fintype arity] [Hashable arity]

def mkSucc
    (prev : KInductionCircuits fsm n) :
    KInductionCircuits fsm (n + 1) :=
  let cInitCarryAssignCirc := prev.cInitCarryAssignCirc
  { cInitCarryAssignCirc := cInitCarryAssignCirc
  , cSuccCarryAssignCirc :=
      (mkCarryAssignCircuitN fsm (n + 2)) |||
      (castCircLe prev.cSuccCarryAssignCirc)
  , cOutAssignCirc :=
      (mkOutputAssignCircuitN fsm (n + 2)) |||
      (castCircLe prev.cOutAssignCirc)
  , cStatesUniqueCirc :=
      mkStateUniqueCircuitN fsm (n + 1) |||
      (castCircLe prev.cStatesUniqueCirc)
  }

theorem IsLawful_mkSucc_of_IsLawful {arity : Type _}
    [DecidableEq arity] [Fintype arity] [Hashable arity]
    {fsm : FSM arity} {n : Nat}
    (prev : KInductionCircuits fsm n)
    (hPrev : prev.IsLawful) :
    (mkSucc prev).IsLawful :=

where
  hCInitCarryAssignCirc := by
    simp only [mkSucc, castCircLe]
    exact hPrev.hCInitCarryAssignCirc
  hCSuccCarryAssignCirc := by
    simp only [mkSucc, castCircLe, Circuit.eval_map]
    simp [Circuit.eval_map]
    intros env
    constructor
    · intros h s i hi
      obtain ⟨h₁, h₂⟩ := h
      rw [hPrev.hCSuccCarryAssignCirc] at h₂
      by_cases hi : i < n + 2
      · simp only [Vars.castLe_stateN_eq_stateN, Circuit.eval_map,
        Vars.castLe_castLe_eq_castLe_self, mkCarryAssignCircuitNAux_eval_eq,
        Vars.castLe_inputs_eq_inputs] at h₂
        rw [h₂ s i hi]
      · have hi : i = n + 2 := by omega
        subst hi
        apply h₁
    · intros h
      constructor
      · intros s
        simp only at h ⊢
        rw [h s _ (by omega)]
      · rw [hPrev.hCSuccCarryAssignCirc]
        intros s i hi
        simp only at h
        simp only [Vars.castLe_stateN_eq_stateN, Circuit.eval_map,
          Vars.castLe_castLe_eq_castLe_self, mkCarryAssignCircuitNAux_eval_eq,
          Vars.castLe_inputs_eq_inputs]
        rw [h s i (by omega)]

  hCOutAssignCirc := by
    simp only [mkSucc, castCircLe]
    simp [Circuit.eval_map]
    intros env
    constructor
    · intros h i hi
      obtain ⟨h₁, h₂⟩ := h
      rw [hPrev.hCOutAssignCirc] at h₂
      by_cases hi : i < n + 2
      · simp at h₁ h₂ ⊢
        rw [h₂ i hi]
      · have : i = n + 2 := by omega
        subst this
        apply h₁
    · intros h
      constructor
      · simp only at h ⊢
        rw [h]
      · rw [hPrev.hCOutAssignCirc]
        intros i hi
        simp
        apply h
  hCStatesUniqueCirc := by
    simp only [mkSucc, castCircLe]
    simp [Circuit.eval_map]
    intros env
    constructor
    · intros h i j hij hjn
      simp [mkStateUniqueCircuitN_eq_false_iff] at h
      simp [hPrev.hCStatesUniqueCirc] at h
      obtain ⟨h₁, h₂⟩ := h
      by_cases hj : j ≤ n
      · grind
      · have : j = n + 1 := by
          omega
        subst this
        grind
    · intros h
      constructor
      · rw [mkStateUniqueCircuitN_eq_false_iff]
        intros i hi
        apply h (j := n + 1)
        omega
        omega
      · simp [hPrev.hCStatesUniqueCirc]
        intros i j hij jn
        simp at h ⊢
        apply h i j
        omega
        omega


RelationOrdering

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

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

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

instFinEnumBV

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

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

NFA'.autSignedCmpSA

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

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

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

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

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

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

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

section fsm

variable {arity : Type} [FinEnum arity]

open BitStream in

end fsm

section nfas_relations

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

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

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

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

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

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

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

:= by
  let getState {w} (bv1 bv2 : BitVec w) : NFA.signedCmpState :=
    if bv2.ult bv1 then .gt else if bv1.ult bv2 then .lt else .eq
  constructor <;> simp [NFA.autSignedCmp, autSignedCmp, autSignedCmpSA, RelationOrdering.srel]
  · cases cmp <;> simp [BitVec.sle_iff_slt_or_eq]; tauto
  · rintro (_ | _ | _) <;> simp
  · rintro (_ | _ | _) a w bv1 bv2 <;> simp [NFA.stepSet]
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro ⟨_, _⟩; use .eq; simp; fin_cases a <;> simp [instFinEnumBV] at * <;> tauto
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .gt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .gt; simp_all
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .lt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .lt; simp_all
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .lt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .lt; simp_all
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [NFA.signedCmpStep, instFinEnumBV]
      · rintro _; fin_cases a <;> simp [NFA.signedCmpStep, instFinEnumBV] at *
        · use .gt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .gt; simp_all


CNFA

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

RawCNFA.statesFinset

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

State

abbrev State := Nat

Std.HashSet.toSet

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

axiom hashMap_missing : ∀ {P : Prop}, P

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

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

@[simp]
lemma Std.HashSet.mem_toSet [BEq α] [Hashable α] (m : HashSet α) : x ∈ m.toSet ↔ x ∈ m

@[simp]
theorem Std.HashSet.mem_attachWith_mem [BEq α] [Hashable α] [LawfulBEq α] (m : HashSet α) {P H} (x : α) h :
    ⟨x, h⟩ ∈ m.attachWith P H ↔ x ∈ m

product

include hinj in
omit [BEq α] [Hashable α] [LawfulBEq α] in
lemma product.prodArray_spec_helper
    (is : Array γ) (hnd : is.toList.Nodup)
    (s : S₁) (hnew : ∀ s₂, f s s₂ ∉ is):
  let motive (a : Array γ) (S : Set S₂)  :=
    a.toList.Nodup ∧
    (∃ r, a.toList = is.toList ++ r ∧ (∀ z ∈ r, ∃ s₁ s₂, z = f s₁ s₂)) ∧
    ∀ s1 s2, f s1 s2 ∈ a ↔ s1 ≠ s ∧ f s1 s2 ∈ is ∨ s1 = s ∧ s2 ∈ S
  let body := m₂.attachWith _ hm₂ |>.fold (init := is) fun is s2 =>
      is.push (f s s2)
  motive body (m₂.attachWith _ hm₂).toSet

include hinj in
lemma product.prodArray'_spec_full {aᵢ : Array γ} (hnd: aᵢ.toList.Nodup) (hnin : ∀ s₁ s₂, f s₁ s₂ ∉ aᵢ) :
    (product.prodArray' f hm₁ hm₂ aᵢ).toList.Nodup ∧
    (∃ r, (product.prodArray' f hm₁ hm₂ aᵢ).toList = aᵢ.toList ++ r ∧ (∀ z ∈ r, ∃ s₁ s₂, z = f s₁ s₂)) ∧
    ∀ s₁ s₂, f s₁ s₂ ∈ product.prodArray' f hm₁ hm₂ aᵢ ↔ (s₁.val ∈ m₁ ∧ s₂.val ∈ m₂)

import Mathlib.Data.Fintype.Prod

import Blase.AutoStructs.Worklist

import Mathlib.Tactic.ApplyFun

open SetRel

section sink

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

end sink

section generic_prod

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

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

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

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

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

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

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

end generic_prod

section product

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

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

lemma product.f_spec {m₁ m₂ : CNFA n} {s₁ : m₁.m.states} {s₂ : m₂.m.states} :
    ∀ a s₁' s₂',
      (a, (s₁', s₂')) ∈ f m₁ m₂ (s₁, s₂) ↔ s₁'.val ∈ m₁.m.tr s₁ a ∧ s₂'.val ∈ m₂.m.tr s₂ a :=

:= by
  suffices heq :
      ∀ as a (hnd : as.toList.Nodup) (hnd' : (FinEnum.toList (α := BitVec n)).Nodup)
       (hnew : ∀ b s₁ s₂, (b, s₁, s₂) ∈ as → b ∉ (FinEnum.toList (α := BitVec n))) s₁' s₂',
        (a, (s₁', s₂')) ∈
          ((FinEnum.toList (α := BitVec n)).foldl (init := as) fun as a =>
          product.prodArray' (λ s₁' s₂' ↦ (a, (s₁', s₂')))
            (fun _ => m₁.wf.trans_tgt_lt (s := s₁) (a := a)) (fun _ => m₂.wf.trans_tgt_lt (s := s₂) (a := a)) as) ↔
        (a, (s₁', s₂')) ∈ as ∨ a ∈ (FinEnum.toList (α := BitVec n)) ∧ s₁'.val ∈ m₁.m.tr s₁ a ∧ s₂'.val ∈ m₂.m.tr s₂ a by
    rintro a s₁' s₂'; rw [f, heq]
    · simp; rintro h; apply Array.not_mem_empty at h; trivial
    · exact List.dedup_eq_self.mp rfl
    · exact FinEnum.nodup_toList
    · rintro _ _ _ h; apply Array.not_mem_empty at h; trivial
  induction (FinEnum.toList (α := BitVec n))
  case nil =>
    simp only [List.foldl_nil, List.not_mem_nil, false_and, Subtype.forall]
    intros _ _ _ _ _ _; tauto
  case cons a as ih =>
    rintro bs b hnd hnd' hnew s₁' s₂'
    obtain ⟨hnd'', ⟨r, hr, hf⟩, hin⟩ := prodArray'_spec_full (fun s₁' s₂' => (a, s₁', s₂'))
      (by rintro _ _ _ _ ⟨rfl, rfl, rfl⟩; tauto)
      (fun _ => m₁.wf.trans_tgt_lt (s := s₁)) (fun _ => m₂.wf.trans_tgt_lt (s := s₂) (a := a))
      hnd
      (by dsimp; rintro s₁ s₂ hin; apply hnew at hin; simp at hin)
    dsimp
    have hmem : ∀ b s₁' s₂', b ≠ a → (b, s₁', s₂') ∈ prodArray' (fun s₁' s₂' => (a, s₁', s₂'))
        (fun _ => m₁.wf.trans_tgt_lt (s := s₁) (a := a))
        (fun _ => m₂.wf.trans_tgt_lt (s := s₂) (a := a)) bs →
          (b, s₁', s₂') ∈ bs := by
      rintro b s₁' s₂' hneq hin
      rw [Array.mem_def] at hin ⊢; rw [hr] at hin; simp only [List.mem_append] at hin
      rcases hin with hin | hin; assumption
      obtain ⟨_, _, ⟨rfl, -, -⟩⟩ := hf _ hin; simp at hneq
    rw [ih] <;> clear ih; rotate_left
    · apply hnd''
    · exact (List.nodup_cons.mp hnd').2
    · rintro b s₁' s₂' hin; by_cases heq : b = a
      · subst heq; apply List.Nodup.notMem hnd'
      · specialize hnew b s₁' s₂' (hmem b s₁' s₂' heq hin); simp at hnew; simp [hnew]
    by_cases heq :  b = a
    · subst heq
      simp [hin]; clear hin; constructor
      · rintro (⟨h1, h2⟩ | ⟨_, h1, h2⟩) <;> right <;> simp [h1, h2]
      · rintro (hc | ⟨h1, h2⟩)
        · apply hnew at hc; simp at hc
        · left; simp [h1, h2]
    · simp only [List.mem_cons, heq, false_or]; constructor
      · rintro (hin | hin)
        · left; apply hmem _ _ _ heq hin
        · right; exact hin
      · rintro (hin | hin)
        · left; rw [Array.mem_def] at hin ⊢; rw [hr]; simp [hin]
        · right; exact hin


syntax "llvm.and"     : MLIR.Pretty.uniform_op

syntax "llvm.and"     : MLIR.Pretty.uniform_op

syntax "llvm.ashr"    : MLIR.Pretty.exact_op

syntax "llvm.add"     : MLIR.Pretty.overflow_op

syntax "llvm.return"  : MLIR.Pretty.uniform_op

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

macro_rules
  | `(mlir_op| $resName:mlir_op_operand = $name:InstCombine.cmp_op_name $x, $y $[: $t]?) => do
    let some opName := extractOpName name.raw
      | Macro.throwUnsupported
    let t ← t.getDM `(mlir_type| _)
    `(mlir_op| $resName:mlir_op_operand = $opName ($x, $y) : ($t, $t) -> (i1) )

macro_rules
  | `(mlir_op| $resName:mlir_op_operand = $name:InstCombine.int_cast_op $x : $t to $t') => do
    let some opName := extractOpName name.raw
      | Macro.throwUnsupported
    `(mlir_op| $resName:mlir_op_operand = $opName ($x) : ($t) -> $t')

macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant( $x $[: $inner_type]?)
      $[: $outer_type]? ) => do

      let outer_type ← outer_type.getDM `(mlir_type| _)
      let inner_type := inner_type.getD outer_type
      `(mlir_op| $res:mlir_op_operand = "llvm.mlir.constant"()
          {value = $x:neg_num : $inner_type} : () -> ($outer_type) )
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant( ${ $x:term }) $[: $t]?) => do
      let t ← t.getDM `(mlir_type| _)
      let x ← `(MLIR.AST.AttrValue.int $x [mlir_type| $t])
      `(mlir_op| $res:mlir_op_operand = "llvm.mlir.constant"() {value = $$($x) } : () -> ($t) )

macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (true) $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (1 : i1) : i1)
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (false) $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant (0 : i1) : i1)
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant $x $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant($x $[: $t]?) $[: $t]?)
  | `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant ${ $x:term } $[: $t]?) =>
      `(mlir_op| $res:mlir_op_operand = llvm.mlir.constant($$($x) $[: $t]?) $[: $t]?)

macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.icmp $p $x, $y $[: $t]?) => do
    let t ← t.getDM `(mlir_type| _)
    match p.getString with
      | "eq" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.eq" ($x, $y) : ($t, $t) -> (i1))
      | "ne" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ne" ($x, $y) : ($t, $t) -> (i1))
      | "slt" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.slt" ($x, $y) : ($t, $t) -> (i1))
      | "sle" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.sle" ($x, $y) : ($t, $t) -> (i1))
      | "sgt" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.sgt" ($x, $y) : ($t, $t) -> (i1))
      | "sge" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.sge" ($x, $y) : ($t, $t) -> (i1))
      | "ult" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ult" ($x, $y) : ($t, $t) -> (i1))
      | "ule" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ule" ($x, $y) : ($t, $t) -> (i1))
      | "ugt" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.ugt" ($x, $y) : ($t, $t) -> (i1))
      | "uge" => `(mlir_op| $res:mlir_op_operand = "llvm.icmp.uge" ($x, $y) : ($t, $t) -> (i1))
      | _ => Macro.throwErrorAt p s!"unexpected predicate {p.getString}"

macro_rules
  | `(mlir_op| $res:mlir_op_operand = llvm.select $c, $x, $y $[: $t]?) => do
    let t ← t.getDM `(mlir_type| _)
    `(mlir_op| $res:mlir_op_operand = "llvm.select" ($c, $x, $y) : (i1, $t, $t) -> ($t))

Lets

inductive Lets (Γ_in : Ctxt d.Ty) (eff : EffectKind) :
    (Γ_out : Ctxt d.Ty) → Type where
  | nil : Lets Γ_in eff Γ_in
  | var (body : Lets Γ_in eff Γ_out) (e : Expr d Γ_out eff t) : Lets Γ_in eff e.outContext

Expr.decidableEq

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

Var

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

Ctxt

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

  deriving Repr, Lean.ToExpr, DecidableEq

Valuation.nil

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

EffectKind

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

Com

inductive Com : Ctxt d.Ty → EffectKind → List d.Ty → Type where
  | rets {Γ} {tys} {eff : EffectKind} (vs : HVector Γ.Var tys) : Com Γ eff tys
  | var (e : Expr Γ eff ty) (body : Com (ty ++ Γ) eff β) : Com Γ eff β

FlatCom

structure FlatCom (d : Dialect) [DialectSignature d]  (Γ_in : Ctxt d.Ty) (eff : EffectKind)
    (Γ_out : Ctxt d.Ty) (ts : List d.Ty) where
  lets : Lets d Γ_in eff Γ_out
  rets : HVector Γ_out.Var ts

HVector.decidableEqReg

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

Var

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

Hom.with

def Hom.with [DecidableEq Ty] {Γ₁ Γ₂ : Ctxt Ty} (f : Γ₁.Hom Γ₂) {ts}
    (v₁ : HVector Γ₁.Var ts) (v₂ : HVector Γ₂.Var ts) : Γ₁.Hom Γ₂ :=
  fun _ w =>
    match v₁.idxOf? w with
    | none => f w
    | some ⟨i, h⟩ => (v₂.get i).cast h

Hom

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

TyDenote

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

Valuation.instAppendHVector

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

Com.outContextHom

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

Com.outContext

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

Com.rec'

def Com.rec' {Γ} (com : Com d Γ eff t) : motive com :=

  Com.rec
    (motive_1 := fun _ _ _ _ => PUnit)
    (motive_2 := fun _ eff' t' c =>
      (h₁ : eff = eff') → (h₂ : t = t') → motive (h₁ ▸ h₂ ▸ c))
    (motive_3 := fun _ _ => PUnit)
    (fun _ _ _ _ _ _ => ⟨⟩)
    (fun v h₁ h₂ =>
      cast (by admit /- proof elided -/
      ) <| rets (h₂ ▸ v))
    (fun e' body' _ r' h₁ h₂ =>
      let e := h₁ ▸ e'
      let body : Com _ _ eff t := cast (by admit /- proof elided -/
      ) body'
      let r : motive body := cast (by admit /- proof elided -/
      ) (r' h₁ h₂)
      cast (by admit /- proof elided -/
      ) <| var e body r)
    ⟨⟩
    (fun _ _ _ _ => ⟨⟩)
    com
    rfl
    rfl

Valuation.cast

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

Com.outContextDiff

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

Expr.outContext

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

Expr.bvars

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

returnTypes

def returnTypes  := Signature.returnTypes ∘ s.signature

Signature

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

DialectSignature

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

Com.bvars

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

Diff

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

Diff.Valid

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

toHom

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

castCtxt

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

Com.denoteLets

def Com.denoteLets : (com : Com d Γ eff ty) → (Γv : Valuation Γ) →
    eff.toMonad d.m (com.outContext.Valuation)
  | .rets _, V => pure V
  | .var e body, V =>
      e.denote V >>= body.denoteLets >>= fun V =>
        return V.cast (by admit /- proof elided -/
        )

DialectDenote

class DialectDenote (d : Dialect) [TyDenote d.Ty] [DialectSignature d] where
  denote : (op : d.Op) → HVector toType (DialectSignature.sig op) →
    (HVector (fun t : Ctxt d.Ty × List d.Ty =>
      t.1.Valuation
      → EffectKind.impure.toMonad d.m (HVector toType t.2))
            (DialectSignature.regSig op)) →
    ((DialectSignature.effectKind op).toMonad d.m
      (HVector toType <| DialectSignature.returnTypes op))

Lets.denote

def Lets.denote [DialectSignature d] [DialectDenote d] {Γ₂}
    (lets : Lets d Γ₁ eff Γ₂) (V : Valuation Γ₁) : (eff.toMonad d.m <| Valuation Γ₂) :=
  match lets with
  | .nil          => return V
  | .var lets' e  => lets'.denote V >>= e.denote

sig

def sig          := Signature.sig ∘ s.signature

regSig

def regSig       := Signature.regSig ∘ s.signature

RegionSignature

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

Ty

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

effectKind

def effectKind   := Signature.effectKind ∘ s.signature

Dialect

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

Op

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

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

HVector.denote

def HVector.denote :
    {l : RegionSignature d.Ty} → (T : HVector (fun t => Com d t.1 .impure t.2) l) →
    HVector (fun t => t.1.Valuation → EffectKind.impure.toMonad d.m (HVector toType t.2)) l
  | _, .nil => HVector.nil
  | _, .cons v vs => HVector.cons (v.denote) (HVector.denote vs)

FlatCom.denote

@[simp] abbrev FlatCom.denote [DialectDenote d]
    (flatCom : FlatCom d Γ eff Γ_out ts)
    (V : Γ.Valuation) : eff.toMonad d.m (HVector toType ts) :=
  flatCom.lets.denote V >>= (return flatCom.rets.map ·)

RegionSignature.map

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

Signature.map

def Signature.map (f : Ty → Ty') : Signature Ty → Signature Ty' :=
  fun sig => {
    sig         := sig.sig.map f
    regSig      := sig.regSig.map f
    returnTypes := sig.returnTypes.map f
  }

map

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

Expr.denote

def Expr.denote {ty} (e : Expr d Γ eff ty) (V : Valuation Γ) :
    eff.toMonad d.m (e.outContext.Valuation) :=
  match e with
  | ⟨op, ty_eq, heff, args, regArgs⟩ => do
      let argsDenote := args.map V
      let val ← EffectKind.liftEffect heff <| DialectDenote.denote op argsDenote regArgs.denote
      return (val ++ V).cast (by admit /- proof elided -/
      )

Expr.op

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

liftEffect

def liftEffect [Pure m] {e1 e2 : EffectKind} {α : Type}
    (hle : e1 ≤ e2) (v1 : e1.toMonad m α) : e2.toMonad m α :=
  match e1, e2, hle with
    | .pure, .pure, _ | .impure, .impure, _ => v1
    | .pure, .impure, _ => Pure.pure v1

toMonad

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

Com.denote

def Com.denote : Com d Γ eff ty → (Γv : Valuation Γ) →
    eff.toMonad d.m (HVector toType ty)
  | .rets vs, Γv     => pure (vs.map Γv)
  | .var e body, V => e.denote V >>= body.denote

Com.ty

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

Valuation.cons

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

cast

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

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

Com.castPureToEff

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

Com.changeEffect

def Com.changeEffect {eff₁ eff₂ : EffectKind} (h : eff₁ ≤ eff₂) :
    Com d Γ eff₁ t → Com d Γ eff₂ t := fun com =>
  Com.rec' (motive := @fun Γ _ => eff₁ ≤ eff₂ → Com d Γ eff₂ t)
            (fun v _h               => rets v)
      (fun e _body castBody h => var (e.changeEffect h) (castBody h))
    com h

Expr.changeEffect

def Expr.changeEffect {eff₁ eff₂ : EffectKind} (h : eff₁ ≤ eff₂) :
    Expr d Γ eff₁ t → Expr d Γ eff₂ t
  | Expr.mk op ty_eq eff_le args regArgs =>
    have heff : DialectSignature.effectKind op ≤ eff₂ := by admit /- proof elided -/

Expr.args

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

Expr.regArgs

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

Regions

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

com

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

def bb0 : Region 0 := [mlir_region|
{
  ^bb0(%arg0: i32):
    %0 = llvm.mlir.constant(8) : i32
    %1 = llvm.mlir.constant(31) : i32
    %2 = llvm.ashr %arg0, %1 : i32
    %3 = llvm.and %2, %0 : i32
    %4 = llvm.add %3, %2 : i32
    llvm.return %4 : i32
  }]

Region

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

MetaLLVM

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

Ty

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

Op

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

MOp

inductive MOp (φ : Nat) : Type
  | unary   (w : Width φ) (op : MOp.UnaryOp φ) :  MOp φ
  | binary  (w : Width φ) (op : MOp.BinaryOp) :  MOp φ
  | select  (w : Width φ) : MOp φ
  | icmp    (c : IntPred) (w : Width φ) : MOp φ

  | const (w : Width φ) (val : ℤ) : MOp φ
deriving Repr, DecidableEq, Inhabited, Lean.ToExpr

binary

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

MOp.BinaryOp

inductive MOp.BinaryOp : Type
  | and
  | or   (disjoint : DisjointFlag := {disjoint := false} )
  | xor
  | shl  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | lshr (exact : ExactFlag := {exact := false} )
  | ashr (exact : ExactFlag := {exact := false} )
  | urem
  | srem
  | add  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | mul  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | sub  (nswnuw : NoWrapFlags := {nsw := false, nuw := false} )
  | sdiv (exact : ExactFlag := {exact := false} )
  | udiv (exact : ExactFlag := {exact := false} )
deriving DecidableEq, Inhabited, Lean.ToExpr

LLVM

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

MTy

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

Width

abbrev Width φ := ConcreteOrMVar Nat φ

ConcreteOrMVar

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

const

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

MOp.UnaryOp

inductive MOp.UnaryOp (φ : Nat) : Type
  | neg
  | not
  | copy
  | freeze
  | trunc (w' : Width φ) (noWrapFlags : NoWrapFlags := {nsw := false, nuw := false} )
  | zext  (w' : Width φ) (nneg : NonNegFlag := {nneg := false} )
  | sext  (w' : Width φ)
deriving Repr, DecidableEq, Inhabited, Lean.ToExpr

select

@[simp_llvm_option]
def select {w : Nat} (c? : IntW 1) (x? y? : IntW w ) : IntW w := do
  let c ← c?
  if c = 1#1 then x? else y?

IntW

def IntW w := PoisonOr <| BitVec w

PoisonOr

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

icmp

@[simp_llvm_option]
def icmp {w : Nat} (c : IntPred) (x y : IntW w) : IntW 1 := do
  let x' ← x
  let y' ← y
  icmp? c x' y'

icmp?

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

IntPred

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

icmp'

@[simp_llvm]
def icmp' {w : Nat} (c : IntPred) (x y : BitVec w) : Bool :=
  match c with
    | .eq => (x == y)
    | .ne => (x != y)
    | .sgt => (x >ₛ y)
    | .sge => (x ≥ₛ y)
    | .slt => (x <ₛ y)
    | .sle => (x ≤ₛ y)
    | .ugt => (x >ᵤ y)
    | .uge => (x ≥ᵤ y)
    | .ult => (x <ᵤ y)
    | .ule => (x ≤ᵤ y)

mkCom

def mkCom [TransformTy d φ] [TransformExpr d φ] [TransformReturn d φ]
  (reg : MLIR.AST.Region φ) :
  ExceptM d (Σ (Γ : Ctxt d.Ty) (eff : EffectKind) (ty : _), Com d Γ eff ty) :=
  match reg.ops with
  | [] => throw <| .generic "Ill-formed region (empty)"
  | coms => BuilderM.runWithEmptyMapping <| do
    let Γ ← declareBindings ∅ reg.args
    let com ← mkComHelper Γ coms
    return ⟨Γ, com⟩

FlatCom.denoteLets

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

Com.toLets

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

Lets.castPureToEff

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

Expr.castPureToEff

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

Expr.returnVars

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

ofFin

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

Com.returnVars

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

Valuation.comap

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

map

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

HVectorLiteral

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

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

TyDenote.toType

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

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

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

@[simp] theorem Lets.denote_var {lets : Lets d Γ_in eff Γ_out} {e : Expr d Γ_out eff t} :
    (lets.var e).denote = fun V_in => lets.denote V_in >>= e.denote

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

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

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

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

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

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

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

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

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

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

theorem map_map {A B C : α → Type*} {l : List α} (t : HVector A l)
    (f : ∀ a, A a → B a) (g : ∀ a, B a → C a) :
    (t.map f).map g = t.map (fun a v => g a (f a v))

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

Zipper.insertPureCom

theorem denote_insertCom {zip : Zipper d Γ_in eff t₁} [LawfulMonad d.m]
    {newCom : Com d zip.Γ_mid eff newTys} {vs : HVector zip.Γ_mid.Var newTys} :
    (zip.insertCom vs newCom).denote = (fun (V_in : Valuation Γ_in) => do
      let V_mid ← zip.top.denote V_in
      let V_newMid ← newCom.denoteLets V_mid
      zip.bot.denote
        (V_newMid.comap <| newCom.outContextHom.with vs newCom.returnVars)
      )

import LeanMLIR.Framework.Basic

import LeanMLIR.Transforms.Rewrite.Match

open Ctxt (Valuation Var Hom)

variable (d : Dialect) [DialectSignature d]

structure Zipper (Γ_in : Ctxt d.Ty) (eff : EffectKind) (tys : List d.Ty) where

  {Γ_mid : Ctxt d.Ty}

  top : Lets d Γ_in eff Γ_mid

  bot : Com d Γ_mid eff tys

namespace Zipper

variable {d}

section Denote

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

def denote (zip : Zipper d Γ_in eff tys) (V_in : Valuation Γ_in) :
    eff.toMonad d.m (HVector toType tys) :=
  (zip.top.denote V_in) >>= zip.bot.denote

end Denote

section ToCom

variable {Γ_mid}

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

end ToCom

section InsertCom

variable [DecidableEq d.Ty]

def insertCom (zip : Zipper d Γ_in eff ty)
    (vs : HVector zip.Γ_mid.Var newTy) (newCom : Com d zip.Γ_mid eff newTy) :
    Zipper d Γ_in eff ty :=
  let top := zip.top.addComToEnd newCom

  let bot := zip.bot.changeVars <| newCom.outContextHom.with vs newCom.returnVars

  { top, bot }

def insertPureCom (zip : Zipper d Γ_in eff ty)
    (vs : HVector zip.Γ_mid.Var newTy)
    (newCom : Com d zip.Γ_mid .pure newTy) : Zipper d Γ_in eff ty :=
  zip.insertCom vs (newCom.castPureToEff eff)

section Lemmas

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

theorem denote_insertPureCom {zip : Zipper d Γ_in eff t₁} [LawfulMonad d.m]
    {newCom : Com d zip.Γ_mid .pure newTys} {vs : HVector zip.Γ_mid.Var newTys} :
    (zip.insertPureCom vs newCom).denote = (fun (V_in : Valuation Γ_in) => do
      let V_mid ← zip.top.denote V_in
      zip.bot.denote
        ((Com.denoteLets newCom V_mid).comap <| newCom.outContextHom.with vs newCom.returnVars)
      ) :=

:= by
  have (V_mid) (h : Com.outContext (Com.castPureToEff eff newCom) = Com.outContext newCom) :
      ((Com.denoteLets newCom V_mid).comap fun x v => v.castCtxt h).comap
        (newCom.castPureToEff eff).outContextHom
      = (Com.denoteLets newCom V_mid).comap newCom.outContextHom := by
    funext t' ⟨v', hv'⟩
    simp only [Com.outContextHom, Com.outContextDiff, Com.bvars_castPureToEff]
    rfl
  funext V
  simp only [insertPureCom, denote_insertCom, Com.denoteLets_castPureToEff,
    Com.returnVars_castPureToEff, Valuation.comap_with, HVector.map_map, pure_bind, this,
    Valuation.comap_outContextHom_denoteLets, Valuation.comap_apply, Var.castCtxt_castCtxt,
    Var.castCtxt_rfl, Com.denoteLets_returnVars]


f

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

instFinEnumBV

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

langMsb

@[simp]
def langMsb : Set (BitVecs 1) := { bvs | bvs.bvs.get 0 |>.msb }

BitVecs

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

head

def head (x : BitStream) : Bool      := x 0

BitStream

def BitStream : Type := Nat → Bool

enc

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

BitVecs'

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

BitVec.ofFn

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

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'

NFA'

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

accepts'

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

ext

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

@[simp]
lemma dec_enc' : dec (enc bvs) = bvs

@[simp]
lemma dec_enc : Function.RightInverse (α := BitVecs' n) enc dec

lemma dec_enc_w (bvs : BitVecs n) : (dec (enc bvs)).w = bvs.w

@[simp]
theorem BitVec.ofFn_getElem {w : Nat} (f : Fin w → Bool) {i : Nat} (hi : i < w) :
    (BitVec.ofFn f)[i] = f ⟨i, hi⟩

@[simp]
theorem BitVec.ofFn_getLsbD {w : Nat} {f : Fin w → Bool} {i : Nat} (hi : i < w) :
    (BitVec.ofFn f).getLsbD i = f ⟨i, hi⟩

theorem BitVec.ofFn_getLsbD_fin {w : Nat} {f : Fin w → Bool} {i : Fin w} :
    (BitVec.ofFn f).getLsbD i = f i

NFA.msbCorrect

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

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

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

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

section fsm

variable {arity : Type} [FinEnum arity]

open BitStream in

end fsm

section nfas_relations

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

def NFA.msbStep (q : NFA.msbState) (a : BitVec 1) : List NFA.msbState :=
  match q, a with
  | .i, 0 => [.i]
  | .i, 1 => [.i, .f]
  | _, _ => []

def NFA.autMsbSet : NFA (BitVec 1) msbState where
  step s a := { s' | s' ∈ msbStep s a }
  start := {.i}
  accept := {.f}

def NFA'.autMsbSet : NFA' 1 := ⟨_, NFA.autMsbSet⟩

def NFA.msbLang : Language (BitVec 1) := { bvs  | bvs.getLast? = some 1 }

def NFA.msbSA (q : msbState) : Language (BitVec 1) :=
  match q with
  | .i => ⊤
  | .f => msbLang

def NFA.msbCorrect : NFA.autMsbSet.correct msbSA msbLang :=

@[simp]
lemma autMsbSet_accepts : NFA'.autMsbSet.accepts = langMsb :=

:= by
  simp [NFA'.accepts, NFA'.accepts', NFA'.autMsbSet]
  rw [NFA.correct_spec NFA.msbCorrect, NFA.msbLang]
  ext bvs; simp only [BitVec.ofNat_eq_ofNat, Set.mem_image, Set.mem_setOf_eq]
  constructor
  · rintro ⟨bvs', hl, heq⟩
    have _ : bvs'.length ≠ 0 := by cases bvs'; tauto; simp
    rw [←heq]
    simp [dec]
    rw [BitVec.msb_eq_getLsbD_last]
    rw [BitVec.ofFn_getLsbD (by omega)]
    simp
    rw [List.getLast?_eq_getElem?] at hl
    rw [List.getElem?_eq_getElem (by omega)] at hl
    injection hl
    simp_all only [BitVec.getElem_one, decide_true]
  · intros h; use enc bvs
    simp only [dec_enc', and_true]
    simp [enc]
    have hw : bvs.w ≠ 0 := by
      rcases bvs with ⟨w, bvs⟩; rintro rfl
      simp_all [BitVec.eq_nil (bvs.head)]
    use ⟨bvs.w - 1, by omega⟩
    simp; rw [List.getLast?_eq_getElem?]
    simp; constructor
    · rw [List.getElem?_eq_getElem (by simp; omega)]; simp
    · ext i hi; rw [BitVec.ofFn_getElem _ (by omega)]
      rw [BitVec.msb_eq_getLsbD_last] at h
      simp [←BitVec.getLsbD_eq_getElem]
      obtain rfl : i = 0 := by omega
      simp_all


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

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 }

BitStream

def BitStream : Type := Nat → Bool

Term.eval

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

Predicate.eval

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

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

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

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

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

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

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

eval

@[simp]
def eval : Circuit α → (α → Bool) → Bool
  | tru, _ => true
  | fals, _ => false
  | var b x, f => if b then f x else !(f x)
  | and c₁ c₂, f => (eval c₁ f) && (eval c₂ f)
  | or c₁ c₂, f => (eval c₁ f) || (eval c₂ f)
  | xor c₁ c₂, f => Bool.xor (eval c₁ f) (eval c₂ f)

sumVarsRight

def sumVarsRight [DecidableEq α] [DecidableEq β] : Circuit (α ⊕ β) → List β
  | tru => []
  | fals => []
  | var _ (Sum.inl _) => []
  | var _ (Sum.inr x) => [x]
  | and c₁ c₂ => (sumVarsRight c₁ ++ sumVarsRight c₂).dedup
  | or c₁ c₂ => (sumVarsRight c₁ ++ sumVarsRight c₂).dedup
  | xor c₁ c₂ => (sumVarsRight c₁ ++ sumVarsRight c₂).dedup

fst

def fst {α β : Type _} [DecidableEq α] [DecidableEq β]
    (c : Circuit (α ⊕ β)) : Circuit α :=
  Circuit.bOr (c.sumVarsRight.pi (λ _ => [true, false]))
  (λ x => Circuit.assignVars c
    (λ i => Sum.rec (λ i _ => Sum.inl i) (λ i hi => Sum.inr (x i (by admit /- proof elided -/
    ))) i))

bOr

def bOr : ∀ (_s : List α) (_f : α → Circuit β), Circuit β
| [], _ => fals
| a::l, f => l.foldl (λ c x => c ||| (f x)) (f a)

assignVars

def assignVars [DecidableEq α] :
    ∀ (c : Circuit α) (_f : ∀ (a : α) (_ha : a ∈ c.vars), β ⊕ Bool), Circuit β
  | tru, _ => tru
  | fals, _ => fals
  | var b x, f =>
    Sum.elim
      (var b)
      (λ c : Bool => if Bool.xor b c then fals else tru)
      (f x (by admit /- proof elided -/
      ))
  | and c₁ c₂, f => (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) &&&
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))
  | or c₁ c₂, f =>  (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) |||
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))
  | xor c₁ c₂, f => (assignVars c₁ (λ x hx => f x (by admit /- proof elided -/
  ))) ^^^
                    (assignVars c₂ (λ x hx => f x (by admit /- proof elided -/
                    )))

vars

def vars [DecidableEq α] : Circuit α → List α
  | tru => []
  | fals => []
  | var _ x => [x]
  | and c₁ c₂ => (vars c₁ ++ vars c₂).dedup
  | or c₁ c₂ => (vars c₁ ++ vars c₂).dedup
  | xor c₁ c₂ => (vars c₁ ++ vars c₂).dedup

evalv

@[simp] def evalv [DecidableEq α] : ∀ (c : Circuit α), (∀ a ∈ vars c, Bool) → Bool
  | tru, _ => true
  | fals, _ => false
  | var b x, f => if b then f x (by admit /- proof elided -/
  ) else !(f x (by admit /- proof elided -/
  ))
  | and c₁ c₂, f => (evalv c₁ (fun i hi => f i (by admit /- proof elided -/
  ))) &&
    (evalv c₂ (fun i hi => f i (by admit /- proof elided -/
    )))
  | or c₁ c₂, f => (evalv c₁ (fun i hi => f i (by admit /- proof elided -/
  ))) ||
    (evalv c₂ (fun i hi => f i (by admit /- proof elided -/
    )))
  | xor c₁ c₂, f => Bool.xor (evalv c₁ (fun i hi => f i (by admit /- proof elided -/
  )))
    (evalv c₂ (fun i hi => f i (by admit /- proof elided -/
    )))

f

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

simplifyOr

def simplifyOr : Circuit α → Circuit α → Circuit α
  | tru, _ => tru
  | _, tru => tru
  | fals, c => c
  | c, fals => c
  | c₁, c₂ => or c₁ c₂

bind

def bind : ∀ (_c : Circuit α) (_f : α → Circuit β), Circuit β
  | tru, _ => tru
  | fals, _ => fals
  | var b x, f => if b then f x else ~~~ (f x)
  | and c₁ c₂, f => (bind c₁ f) &&& (bind c₂ f)
  | or c₁ c₂, f => (bind c₁ f) ||| (bind c₂ f)
  | xor c₁ c₂, f => (bind c₁ f) ^^^ (bind c₂ f)

lemma le_def : ∀ (c₁ c₂ : Circuit α), c₁ ≤ c₂ ↔ ∀ f, eval c₁ f → eval c₂ f

theorem eval_fst {α β : Type _} [DecidableEq α] [DecidableEq β]
    (c : Circuit (α ⊕ β)) (g : α → Bool) :
    c.fst.eval g ↔ ∃ g' : β → Bool, c.eval (Sum.elim g g')

lemma eval_assignVars [DecidableEq α] : ∀ {c : Circuit α}
    {f : ∀ (a : α) (_ha : a ∈ c.vars), β ⊕ Bool} {g : β → Bool},
    eval (assignVars c f) g = evalv c (λ a ha => Sum.elim g id (f a ha))

@[simp] lemma eval_bOr :
  ∀ {s : List α} {f : α → Circuit β} {g : β → Bool},
    eval (bOr s f) g = ∃ a ∈ s, eval (f a) g

@[simp] lemma eval_foldl_or :
  ∀ (s : List α) (f : α → Circuit β) (c : Circuit β) (g : β → Bool),
    (eval (s.foldl (λ c x => c ||| (f x)) c) g : Prop) ↔
      eval c g ∨ (∃ a ∈ s, eval (f a) g)

@[simp] lemma eval_or : ∀ (c₁ c₂ : Circuit α) (f : α → Bool),
    (eval (c₁ ||| c₂) f) = ((eval c₁ f) || (eval c₂ f))

lemma eval_eq_evalv [DecidableEq α] : ∀ (c : Circuit α) (f : α → Bool),
    eval c f = evalv c (λ x _ => f x)

lemma eval_bind : ∀ (c : Circuit α) (f : α → Circuit β) (g : β → Bool),
    eval (bind c f) g = eval c (λ a => eval (f a) g)

FSM.State

abbrev State : Type := p.α → Bool

decideIfZerosAux

theorem carry_changeInitCarry_succ
    (p : FSM arity) (c : p.α → Bool) (x : arity → BitStream) : ∀ n,
    (p.changeInitCarry c).carry x (n+1) =
      (p.changeInitCarry (p.nextBit c (fun a => x a 0)).1).carry
        (fun a i => x a (i+1)) n
  | 0 => by simp [carry, changeInitCarry, nextBit]
  | n+1 => by
    rw [carry, carry_changeInitCarry_succ p _ _ n]
    simp [nextBit, carry, changeInitCarry]

theorem eval_changeInitCarry_succ
    (p : FSM arity) (c : p.α → Bool) (x : arity → BitStream) (n : ℕ) :
    (p.changeInitCarry c).eval x (n+1) =
      (p.changeInitCarry (p.nextBit c (fun a => x a 0)).1).eval
        (fun a i => x a (i+1)) n

theorem evalAux_eq_zero_of_set {arity : Type _} (p : FSM arity)
    (R : Set (p.α → Bool)) (hR : ∀ x s, (p.nextBit s x).1 ∈ R → s ∈ R)
    (hi : p.initCarry ∉ R) (hr1 : ∀ x s, (p.nextBit s x).2 = true → s ∈ R)
    (x : arity → BitStream) (n : ℕ) : p.eval x n = false ∧ p.carry x n ∉ R

theorem eval_eq_zero_of_set {arity : Type _} (p : FSM arity)
    (R : Set (p.α → Bool)) (hR : ∀ x s, (p.nextBit s x).1 ∈ R → s ∈ R)
    (hi : p.initCarry ∉ R) (hr1 : ∀ x s, (p.nextBit s x).2 = true → s ∈ R) :
    p.eval = fun _ _ => false

theorem decideIfZeroAux_wf {α : Type _} [Fintype α] [DecidableEq α]
    {c c' : Circuit α} (h : ¬c' ≤ c) : card_compl (c' ||| c) < card_compl c

import Mathlib.Data.FinEnum

import Mathlib.Data.Fintype.Card

import Mathlib.Data.Fintype.Sum

import Mathlib.Data.Fintype.Sigma

import Mathlib.Data.Fintype.Pi

import Mathlib.Data.Fintype.BigOperators

import Mathlib.Tactic.Zify

import Mathlib.Tactic.Ring

import Blase.FinEnum

import Blase.Fast.Defs

import Blase.Fast.Circuit

import Blase.Vars

open Sum

section FSM

variable {α β α' β' : Type} {γ : β → Type}

structure FSM (arity : Type) : Type 1 where

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

  ( initCarry : α → Bool )

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

open Lean in

open Lean in

open Lean in

open Lean in

namespace FSM

variable {arity : Type} (p : FSM arity)

abbrev State : Type := p.α → Bool

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)

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

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

def changeInitCarry (p : FSM arity) (c : p.α → Bool) : FSM arity :=
  { p with initCarry := c }

section EvalInduction

end EvalInduction

end FSM

namespace FSM

end FSM

open Term

def card_compl [Fintype α] [DecidableEq α] (c : Circuit α) : ℕ :=
  Finset.card $ (@Finset.univ (α → Bool) _).filter (fun a => c.eval a = false)

def decideIfZerosAux {arity : Type _} [DecidableEq arity]
    (p : FSM arity) (c : Circuit p.α) : Bool :=

  if c.eval p.initCarry
  then false
  else

    have c' := (c.bind (p.nextStateCirc)).fst
    if h : c' ≤ c then true
    else
      have _wf : card_compl (c' ||| c) < card_compl c :=
        decideIfZeroAux_wf h
      decideIfZerosAux p (c' ||| c)
  termination_by card_compl c

theorem decideIfZerosAux_correct {arity : Type _} [DecidableEq arity]
    (p : FSM arity) (c : Circuit p.α)
    (hc : ∀ s, c.eval s = true →
      ∃ m y, (p.changeInitCarry s).eval y m = true)
    (hc₂ : ∀ (x : arity → Bool) (s : p.α → Bool),
      (FSM.nextBit p s x).snd = true → Circuit.eval c s = true) :
    decideIfZerosAux p c = true ↔ ∀ n x, p.eval x n = false :=

:= by
  rw [decideIfZerosAux]
  split_ifs with h
  · simp
    exact hc p.initCarry h
  · dsimp
    split_ifs with h'
    · simp only [true_iff]
      intro n x
      rw [p.eval_eq_zero_of_set {x | c.eval x = true}]
      · intro y s
        simp [Circuit.le_def, Circuit.eval_fst, Circuit.eval_bind] at h'
        simp [FSM.nextBit]
        apply h'
      · assumption
      · exact hc₂
    · let c' := (c.bind (p.nextStateCirc)).fst
      have _wf : card_compl (c' ||| c) < card_compl c :=
        decideIfZeroAux_wf h'
      apply decideIfZerosAux_correct p (c' ||| c)
      simp [c', Circuit.eval_fst, Circuit.eval_bind]
      intro s hs
      rcases hs with ⟨x, hx⟩ | h
      · rcases hc _ hx with ⟨m, y, hmy⟩
        use (m+1)
        use fun a i => Nat.casesOn i x (fun i a => y a i) a
        rw [FSM.eval_changeInitCarry_succ]
        rw [← hmy]
        simp only [FSM.nextBit, Nat.rec_zero]
      · exact hc _ h
      · intro x s h
        have := hc₂ _ _ h
        simp only [Bool.or_eq_true, Circuit.eval_or, this, or_true]
termination_by card_compl c