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 }

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

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

ext

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

CNFA.ofFSM

@[simp]
theorem BitVec.cast_inj (h : w = w') {x y : BitVec w} : BitVec.cast h x = BitVec.cast h y ↔ x = y

@[simp]
theorem BitVec.append_inj {x1 x2 : BitVec w} {y1 y2 : BitVec w'} :
    x1 ++ y1 = x2 ++ y2 ↔ x1 = x2 ∧ y1 = y2

@[simp]
lemma BitVec.cons_inj : cons b1 bv1 = cons b2 bv2 ↔ (b1 = b2) ∧ bv1 = bv2

@[simp]
lemma bitVecToFinFun_rinv (c : carry → Bool) [FinEnum carry]:
    bitVecToFinFun (finFunToBitVec c) = c

@[simp]
lemma bitVecToFinFun_linv [FinEnum ar] (bv : BitVec $ FinEnum.card ar) :
    finFunToBitVec (bitVecToFinFun bv) = bv

@[simp]
lemma bitVecToFinFun_inj [FinEnum ar] : Function.Injective (bitVecToFinFun (ar := ar))

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

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 }

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

@[simp]
lemma CNFA.ofFSM.f_spec {p : FSM arity} {s s' : BitVec (FinEnum.card p.α)} :
    (a, s') ∈ f p s ↔ bitVecToFinFun s' ∈ (NFA.ofFSM p).step (bitVecToFinFun s) a :=

:= by
  let motive (as : List (BitVec (FinEnum.card arity))) := ∀ (acc : Array _) a s',
    ((a, s') ∈ as.foldl (init := acc) (process p s))
      ↔ (a, s') ∈ acc ∨ (a.setWidth (FinEnum.card arity) ∈ as) ∧
          bitVecToFinFun s' ∈ (NFA.ofFSM p).step (bitVecToFinFun s) a
  suffices h : motive (FinEnum.toList (BitVec (FinEnum.card arity))) by
    specialize h #[]
    simp at h
    rw [←h]
    rfl
  generalize FinEnum.toList (BitVec (FinEnum.card arity)) = qs
  induction qs
  case nil => simp [motive]
  case cons a as ih =>
    rintro acc b s'
    simp only [List.foldl_cons]; rw [ih]
    simp [process]
    constructor
    · rintro ((hacc | ⟨rfl, rfl⟩) | ⟨hin₁, hin₂⟩)
      · exact .inl hacc
      · right; simp [NFA.ofFSM]; constructor <;> rfl
      · right; simp_all only [or_true, and_self, motive]
    · rintro (hacc | ⟨(rfl | hold), hst⟩)
      · tauto
      · simp [NFA.ofFSM] at hst; left; right
        rcases hst with ⟨hs', hb⟩
        constructor
        · rw [←BitVec.cons_msb_setWidth b]
          simp_all only [BitVec.setWidth_cons, BitVec.cons_inj, and_true]; rfl
        · apply_fun bitVecToFinFun <;> simp only [hs', bitVecToFinFun_rinv, bitVecToFinFun_inj]
          rfl
      · tauto


syntax "max" : MLIR.Pretty.uniform_op

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

Predicate.evalUlt

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

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

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

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

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

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

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

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

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

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

Predicate.arity

@[simp] def Predicate.arity : Predicate → Nat
| .width _ _ => 0
| .binary .eq t1 t2 => max t1.arity t2.arity
| .binary .neq t₁ t₂ => max t₁.arity t₂.arity
| .binary .ult t₁ t₂ => max t₁.arity t₂.arity
| .binary .ule t₁ t₂ => t₁.arity ⊔ t₂.arity ⊔ (t₁.arity ⊔ t₂.arity)
| .binary .slt t₁ t₂ => (t₁.arity ⊔ t₂.arity ⊔ (t₁.arity ⊔ t₂.arity))
| .binary .sle t₁ t₂ => (t₁.arity ⊔ t₂.arity ⊔ (t₁.arity ⊔ t₂.arity) ⊔ (t₁.arity ⊔ t₂.arity))
| .lor p q => max p.arity q.arity
| .land p q => max p.arity q.arity

BinaryPredicate

inductive BinaryPredicate
| eq
| neq
| ult
| ule
| slt
| sle
deriving Repr, Lean.ToExpr

WidthPredicate

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

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

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

theorem and_comm (x y : BitStream) : x &&& y = y &&& x

lemma Term.evalFin_eq_eval (t : Term)
   (varsList : List BitStream) (varsFin : Fin t.arity → BitStream)
   (hvars : ∀ (i : Fin t.arity), varsList.getD i default = (varsFin i)) :
    Term.evalFin t varsFin = Term.eval t varsList

import Mathlib.Data.Fintype.Card

import Mathlib.Data.Fintype.Sum

import Mathlib.Data.Fintype.Sigma

import Mathlib.Data.Fintype.BigOperators

import Mathlib.Tactic.Zify

import Mathlib.Tactic.Ring

import Blase.Fast.Defs

import Blase.Fast.BitStream

open Term

lemma Predicate.evalFin_eq_eval (p : Predicate)
   (varsList : List BitStream) (varsFin : Fin p.arity → BitStream)
   (hvars : ∀ (i : Fin p.arity), varsList.getD i default = (varsFin i)) :
    Predicate.evalFin p varsFin  = Predicate.eval p varsList :=

:= by
  induction p generalizing varsList <;>
    dsimp -failIfUnchanged [Predicate.evalFin, Predicate.eval, Predicate.arity] at *
  case width rel n =>
    rcases rel <;> dsimp -failIfUnchanged [Predicate.evalFin, Predicate.eval, Predicate.arity] at *
  case binary ap t₁ t₂ =>
    rcases ap <;>
    · dsimp [Predicate.evalFin, Predicate.eval, Predicate.arity] at *
      simp [evalEq, evalNeq, evalUlt, evalSlt, evalLor]
      rw [Term.evalFin_eq_eval _ varsList]
      · rw [Term.evalFin_eq_eval _ varsList]
        try rw [BitStream.and_comm]
        · intros i
          rw [hvars ⟨i, by omega⟩]
          rfl
      · intros i
        rw [hvars ⟨i, by omega⟩]
        rfl
  case land p q hp hq =>
    simp [evalLand]
    rw [hp varsList]
    · rw [hq varsList]
      · intros i
        rw [hvars ⟨i, by omega⟩]
        rfl
    · intros i
      rw [hvars ⟨i, by omega⟩]
      rfl
  case lor p q hp hq =>
    simp [evalLor]
    rw [hp varsList]
    · rw [hq varsList]
      · intros i
        rw [hvars ⟨i, by omega⟩]
        rfl
    · intros i
      rw [hvars ⟨i, by omega⟩]
      rfl


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

@[simp]
theorem ofBitVecZext_eq_getLsbD (x : BitVec w) (i : Nat) :
  ofBitVecZext x i = x.getLsbD i

@[simp] theorem addAux_zero (x y : BitStream) : (x.addAux y 0) =
  ((x 0) ^^ (y 0), (x 0) && (y 0))

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 ofBitVecZext {w} (x : BitVec w) : BitStream :=
  fun i => x.getLsbD i

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 ofBitVecZext_add_EqualUpTo :
    ofBitVecZext (x + y) ≈ʷ (ofBitVecZext x) + (ofBitVecZext y) :=

:= by
  intros n a
  have add_lemma : ⟨(x + y).getLsbD n, BitVec.carry (n + 1) x y false ⟩ = (ofBitVecZext x).addAux (ofBitVecZext y) n := by
    induction n
    case zero =>
      simp only [zero_add, addAux_zero, ofBitVecZext_eq_getLsbD, Prod.mk.injEq]
      simp only [BitVec.getLsbD, BitVec.toNat_add, Nat.testBit_zero, Nat.mod_two_pos_mod_two_eq_one,
        a, true_and]
      simp only [add_odd_iff_neq, ne_eq, eq_iff_iff, decide_not, Bool.decide_iff_dist,
        Bool.not_eq_eq_eq_not, BitVec.carry, pow_one, Bool.toNat_false, add_zero, ge_iff_le,
        two_le_add_iff_odd_and_odd, Bool.decide_and, and_true]
      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]


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

product

def product (final? : Prop → Prop → Prop) (M N : NFA' n) : NFA' n where
  σ := _
  M := M.M.product final? N.M

NFA'

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

nfa

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

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

include hinj in
lemma product.prodArray_spec_full :
    (product.prodArray f hm₁ hm₂).toList.Nodup ∧
    ∀ s₁ s₂, f s₁ s₂ ∈ product.prodArray f hm₁ hm₂ ↔ (s₁.val ∈ m₁ ∧ s₂.val ∈ m₂)

include hinj in
@[simp]
lemma product.prodArray_spec :
    ∀ s₁ s₂, f s₁ s₂ ∈ product.prodArray f hm₁ hm₂ ↔ (s₁.val ∈ m₁ ∧ s₂.val ∈ m₂)

@[simp]
lemma product.inits_spec :
    ∀ s₁ s₂, (s₁, s₂) ∈ inits m₁ m₂ ↔ (s₁.val ∈ m₁.m.initials ∧ s₂.val ∈ m₂.m.initials)

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

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)

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

end generic_prod

section product

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

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

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

noncomputable def to_prop (f : Bool → Bool → Bool) (p1 p2 : Prop) : Prop :=
  f (@Decidable.decide p1 (Classical.propDecidable _)) (@Decidable.decide p2 (Classical.propDecidable _))

lemma product.sim {m1 m2 : CNFA n}:
    m1.Sim M1 → m2.Sim M2 →
    (nfa (product.inits m1 m2) (final final? m1 m2) (f m1 m2)).Bisim (M1.M.product (to_prop final?) M2.M) :=

:= by
  rintro ⟨R₁, hsim₁⟩ ⟨R₂, hsim₂⟩
  let R : SetRel (m1.m.states × m2.m.states) (M1.σ × M2.σ) :=
    {((s₁, s₂), (q₁, q₂)) | s₁.val ~[R₁] q₁ ∧ s₂.val ~[R₂] q₂ }
  use R; constructor
  · rintro ⟨s₁, s₂⟩ ⟨q₁, q₂⟩ ⟨hR₁, hR₂⟩
    simp [nfa, to_prop, final]
    rw [←hsim₁.accept hR₁, ←hsim₂.accept hR₂]; congr
  · constructor
    · rintro ⟨s₁, s₂⟩ hstart
      simp [nfa, inits_spec] at hstart; rcases hstart with ⟨h₁, h₂⟩
      obtain ⟨q₁, hq₁, hR₁⟩ := hsim₁.initial₁ h₁
      obtain ⟨q₂, hq₂, hR₂⟩ := hsim₂.initial₁ h₂
      use (q₁, q₂); simp [NFA.product, R, *]
    · rintro ⟨q₁, q₂⟩ ⟨hst₁, hst₂⟩
      apply hsim₁.initial₂ at hst₁; obtain ⟨s₁, hi₁, hR₁⟩ := hst₁
      apply hsim₂.initial₂ at hst₂; obtain ⟨s₂, hi₂, hR₂⟩ := hst₂
      simp only [nfa, Set.mem_setOf_eq, Prod.exists, inits_spec]
      have hin₁ : s₁ ∈ m1.m.states := by apply m1.wf.initials_lt hi₁
      have hin₂ : s₂ ∈ m2.m.states := by apply m2.wf.initials_lt hi₂
      use ⟨s₁, hin₁⟩, ⟨s₂, hin₂⟩
      simp_all [R]
  · rintro ⟨s₁, s₂⟩ ⟨q₁, q₂⟩ a ⟨s₁', s₂'⟩ ⟨hR₁, hR₂⟩ hst
    simp [nfa, f_spec] at hst; rcases hst with ⟨hst₁, hst₂⟩
    obtain ⟨q₁', hst₁, hR₁'⟩ := hsim₁.trans_match₁ hR₁ hst₁
    obtain ⟨q₂', hst₂, hR₂'⟩ := hsim₂.trans_match₁ hR₂ hst₂
    simp [NFA.product]
    use q₁', q₂'
    simp_all [R]
  · rintro ⟨s₁, s₂⟩ ⟨q₁, q₂⟩ a ⟨q₁', q₂'⟩ ⟨hR₁, hR₂⟩ hst
    simp [NFA.product] at hst; rcases hst with ⟨hst₁, hst₂⟩
    obtain ⟨s₁', hst₁, hR₁'⟩ := hsim₁.trans_match₂ hR₁ hst₁ (by simp) (by simp)
    obtain ⟨s₂', hst₂, hR₂'⟩ := hsim₂.trans_match₂ hR₂ hst₂ (by simp) (by simp)
    have hin₁ : s₁' ∈ m1.m.states := by apply m1.wf.trans_tgt_lt hst₁;
    have hin₂ : s₂' ∈ m2.m.states := by apply m2.wf.trans_tgt_lt hst₂
    simp only [nfa, Set.mem_setOf_eq, Prod.exists, f_spec]
    use ⟨s₁', hin₁⟩, ⟨s₂', hin₂⟩
    simp_all [R]


syntax "neg" : MLIR.Pretty.uniform_op

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

Mapping

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

Ctxt

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

  deriving Repr, Lean.ToExpr, DecidableEq

Var

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

Lets

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

Expr.decidableEq

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

Var

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

Valuation.nil

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

EffectKind

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

appendCases

@[elab_as_elim]
def appendCases
    {motive : (Γ ++ Δ).Var t → Sort u}
    (left : (v : Var Γ t) → motive (appendInl v))
    (right : (v : Var Δ t) → motive (appendInr v)) :
    (v : (Γ ++ Δ).Var t) → motive v
  | ⟨idx, h⟩ =>
    if hv : idx < Γ.length then
      left ⟨idx, by admit /- proof elided -/
      ⟩
    else
      let v' : Var _ _ := ⟨idx - Γ.length, by admit /- proof elided -/
      ⟩
      have eq : v'.appendInr = ⟨idx, h⟩ := by admit /- proof elided -/

last

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

Valuation.cons

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

HVector.decidableEqReg

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

appendInl

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

sig

def sig          := Signature.sig ∘ s.signature

Signature

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

DialectSignature

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

effectKind

def effectKind   := Signature.effectKind ∘ s.signature

returnTypes

def returnTypes  := Signature.returnTypes ∘ s.signature

Expr.ty

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

Expr.op

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

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

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

theorem eq.ty_eq {v : Γ.Var t} {w : Γ.Var u} (h : v.eq w) : t = u

@[simp]
theorem Expr.op_mk {Γ : Ctxt d.Ty} {ty} {eff : EffectKind} (op : d.Op)
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) (DialectSignature.sig op))
    (regArgs) :
    (Expr.mk op ty_eq eff_le args regArgs).op = op

@[simp]
theorem Expr.regArgs_mk {Γ : Ctxt d.Ty} {ty eff op}
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) (DialectSignature.sig op)) (regArgs) :
    (Expr.mk op ty_eq eff_le args regArgs).regArgs = regArgs

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

@[simp]
theorem Expr.args_mk {Γ : Ctxt d.Ty} {ty eff op}
    (ty_eq : ty = DialectSignature.returnTypes op)
    (eff_le : DialectSignature.effectKind op ≤ eff)
    (args : HVector (Var Γ) (DialectSignature.sig op)) (regArgs) :
    (Expr.mk op ty_eq eff_le args regArgs).args = args

matchVar

@[simp]
theorem MatchVar.liftM_bind_eq_some_iff (x? : Option α)
    (f : α → MatchVarM Δ Γ β) :
    ((liftM x? >>= f) mapIn = some mapOut)
    ↔ ( ∃ h : x?.isSome,
        f (x?.get h) mapIn = some mapOut )

import LeanMLIR.Framework

import LeanMLIR.Transforms.Rewrite.Mapping

open Ctxt (Var VarSet Valuation Hom)

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

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

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

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

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

open MatchVarM

variable [DecidableEq d.Op]

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

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

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

end

section MatchVar

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

theorem matchVar_appendInl {w : Var ⟨te⟩ t} :
    matchVar lets v (.var matchLets matchExpr) w.appendInl ma = some ma' →
    ∃ args,
      lets.getPureExpr v
        = some ⟨_, w, matchExpr.op, matchExpr.ty_eq, matchExpr.eff_le, args, matchExpr.regArgs⟩
      ∧ matchArg lets matchLets args matchExpr.args ma = some ma' :=

:= by
  unfold matchVar
  simp only [Var.appendCases_appendInl]
  simp only [MatchVar.liftM_bind_eq_some_iff, forall_exists_index]
  generalize h_e? : lets.getPureExpr v = e?
  intro h_isSome h
  split_ifs at h
  case neg => contradiction
  case neg => contradiction
  case pos h_e? hw =>
    rcases matchExpr with ⟨mOp, _, _, mArgs, mRegArgs⟩
    rcases h_e? with ⟨(rfl : _ = mOp), (rfl : _ = mRegArgs)⟩
    rcases Option.isSome_iff_exists.mp h_isSome with ⟨⟨tys', w', e'⟩, rfl⟩
    simp only [Option.get_some, Expr.op_mk, Expr.regArgs_mk, Option.some.injEq, Sigma.mk.injEq,
      Expr.args_mk]
    obtain rfl : te = tys' := by
      simpa [e'.ty_eq]
    rcases e' with ⟨_op, _ty_eq, _eff_le, args, regArgs⟩
    obtain rfl : w = w' := by
      simp_all [Option.get_some]

    exact ⟨args, ⟨rfl, by rfl⟩, h⟩


lemma RawCNFA.reverse_spec {m : RawCNFA A} (hwf : m.WF) :
    let m' := m.reverse
    m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.finals ∧ m'.finals = m.initials ∧
      ∀ s a s', s' ∈ m'.tr s a ↔ s ∈ m.tr s' a

RawCNFA.statesFinset

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

State

abbrev State := Nat

RawCNFA.states

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

RawCNFA.tr

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

Std.HashSet.toSet

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

Std.HashMap.toPFun

def Std.HashMap.toPFun [BEq α] [Hashable α] (m : HashMap α β) (x : α) : Option β := m[x]?

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 }

theorem Std.HashMap.fold_induction [BEq α] [LawfulBEq α] [DecidableEq α] [Hashable α]
  {f : γ → α → β → γ} {m : HashMap α β} {motive : γ → (α → Option β) → Prop} :
    motive b (λ _ ↦ none) →
    (∀ b x y m, m x = none → motive b m → motive (f b x y) (Function.update m x y)) →
    motive (m.fold f b) m.toPFun

theorem Std.HashMap.toPFun_toList[BEq α] [LawfulBEq α] [Hashable α] (m : HashMap α β) :
    m.toPFun = λ k ↦ m.toList.find? (λ x ↦ x.1 == k) |>.map Prod.snd

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

RawCNFA.reverse

lemma RawCNFA.reverse_spec_procesState {m : RawCNFA A} (hwf : m.WF) s₀ a₀ ss' (hs₀ : s₀ ∈ m.states) :
    let motive m' ss' :=
      (∀ s ∈ ss', s ∈ m.states) →
        m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.initials ∧ m'.finals = m.finals ∧
          ∀ s a s', s' ∈ m'.tr s a ↔ (s' ∈ m.tr s a ∨ s' = s₀ ∧ a = a₀ ∧ s ∈ ss')
    motive (RawCNFA.reverse.processState m (s₀, a₀) ss') ss'.toSet

import Mathlib.Data.Fintype.Prod

import Blase.AutoStructs.Worklist

import Mathlib.Tactic.ApplyFun

open SetRel

section sink

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

end sink

section generic_prod

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

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

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

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

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

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

end generic_prod

section product

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

end product

section determinization

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

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

:= m.reverse
    m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.finals ∧ m'.finals = m.initials ∧
      ∀ s a s', s' ∈ m'.tr s a ↔ s ∈ m.tr s' a := by
  let motive (m' : RawCNFA A) (trs : (State × A) → Option (Std.HashSet State)) :=
    (∀ s a ss', trs (s, a) = some ss' → s ∈ m.states ∧ ∀ s' ∈ ss', s' ∈ m.states) →
    m'.WF ∧ m'.stateMax = m.stateMax ∧ m'.initials = m.finals ∧ m'.finals = m.initials ∧
      ∀ s a s', s' ∈ m'.tr s a ↔ ∃ ss', trs (s', a) = some ss' ∧ s ∈ ss' -- s' -a-> s is a transition
  suffices h : motive (m.reverse) (m.trans.toPFun) by
    specialize h (by rintro s a ss' heq; exact ⟨hwf.trans_src_lt _ (Std.HashMap.mem_of_getElem? heq), fun s' hin => hwf.trans_tgt_lt_internal heq hin⟩)
    simp_all

  apply Std.HashMap.fold_induction
  · simp [motive, tr]
    constructor <;> simp_all [tr]
  · rintro m' ⟨s, a⟩ ss' trs hnew ih
    rintro hsts
    specialize ih (by rintro s₀ a₀ ss₀ heq; apply hsts _ a₀; simp [Function.update_apply]; split <;> simp_all)
    rcases ih with ⟨hwf', hm', his', hfs', htrs'⟩
    specialize hsts s a ss' (by simp)
    have hs : s ∈ m'.states := by simp_all [states]
    have hss' : ∀ s ∈ ss'.toSet, s ∈ m'.states := by simp_all [states]
    obtain ⟨hwf'', hm'', his'', hfs'', htrs''⟩ := reverse_spec_procesState hwf' s a ss' hs hss'
    simp_all only [true_and]
    rintro s₁ b s₂
    by_cases hcond : s₂ = s ∧ b = a
    · rcases hcond with ⟨rfl, rfl, hin⟩
      simp_all
    · simp_all


syntax "neg" : MLIR.Pretty.uniform_op

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

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

Zipper

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

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 }

matchArgRes

def matchArgRes (lets : Lets d Γ_in eff Γ_out)
    (matchLets : Lets d Δ_in .pure Δ_out)
    (vs : HVector Γ_out.Var ts)
    (ws : HVector Δ_out.Var ts) :
    Option (MatchArgResult lets matchLets vs ws ∅) := do
  (matchArg lets matchLets vs ws ∅).attach.map fun ⟨⟨_, _⟩, h⟩ => .mk h

matchArg

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

matchVar

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

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

MatchVar

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

MatchVarM

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

Mapping

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

MatchVarM.unifyVars

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

appendCases

@[elab_as_elim]
def appendCases
    {motive : (Γ ++ Δ).Var t → Sort u}
    (left : (v : Var Γ t) → motive (appendInl v))
    (right : (v : Var Δ t) → motive (appendInr v)) :
    (v : (Γ ++ Δ).Var t) → motive v
  | ⟨idx, h⟩ =>
    if hv : idx < Γ.length then
      left ⟨idx, by admit /- proof elided -/
      ⟩
    else
      let v' : Var _ _ := ⟨idx - Γ.length, by admit /- proof elided -/
      ⟩
      have eq : v'.appendInr = ⟨idx, h⟩ := by admit /- proof elided -/

last

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

Valuation.cons

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

MatchArgResult

def MatchArgResult := { mapOut : Mapping _ _ //
  ∃ (mapIn' mapOut' : Mapping _ _),
    mapIn.entries ⊆ mapIn'.entries
    ∧ mapOut'.entries ⊆ mapOut.entries
    ∧ matchArg lets matchLets vs ws mapIn' = some ((), mapOut') }

toHom

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

Hom

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

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

Expr.pdenoteOp

@[simp] abbrev Expr.pdenoteOp :
    Expr d Γ .pure ty → Γ.Valuation → (HVector toType ty) :=
  Expr.denoteOp

Expr.ty

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

Expr.denoteOp

def Expr.denoteOp (e : Expr d Γ eff ty) (V : Γ.Valuation) :
    eff.toMonad d.m (HVector toType ty) :=
  EffectKind.liftEffect e.eff_le <| cast (by admit /- proof elided -/
  ) <|
    DialectDenote.denote e.op (e.args.map V) e.regArgs.denote

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

Signature

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

DialectSignature

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

regSig

def regSig       := Signature.regSig ∘ s.signature

RegionSignature

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

Ty

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

effectKind

def effectKind   := Signature.effectKind ∘ s.signature

returnTypes

def returnTypes  := Signature.returnTypes ∘ s.signature

Dialect

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

Op

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

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

TyDenote

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

Valuation.instAppendHVector

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

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)

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

Expr.outContext

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

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

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

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

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

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

toCom

def toCom (zip : Zipper d Γ_in eff ty) : Com d Γ_in eff ty :=
  go zip.top zip.bot
  where
    go : {Γ_mid : _} → Lets d Γ_in eff Γ_mid → Com d Γ_mid eff ty → Com d Γ_in eff ty
      | _, .nil, com          => com
      | _, .var body e, com  => go body (.var e com)

denote

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

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

MatchVarResult.toHom

def MatchVarResult.toHom
    (map : MatchVarResult lets v matchLets w mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.vars w) :
    Δ_in.Hom Γ_out :=
  map.val.toHom <| map.isTotal_of hvars

def MatchVarResult.isTotal_of
    (map : MatchVarResult lets v matchLets w mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.vars w) :
    map.val.IsTotal :=

def MatchArgResult.isTotal_of
    (map : MatchArgResult lets matchLets vs ws mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.varsOfVec ws) :
    map.val.IsTotal :=

MatchArgResult.toHom

def MatchArgResult.toHom
    (map : MatchArgResult lets matchLets vs ws mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.varsOfVec ws) :
    Δ_in.Hom Γ_out :=
  map.val.toHom <| map.isTotal_of hvars

variable
  {Γ_in Γ_out Δ_in Δ_out : Ctxt d.Ty}
  {lets : Lets d Γ_in eff Γ_out}
  {matchTy}
  {v : Var Γ_out matchTy}
  {matchLets : Lets d Δ_in .pure Δ_out}
  {w : Var Δ_out matchTy}
in

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)

MatchVarResult

def MatchVarResult := { mapOut : Mapping _ _ //
  ∃ (mapIn' mapOut' : Mapping _ _),
    mapIn.entries ⊆ mapIn'.entries
    ∧ mapOut'.entries ⊆ mapOut.entries
    ∧ matchVar lets v matchLets w mapIn' = some ((), mapOut') }

toArgResult

noncomputable def toArgResult
    (mapOut : MatchVarResult lets v (.var matchLets matchExpr) w.appendInl mapIn) :
    let args := mapOut.getPureExpr_eq_some.choose
    MatchArgResult lets matchLets args matchExpr.args mapIn :=
  ⟨mapOut.1, by admit /- proof elided -/
  ⟩

appendInl

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

eqvVarLeft

def eqvVarLeft  :
    MatchVarResult lets v (.var matchLets matchExpr) w.appendInr ma
    ≃ MatchVarResult lets v matchLets w ma where
  toFun := fun ⟨x, h⟩ => ⟨x, by admit /- proof elided -/
  ⟩
  invFun := fun ⟨x, h⟩ => ⟨x, by admit /- proof elided -/
  ⟩

Expr.changeVars

def Expr.changeVars (varsMap : Γ.Hom Γ') {ty} (e : Expr d Γ eff ty) :
    Expr d Γ' eff ty :=
  ⟨e.op, e.ty_eq, e.eff_le, e.args.map varsMap, e.regArgs⟩

Com.changeVars

def Com.changeVars : Com d Γ eff ty →
    (varsMap : Γ.Hom Γ') →
    Com d Γ' eff ty
  |  .rets e => fun varsMap => .rets (e.map varsMap)
  |  .var e body => fun varsMap => .var (e.changeVars varsMap)
      (body.changeVars (fun _ v => varsMap.append v))

Hom.append

def Hom.append {ζ : Ctxt Ty} (f : Γ.Hom Δ) : Hom (ζ ++ Γ) (ζ ++ Δ) :=
  fun _ => Var.appendCases
    (fun v => v.appendInl)
    (fun v => (f v).appendInr)

append

def append (d₁ : Diff Γ₁ Γ₂) (d₂ : Diff Γ₂ Γ₃) : Diff Γ₁ Γ₃ :=
  {val := d₁.val + d₂.val,  property := append_valid d₁.property d₂.property}

insertPureCom

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)

insertCom

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 }

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

Com.outContextHom

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

Com.outContextDiff

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

Expr.bvars

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

Com.bvars

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

castCtxt

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

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

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⟩

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] theorem denote_toCom [LawfulMonad d.m] (zip : Zipper d Γ_in eff ty) :
    zip.toCom.denote = zip.denote

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

theorem denote_matchLets_of
    (map : MatchArgResult lets matchLets vs ws mapIn)
    (hvars : ∀ t (v : Var Δ_in t), ⟨t, v⟩ ∈ matchLets.varsOfVec ws)
    (V : lets.ValidDenotation) :
    ws.map (matchLets.denote (V.val.comap <| map.toHom h)) = vs.map V.val

theorem denote_matchArg
    {vs ws : HVector (Var _) ts}
    (mapOut : MatchArgResult lets matchLets vs ws mapIn)
    (V : lets.ValidDenotation) :
    HVector.map (matchLets.denote (mapOut.val.mapValuation V.val)) ws = HVector.map (V.val) vs

theorem HVector.map_eq_map_of_matchArg
    {lets : Lets d Γ_in eff Γ_out}
    {matchLets : Lets d Δ_in .pure Δ_out}
    {ma : Mapping Δ_in Γ_out}
    {l : List d.Ty} {args₁ : HVector _ l} {args₂ : HVector _ l}
    (mapOut : MatchArgResult lets matchLets args₁ args₂ ma)
    (f₁ f₂ : (t : d.Ty) → Var _ t → ⟦t⟧)
    (hf : ∀ {t v₁ v₂},
      (mapOut' : MatchVarResult lets v₁ matchLets v₂ ma)
      → mapOut'.val = mapOut.val
      → f₂ t v₂ = f₁ t v₁) :
    HVector.map f₂ args₂ = HVector.map f₁ args₁

theorem denote_matchVar
    {v w : Var _ t}
    (mapOut : MatchVarResult lets v matchLets w mapIn)
    (V : lets.ValidDenotation) :
    (matchLets.denote (mapOut.val.mapValuation V.val) w)
    = V.val v

theorem getPureExpr_eq_some
    (mapOut : MatchVarResult lets v (.var matchLets matchExpr) w.appendInl mapIn) :
    ∃ args, lets.getPureExpr v = some ⟨te, w, ⟨
        matchExpr.op,
        matchExpr.ty_eq,
        matchExpr.eff_le,
        args,
        matchExpr.regArgs
      ⟩⟩

theorem matchVar_appendInl {w : Var ⟨te⟩ t} :
    matchVar lets v (.var matchLets matchExpr) w.appendInl ma = some ma' →
    ∃ args,
      lets.getPureExpr v
        = some ⟨_, w, matchExpr.op, matchExpr.ty_eq, matchExpr.eff_le, args, matchExpr.regArgs⟩
      ∧ matchArg lets matchLets args matchExpr.args ma = some ma'

@[simp]
theorem MatchVar.liftM_bind_eq_some_iff (x? : Option α)
    (f : α → MatchVarM Δ Γ β) :
    ((liftM x? >>= f) mapIn = some mapOut)
    ↔ ( ∃ h : x?.isSome,
        f (x?.get h) mapIn = some mapOut )

@[simp] theorem Com.denote_changeVars
    (varsMap : Γ.Hom Γ') (c : Com d Γ eff ty) :
    (c.changeVars varsMap).denote =
    fun V => c.denote (V.comap varsMap)

theorem denote_insertPureCom_eq_of [LawfulMonad d.m]
    {zip : Zipper d Γ_in eff tys} {vs}
    {newCom : Com d zip.Γ_mid .pure newTys} {V_in : Valuation Γ_in}
    (h : ∀ V : zip.top.ValidDenotation,
        newCom.denote V.val = vs.map V.val) :
    (zip.insertPureCom vs newCom).denote V_in = zip.denote V_in

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

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

SplitProgramResult

structure SplitProgramResult extends Zipper d Γ eff t where
  {midTypes : List d.Ty}
  midRet : HVector toZipper.Γ_mid.Var midTypes

splitProgramAt

def splitProgramAt (pos : ℕ) (prog : Com d Γ eff t) :
    Option (SplitProgramResult d Γ eff t) :=
  splitProgramAtAux pos .nil prog

rewriteAt

theorem denote_splitProgramAtAux [LawfulMonad d.m] :
    {pos : ℕ} → {lets : Lets d Γ₁ eff Γ₂} →
    {prog : Com d Γ₂ eff t} →
    {res : _} → (hres : res ∈ splitProgramAtAux pos lets prog) →
    (V : Valuation Γ₁) →
    res.denote V = (lets.denote V) >>= prog.denote
  | 0, lets, .var e body, res, hres, V => by
    obtain rfl

@[simp]
theorem denote_splitProgramAt [LawfulMonad d.m] {pos : ℕ} {prog : Com d Γ eff t}
    {res : _} (hres : res ∈ splitProgramAt pos prog) :
    res.denote = prog.denote

import LeanMLIR.Framework

import LeanMLIR.Framework.Zipper

import LeanMLIR.Transforms.Rewrite.Match

open Ctxt (Var VarSet Valuation)

variable {d} [DialectSignature d] [TyDenote d.Ty] [DialectDenote d] [Monad d.m]

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

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

section SplitProgram

variable (d Γ eff t) in

structure SplitProgramResult extends Zipper d Γ eff t where
  {midTypes : List d.Ty}
  midRet : HVector toZipper.Γ_mid.Var midTypes

def splitProgramAtAux : (pos : ℕ) → (lets : Lets d Γ₁ eff Γ₂) →
    (prog : Com d Γ₂ eff t) →
    Option (SplitProgramResult d Γ₁ eff t)
  | 0, lets, .var e body => some {
      top := lets.var e
      bot := body
      midRet := e.returnVars
    }
  | _, _, .rets _ => none
  | n+1, lets, .var e body =>
    splitProgramAtAux n (lets.var e) body

def splitProgramAt (pos : ℕ) (prog : Com d Γ eff t) :
    Option (SplitProgramResult d Γ eff t) :=
  splitProgramAtAux pos .nil prog

end SplitProgram

def rewriteAt
    (lhs rhs : Com d Γ₁ .pure ts₁)
    (hlhs : ∀ t (v : Var Γ₁ t), ⟨t, v⟩ ∈ lhs.vars)
    (pos : ℕ) (target : Com d Γ₂ eff t₂) :
    Option (Com d Γ₂ eff t₂) := do
  let splitRes ← splitProgramAt pos target
  if h : ts₁ = splitRes.midTypes then
    let m ← matchArgRes splitRes.top lhs.toLets splitRes.midRet (h ▸ lhs.returnVars)
    let m := m.toHom <| by
      subst h; exact hlhs
    let rhs := rhs.changeVars m
    let zip := splitRes.insertPureCom splitRes.midRet (cast (by admit /- proof elided -/
    ) rhs)
    return zip.toCom
  else none

theorem denote_rewriteAt [LawfulMonad d.m]
    {lhs rhs : Com d Γ₁ .pure t₁}
    (hl : lhs.denote = rhs.denote)
    {hlhs : ∀ t (v : Var Γ₁ t), ⟨t, v⟩ ∈ lhs.vars}
    {pos : ℕ} {target : Com d Γ₂ eff t₂}
    {rew : Com d Γ₂ eff t₂}
    (hrew : rew ∈ rewriteAt lhs rhs hlhs pos target) :
    rew.denote = target.denote :=

:= by
  funext V
  simp only [rewriteAt, Option.pure_def, Option.bind_eq_bind, Option.mem_def,
    Option.bind_eq_some_iff, Option.dite_none_right_eq_some, Option.some.injEq] at hrew
  rcases hrew with ⟨res, h_split, rfl, varMap', -, rfl⟩
  dsimp only at varMap'
  simp only [cast_eq, Zipper.denote_toCom, ← denote_splitProgramAt h_split]
  apply Zipper.denote_insertPureCom_eq_of
  intro V
  simp only [Expr.pdenoteOp, Com.denote_changeVars, ← hl]
  rw [← Com.denoteLets_returnVars, Com.denoteLets_eq]
  apply denote_matchLets_of varMap' hlhs