import Blase.MultiWidth.Defs
import Blase.Vars
import Blase.KInduction.KInduction
import Lean
import Blase.Blase.Fast.BitStream
import Blase.Fast.FiniteStateMachine

syntax "min" : MLIR.Pretty.uniform_op

syntax "min" : 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

FSM

structure FSM (arity : Type) : Type 1 where

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

  ( initCarry : α → Bool )

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

Var

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

composeBinaryAux'

def composeBinaryAux'
    (p : FSM Bool)
    (qtrue : FSM α)
    (qfalse : FSM α) :
    FSM α :=
  p.compose (α)
    (λ _ => α)
    (λ _ i => i)
    (λ b => match b with
      | true => qtrue
      | false => qfalse)

compose

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

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

xor

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

composeUnaryAux

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

ls

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

latchImmediate

def latchImmediate (initVal : Bool) : FSM Bool where
  α := Unit
  initCarry := fun _ => initVal
  outputCirc :=
    let xval := Circuit.var true (inr false)
    let control := Circuit.var true (inr true)
    let state := Circuit.var true (inl ())
    Circuit.ite control xval state
  nextStateCirc := fun () =>
    let xval := Circuit.var true (inr false)
    let control := Circuit.var true (inr true)
    let state := Circuit.var true (inl ())
    Circuit.ite control xval state

ite

def ite (cond t f : Circuit α) : Circuit α :=
  (cond &&& t) ||| (~~~ cond &&& f)

TermFSM

structure TermFSM (wcard tcard bcard ncard icard pcard : Nat) (t : Nondep.Term) where
  toFsmZext : FSM (StateSpace wcard tcard bcard ncard icard pcard)
  width : NatFSM wcard tcard bcard ncard icard pcard t.width

NatFSM

structure NatFSM (wcard tcard bcard ncard icard pcard : Nat) (v : Nondep.WidthExpr) where
  toFsm : FSM (StateSpace wcard tcard bcard ncard icard pcard)

StateSpace

inductive StateSpace (wcard tcard bcard ncard icard pcard : Nat)
| widthVar (v : Fin wcard)
| termVar (v : Fin tcard)
| predVar (v : Fin pcard)
| boolVar (v : Fin bcard)
deriving DecidableEq, Repr, Hashable

Term

inductive Term
| ofNat (w : WidthExpr) (n : Nat) : Term
| var (v : Nat) (w : WidthExpr) : Term
| add (w : WidthExpr) (a b : Term) : Term
| zext (a : Term) (wnew : WidthExpr) : Term
| setWidth (a : Term) (wnew : WidthExpr) : Term
| sext (a : Term) (wnew : WidthExpr) : Term
| bor (w : WidthExpr) (a b : Term) : Term
| band (w : WidthExpr) (a b : Term) : Term
| bxor (w : WidthExpr) (a b : Term) : Term
| bnot (w : WidthExpr)  (a : Term) : Term
| boolVar (v : Nat) : Term
| boolConst (b : Bool) : Term
| shiftl (w : WidthExpr) (a : Term) (k : Nat) : Term
| bvOfBool (b : Term) : Term
| binWidthRel (k : WidthBinaryRelationKind) (wa wb : WidthExpr) : Term
| binRel (k : BinaryRelationKind) (w : WidthExpr)
    (a : Term) (b : Term) : Term
| or (p1 p2 : Term) : Term
| and (p1 p2 : Term) : Term
| pvar (v : Nat) : Term
| boolBinRel (k : BoolBinaryRelationKind)
    (a b : Term) : Term
deriving DecidableEq, Inhabited, Repr, Lean.ToExpr

WidthExpr

inductive WidthExpr where
| const : Nat → WidthExpr
| var : Nat → WidthExpr
| max : WidthExpr → WidthExpr → WidthExpr
| min : WidthExpr → WidthExpr → WidthExpr
| addK : WidthExpr → Nat → WidthExpr
| kadd : Nat → WidthExpr → WidthExpr
deriving Inhabited, Repr, Hashable, DecidableEq, Lean.ToExpr

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

WidthExpr

inductive WidthExpr where
| const : Nat → WidthExpr
| var : Nat → WidthExpr
| max : WidthExpr → WidthExpr → WidthExpr
| min : WidthExpr → WidthExpr → WidthExpr
| addK : WidthExpr → Nat → WidthExpr
| kadd : (k : Nat) → (v : WidthExpr) → WidthExpr
deriving Inhabited, Repr, Hashable, DecidableEq, Lean.ToExpr

HNatFSMToBitstream

structure HNatFSMToBitstream {wcard : Nat} {v : WidthExpr wcard} {tcard : Nat} {bcard : Nat} {pcard : Nat}
   (fsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v)) : Prop where
  heq :
    ∀ (wenv : Fin wcard → Nat)
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
    (henv : HWidthEnv fsmEnv wenv) →
      fsm.toFsm.eval fsmEnv =
      BitStream.ofNatUnary (v.toNat wenv)

WidthExpr.toNat

def WidthExpr.toNat (e : WidthExpr wcard) (env : WidthExpr.Env wcard) : Nat :=
  match e with
  | .const n => n
  | .var v => env v
  | .min v w => Nat.min (v.toNat env) (w.toNat env)
  | .max v w => Nat.max (v.toNat env) (w.toNat env)
  | .addK v k => v.toNat env + k
  | .kadd k v => k + v.toNat env

WidthExpr

inductive WidthExpr (wcard : Nat) : Type
| const (n : Nat) :  WidthExpr wcard
| var : (v : Fin wcard) → WidthExpr wcard
| min : (v w : WidthExpr wcard) → WidthExpr wcard
| max : (v w : WidthExpr wcard) → WidthExpr wcard
| addK : (v : WidthExpr wcard) → (k : Nat) → WidthExpr wcard
| kadd : (k : Nat) → (v : WidthExpr wcard) → WidthExpr wcard

WidthExpr.Env

abbrev WidthExpr.Env (wcard : Nat) : Type :=
  Fin wcard → Nat

HWidthEnv

structure HWidthEnv {wcard tcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (wenv : Fin wcard → Nat) : Prop where
    heq_width : ∀ (v : Fin wcard),
      fsmEnv (StateSpace.widthVar v) = BitStream.ofNatUnary (wenv v)

HPredicateEnv

structure HPredicateEnv {wcard tcard bcard ncard icard pcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (penv : Fin pcard → Prop) : Prop where
    heq_width : ∀ (v : Fin pcard),
      fsmEnv (StateSpace.predVar v) = BitStream.ofProp (penv v)

BitStream

def BitStream : Type := Nat → Bool

ofNatUnary

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

HPredFSMToBitStream

structure HPredFSMToBitStream {pcard : Nat}
  {tctx : Term.Ctx wcard tcard}
  {p : Term bcard ncard icard pcard tctx .prop}
  (fsm : TermFSM wcard tcard bcard ncard icard pcard
    (.ofDepTerm p)) : Prop where
  heq :
    ∀ {wenv : WidthExpr.Env wcard}
      (benv : Term.BoolEnv bcard)
      (nenv : Term.NatEnv ncard)
      (ienv : Term.IntEnv icard)
      (penv : Predicate.Env pcard) (tenv : tctx.Env wenv)
      (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
      (htenv : HTermEnv fsmEnv tenv benv) →
      (hpenv : HPredicateEnv fsmEnv penv) →
        p.toBV benv nenv ienv penv tenv  ↔ (fsm.toFsmZext.eval fsmEnv = .negOne)

Term.Ctx

abbrev Term.Ctx (wcard : Nat) (tcard : Nat) : Type :=
  Fin tcard → WidthExpr wcard

Term.BoolEnv

def Term.BoolEnv (bcard : Nat) : Type := Fin bcard → Bool

Term.IntEnv

def Term.IntEnv (icard : Nat) : Type := Fin icard → Nat

HTermFSMToBitStream

structure HTermFSMToBitStream {w : WidthExpr wcard}
  {tctx : Term.Ctx wcard tcard}
  {t : Term bcard ncard icard pcard tctx (.bv w)}
  (fsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm t)) : Prop where
  heq :
    ∀ {wenv : WidthExpr.Env wcard}
      (benv : Term.BoolEnv bcard)
      (nenv : Term.NatEnv ncard)
      (ienv : Term.IntEnv icard)
      (penv : Predicate.Env pcard) (tenv : tctx.Env wenv)
      (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
      (henv : HTermEnv fsmEnv tenv benv) →
        fsm.toFsmZext.eval fsmEnv =
        BitStream.ofBitVecZext (t.toBV benv nenv ienv penv tenv)

Predicate.Env

def Predicate.Env (pcard : Nat) : Type :=
  Fin pcard → Prop

TermKind

inductive TermKind (wcard : Nat) : Type
| bool
| bv (w : WidthExpr wcard)  : TermKind wcard
| prop
| nat
| int

HTermEnv

structure HTermEnv {wcard tcard bcard : Nat}
    {wenv : Fin wcard → Nat} {tctx : Term.Ctx wcard tcard}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (tenv : tctx.Env wenv)
    (benv : Term.BoolEnv bcard) : Prop
  extends HWidthEnv fsmEnv wenv where
    heq_term : ∀ (v : Fin tcard),
      fsmEnv (StateSpace.termVar v) = BitStream.ofBitVecZext (tenv v)
    heq_bool : ∀ (v : Fin bcard),
      fsmEnv (StateSpace.boolVar v) = BitStream.ofBool (benv v)

BitStream.ofBool

noncomputable def BitStream.ofBool (b : Bool) : BitStream := fun _i => b

Term.Ctx.Env

abbrev Term.Ctx.Env
  (tctx : Term.Ctx wcard tcard)
  (wenv : WidthExpr.Env wcard) :=
  (v : Fin tcard) → BitVec ((tctx v).toNat wenv)

ofBitVecZext

abbrev ofBitVecZext {w} (x : BitVec w) : BitStream :=
  fun i => x.getLsbD i

Term.NatEnv

def Term.NatEnv (ncard : Nat) : Type := Fin ncard → Nat

Term.toBV

def Term.toBV {wenv : WidthExpr.Env wcard}
    {tctx : Term.Ctx wcard tcard}
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (penv : Predicate.Env pcard)
    (tenv : tctx.Env wenv)
    (t : Term bcard ncard icard pcard tctx k) : k.denote wenv :=
match t with
| .ofNat w n => BitVec.ofNat (w.toNat wenv) n
| .boolConst b => b
| .var v => tenv.get v.1 v.2
| .add (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a + b
| .zext a v => (a.toBV benv nenv ienv penv tenv).zeroExtend (v.toNat wenv)
| .setWidth a v => (a.toBV benv nenv ienv penv tenv).zeroExtend (v.toNat wenv)
| .sext a v => (a.toBV benv nenv ienv penv tenv).signExtend (v.toNat wenv)
| .bor a b (w := w) =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a ||| b
| .band (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a &&& b
| .bxor (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a ^^^ b
| .bnot (w := w) a =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    ~~~ a
| .boolVar v => benv v
| .shiftl (w := w) a k =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    a <<< k
| .bvOfBool b => BitVec.ofBool (b.toBV benv nenv ienv penv tenv)

| .binWidthRel rel wa wb =>
  match rel with
  | .eq => wa.toNat wenv = wb.toNat wenv
  | .le => wa.toNat wenv ≤ wb.toNat wenv
| .binRel rel _w a b =>
  match rel with
  | .eq => a.toBV benv nenv ienv penv tenv = b.toBV benv nenv ienv penv tenv
  | .ne => a.toBV benv nenv ienv penv tenv ≠ b.toBV benv nenv ienv penv tenv
  | .ult => (a.toBV benv nenv ienv penv tenv).ult (b.toBV benv nenv ienv penv tenv) = true
  | .ule => (a.toBV benv nenv ienv penv tenv).ule (b.toBV benv nenv ienv penv tenv) = true
  | .slt => (a.toBV benv nenv ienv penv tenv).slt (b.toBV benv nenv ienv penv tenv) = true
  | .sle => (a.toBV benv nenv ienv penv tenv).sle (b.toBV benv nenv ienv penv tenv) = true
| .and p1 p2 => p1.toBV benv nenv ienv penv tenv  ∧ p2.toBV benv nenv ienv penv tenv
| .or p1 p2 => p1.toBV benv nenv ienv penv tenv ∨ p2.toBV benv nenv ienv penv tenv
| .boolBinRel rel a b =>
  match rel with

  | .eq => (a.toBV benv nenv ienv penv tenv) = (b.toBV benv nenv ienv penv tenv)
| .pvar v => penv v

Term

inductive Term {wcard tcard : Nat} (bcard : Nat) (ncard : Nat) (icard : Nat) (pcard : Nat)
  (tctx : Term.Ctx wcard tcard) : TermKind wcard → Type

| ofNat (w : WidthExpr wcard) (n : Nat) : Term bcard ncard icard pcard tctx (.bv w)

| var (v : Fin tcard) : Term bcard ncard icard pcard tctx (.bv (tctx v))

| add (a : Term bcard ncard icard pcard tctx (.bv w))
  (b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| shiftl (a : Term bcard ncard icard pcard tctx (.bv w)) (k : Nat) : Term bcard ncard icard pcard tctx (.bv w)

| bor (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| band (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| bxor (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| bnot (a : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| zext (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| setWidth (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| sext (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| bvOfBool (b : Term bcard ncard icard pcard tctx .bool) : Term bcard ncard icard pcard tctx (.bv (.const 1))

| boolConst (b : Bool) : Term bcard ncard icard pcard tctx .bool
| boolVar (v : Fin bcard) : Term bcard ncard icard pcard tctx .bool
| binWidthRel (k : WidthBinaryRelationKind) (wa wb : WidthExpr wcard) :
    Term bcard ncard icard pcard tctx .prop
| binRel
    (k : BinaryRelationKind)
    (w : WidthExpr wcard)
    (a : Term bcard ncard icard pcard tctx (.bv w))
    (b : Term bcard ncard icard pcard tctx (.bv w)) :
    Term bcard ncard icard pcard tctx .prop
| and (p1 p2 : Term bcard ncard icard pcard tctx (.prop)) : Term bcard ncard icard pcard tctx (.prop)
| or (p1 p2 : Term bcard ncard icard pcard tctx (.prop)) : Term bcard ncard icard pcard tctx (.prop)
| pvar (v : Fin pcard) : Term bcard ncard icard pcard tctx (.prop)

| boolBinRel
  (k : BoolBinaryRelationKind)
  (a b : Term bcard ncard icard pcard tctx .bool) :
  Term bcard ncard icard pcard tctx (.prop)

Term.Ctx.Env.get

def Term.Ctx.Env.get {tcard : Nat}
  {wcard : Nat} {wenv : Fin wcard → Nat}
  {tctx : Term.Ctx wcard tcard}
  (tenv : tctx.Env wenv) (i : Nat) (hi : i < tcard) :
  BitVec ((tctx ⟨i, hi⟩).toNat wenv) :=
  tenv ⟨i, hi⟩

BinaryRelationKind

inductive BinaryRelationKind
| eq
| ne
| ule
| slt
| sle
| ult
deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

WidthBinaryRelationKind

inductive WidthBinaryRelationKind
| eq
| le

deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

ofBitVecSext

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

zero

abbrev zero   : BitStream := fun _ => false

carry'

def carry' (initCarry : Bool) (x y : BitStream) : BitStream :=
  fun n =>
  match n with
  | 0 => initCarry
  | n + 1 => (addAux' initCarry x y n).2

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)

protected theorem carry'_eq_carry (x y : BitStream) (c : Bool)
    (x' y' : BitVec w)
    (hx : ∀ i, i < n → x'.getLsbD i = x i)
    (hy : ∀ i,  i < n → y'.getLsbD i = y i) :
    carry' c x y n = (BitVec.carry n x' y' c)

@[simp] theorem carry'_succ (initCarry : Bool) (x y : BitStream) :
    (carry' initCarry x y (i + 1)) =
    let out

MultiWidth.fsmTermSlt

@[simp]
theorem eval_fsmMsb_eq {wcard bcard tcard : Nat}
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    {tctx : Term.Ctx wcard tcard}
    (tenv : Term.Ctx.Env tctx wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (w : WidthExpr wcard)
    (x : Term bcard ncard icard pcard tctx (.bv w))
    (xfsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm x))
    (hxfsm : HTermFSMToBitStream xfsm)
    (wfsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    (hwfsm : HNatFSMToBitstream wfsm)
    (htenv : HTermEnv fsmEnv tenv benv) :
    (fsmMsb xfsm.toFsmZext wfsm.toFsm).eval fsmEnv = (fun i =>
      BitStream.ofBitVecZext (x.toBV benv nenv ienv penv tenv) (min i (w.toNat wenv - 1)))

theorem eval_fsmMsb_eq_BitStream_ofBitVecSext {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    {tctx : Term.Ctx wcard tcard}
    (tenv : Term.Ctx.Env tctx wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (w : WidthExpr wcard)
    (x : Term bcard ncard icard pcard tctx (.bv w))
    (xfsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm x))
    (hxfsm : HTermFSMToBitStream xfsm)
    (wfsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    (hwfsm : HNatFSMToBitstream wfsm)
    (htenv : HTermEnv fsmEnv tenv benv) :
    (fsmMsb xfsm.toFsmZext wfsm.toFsm).eval fsmEnv =
      BitStream.ofBitVecSext (x.toBV benv nenv ienv penv tenv)

import Blase.Fast.FiniteStateMachine

import Blase.Vars

import Blase.MultiWidth.Defs

import Blase.KInduction.KInduction

import Lean

namespace MultiWidth

def fsmMsb (x w : FSM α) : FSM α :=
  composeBinaryAux'
    (FSM.latchImmediate false)
    (qfalse := x)
    (qtrue := w)

def fsmCarry'' (initialCarryVal : Bool): FSM Bool :=
  let outputCirc :=
    let carry := Circuit.var true (Sum.inl ())
    let a := Circuit.var true (Sum.inr true)
    let b := Circuit.var true (Sum.inr false)

    ((a &&& b) ||| (a &&& carry) ||| (b &&& carry))
  { α := Unit,

    initCarry := fun () => initialCarryVal,
    outputCirc := Circuit.var true (Sum.inl ()) ,
    nextStateCirc := fun () => outputCirc
  }

def fsmMsbEq (a : FSM α) (b : FSM α) : FSM α :=
  composeUnaryAux (FSM.ls false) <|
    composeBinaryAux' FSM.xor a b

def fsmTermSlt
  {wcard tcard : Nat}
  {w : Nondep.WidthExpr}
  {a b : Nondep.Term}
  (wfsm : NatFSM wcard tcard bcard ncard icard pcard w)
  (afsm : TermFSM wcard tcard bcard ncard icard pcard a)
  (bfsm : TermFSM wcard tcard bcard ncard icard pcard b)
  : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
    let afsm := afsm.toFsmZext
    let bfsm := bfsm.toFsmZext
    let afsm := fsmMsb afsm wfsm.toFsm
    let bfsm := fsmMsb bfsm wfsm.toFsm
    let carryFsm :=
      (~~~ (composeBinaryAux' (fsmCarry'' true)  afsm (~~~ bfsm)))
    let xorFsm := fsmMsbEq afsm bfsm
    let val := xorFsm ^^^ carryFsm
    val

theorem eval_fsmTermSlt_eq_decide_slt {wcard tcard : Nat}
    (tctx : Term.Ctx wcard tcard)
    {wenv : WidthExpr.Env wcard}
    (tenv : tctx.Env wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (w : WidthExpr wcard)
    (a : Term bcard ncard icard pcard tctx (.bv w))
    (b : Term bcard ncard icard pcard tctx (.bv w))
    (wfsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    (hwfsm : HNatFSMToBitstream wfsm)
    (afsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm a))
    (hafsm : HTermFSMToBitStream afsm)
    (bfsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm b))
    (hbfsm : HTermFSMToBitStream bfsm)
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (henv : HTermEnv fsmEnv tenv benv)
    :
    ((fsmTermSlt
      wfsm
      afsm
      bfsm)).eval fsmEnv i =
       decide (((a.toBV benv nenv ienv penv tenv).signExtend i).slt
          ((b.toBV benv nenv ienv penv tenv).signExtend i)) :=

:= by
  simp [fsmTermSlt]
  have := BitVec.slt_eq_not_carry
    (x := (a.toBV benv nenv ienv penv tenv).signExtend i)
    (y := (b.toBV benv nenv ienv penv tenv).signExtend i)
  rw [this]
  clear this
  simp [eval_fsmMsb_eq_BitStream_ofBitVecSext
        (hxfsm := hafsm) (hwfsm := hwfsm)
        (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv)
        (tenv := tenv) (htenv := henv)]
  simp [eval_fsmMsb_eq_BitStream_ofBitVecSext (hxfsm := hbfsm) (hwfsm := hwfsm)
    (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv)
    (tenv := tenv) (htenv := henv)]
  rw [BitStream.carry'_eq_carry
      (x' := BitVec.signExtend i (Term.toBV benv nenv ienv penv tenv a))
      (y' := ~~~ BitVec.signExtend i (Term.toBV benv nenv ienv penv tenv b))]
  -- simp [fsmMsbEq]
  · rcases i with rfl | i
    · simp
      simp [BitVec.of_length_zero]
      simp [fsmMsbEq]
    · simp [fsmMsbEq]
      simp [eval_fsmMsb_eq_BitStream_ofBitVecSext
            (hxfsm := hafsm) (hwfsm := hwfsm)
            (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv)
            (tenv := tenv) (htenv := henv)]
      simp [eval_fsmMsb_eq_BitStream_ofBitVecSext (hxfsm := hbfsm) (hwfsm := hwfsm)
        (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv)
        (tenv := tenv) (htenv := henv)]
      simp [BitStream.ofBitVecSext]
      by_cases hi : i < w.toNat wenv
      · simp [BitVec.msb_eq_getLsbD_last, BitVec.getElem_signExtend]
        simp [hi]
        grind [Bool]
      · simp [BitVec.msb_eq_getLsbD_last, BitVec.getElem_signExtend]
        simp [hi]
        grind [Bool]
  · intros j
    intros hj
    simp [hj]
    simp [BitVec.getElem_signExtend]
    simp [BitStream.ofBitVecSext]
    by_cases hw : j < w.toNat wenv
    · simp [hw]
    · simp [hw]
  · intros j
    simp
    intros hj
    simp [hj]
    simp [BitVec.getElem_signExtend]
    simp [BitStream.ofBitVecSext]
    by_cases hw : j < w.toNat wenv
    · simp [hw]
    · simp [hw]


syntax "min" : MLIR.Pretty.uniform_op

syntax "min" : 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

FSM

structure FSM (arity : Type) : Type 1 where

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

  ( initCarry : α → Bool )

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

Var

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

composeBinaryAux'

def composeBinaryAux'
    (p : FSM Bool)
    (qtrue : FSM α)
    (qfalse : FSM α) :
    FSM α :=
  p.compose (α)
    (λ _ => α)
    (λ _ i => i)
    (λ b => match b with
      | true => qtrue
      | false => qfalse)

compose

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

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

xor

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

composeUnaryAux

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

ls

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

latchImmediate

def latchImmediate (initVal : Bool) : FSM Bool where
  α := Unit
  initCarry := fun _ => initVal
  outputCirc :=
    let xval := Circuit.var true (inr false)
    let control := Circuit.var true (inr true)
    let state := Circuit.var true (inl ())
    Circuit.ite control xval state
  nextStateCirc := fun () =>
    let xval := Circuit.var true (inr false)
    let control := Circuit.var true (inr true)
    let state := Circuit.var true (inl ())
    Circuit.ite control xval state

ite

def ite (cond t f : Circuit α) : Circuit α :=
  (cond &&& t) ||| (~~~ cond &&& f)

TermFSM

structure TermFSM (wcard tcard bcard ncard icard pcard : Nat) (t : Nondep.Term) where
  toFsmZext : FSM (StateSpace wcard tcard bcard ncard icard pcard)
  width : NatFSM wcard tcard bcard ncard icard pcard t.width

NatFSM

structure NatFSM (wcard tcard bcard ncard icard pcard : Nat) (v : Nondep.WidthExpr) where
  toFsm : FSM (StateSpace wcard tcard bcard ncard icard pcard)

StateSpace

inductive StateSpace (wcard tcard bcard ncard icard pcard : Nat)
| widthVar (v : Fin wcard)
| termVar (v : Fin tcard)
| predVar (v : Fin pcard)
| boolVar (v : Fin bcard)
deriving DecidableEq, Repr, Hashable

Term

inductive Term
| ofNat (w : WidthExpr) (n : Nat) : Term
| var (v : Nat) (w : WidthExpr) : Term
| add (w : WidthExpr) (a b : Term) : Term
| zext (a : Term) (wnew : WidthExpr) : Term
| setWidth (a : Term) (wnew : WidthExpr) : Term
| sext (a : Term) (wnew : WidthExpr) : Term
| bor (w : WidthExpr) (a b : Term) : Term
| band (w : WidthExpr) (a b : Term) : Term
| bxor (w : WidthExpr) (a b : Term) : Term
| bnot (w : WidthExpr)  (a : Term) : Term
| boolVar (v : Nat) : Term
| boolConst (b : Bool) : Term
| shiftl (w : WidthExpr) (a : Term) (k : Nat) : Term
| bvOfBool (b : Term) : Term
| binWidthRel (k : WidthBinaryRelationKind) (wa wb : WidthExpr) : Term
| binRel (k : BinaryRelationKind) (w : WidthExpr)
    (a : Term) (b : Term) : Term
| or (p1 p2 : Term) : Term
| and (p1 p2 : Term) : Term
| pvar (v : Nat) : Term
| boolBinRel (k : BoolBinaryRelationKind)
    (a b : Term) : Term
deriving DecidableEq, Inhabited, Repr, Lean.ToExpr

WidthExpr

inductive WidthExpr where
| const : Nat → WidthExpr
| var : Nat → WidthExpr
| max : WidthExpr → WidthExpr → WidthExpr
| min : WidthExpr → WidthExpr → WidthExpr
| addK : WidthExpr → Nat → WidthExpr
| kadd : Nat → WidthExpr → WidthExpr
deriving Inhabited, Repr, Hashable, DecidableEq, Lean.ToExpr

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

WidthExpr

inductive WidthExpr where
| const : Nat → WidthExpr
| var : Nat → WidthExpr
| max : WidthExpr → WidthExpr → WidthExpr
| min : WidthExpr → WidthExpr → WidthExpr
| addK : WidthExpr → Nat → WidthExpr
| kadd : (k : Nat) → (v : WidthExpr) → WidthExpr
deriving Inhabited, Repr, Hashable, DecidableEq, Lean.ToExpr

HNatFSMToBitstream

structure HNatFSMToBitstream {wcard : Nat} {v : WidthExpr wcard} {tcard : Nat} {bcard : Nat} {pcard : Nat}
   (fsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v)) : Prop where
  heq :
    ∀ (wenv : Fin wcard → Nat)
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
    (henv : HWidthEnv fsmEnv wenv) →
      fsm.toFsm.eval fsmEnv =
      BitStream.ofNatUnary (v.toNat wenv)

WidthExpr.toNat

def WidthExpr.toNat (e : WidthExpr wcard) (env : WidthExpr.Env wcard) : Nat :=
  match e with
  | .const n => n
  | .var v => env v
  | .min v w => Nat.min (v.toNat env) (w.toNat env)
  | .max v w => Nat.max (v.toNat env) (w.toNat env)
  | .addK v k => v.toNat env + k
  | .kadd k v => k + v.toNat env

WidthExpr

inductive WidthExpr (wcard : Nat) : Type
| const (n : Nat) :  WidthExpr wcard
| var : (v : Fin wcard) → WidthExpr wcard
| min : (v w : WidthExpr wcard) → WidthExpr wcard
| max : (v w : WidthExpr wcard) → WidthExpr wcard
| addK : (v : WidthExpr wcard) → (k : Nat) → WidthExpr wcard
| kadd : (k : Nat) → (v : WidthExpr wcard) → WidthExpr wcard

WidthExpr.Env

abbrev WidthExpr.Env (wcard : Nat) : Type :=
  Fin wcard → Nat

HWidthEnv

structure HWidthEnv {wcard tcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (wenv : Fin wcard → Nat) : Prop where
    heq_width : ∀ (v : Fin wcard),
      fsmEnv (StateSpace.widthVar v) = BitStream.ofNatUnary (wenv v)

HPredicateEnv

structure HPredicateEnv {wcard tcard bcard ncard icard pcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (penv : Fin pcard → Prop) : Prop where
    heq_width : ∀ (v : Fin pcard),
      fsmEnv (StateSpace.predVar v) = BitStream.ofProp (penv v)

BitStream

def BitStream : Type := Nat → Bool

ofNatUnary

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

HPredFSMToBitStream

structure HPredFSMToBitStream {pcard : Nat}
  {tctx : Term.Ctx wcard tcard}
  {p : Term bcard ncard icard pcard tctx .prop}
  (fsm : TermFSM wcard tcard bcard ncard icard pcard
    (.ofDepTerm p)) : Prop where
  heq :
    ∀ {wenv : WidthExpr.Env wcard}
      (benv : Term.BoolEnv bcard)
      (nenv : Term.NatEnv ncard)
      (ienv : Term.IntEnv icard)
      (penv : Predicate.Env pcard) (tenv : tctx.Env wenv)
      (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
      (htenv : HTermEnv fsmEnv tenv benv) →
      (hpenv : HPredicateEnv fsmEnv penv) →
        p.toBV benv nenv ienv penv tenv  ↔ (fsm.toFsmZext.eval fsmEnv = .negOne)

Term.Ctx

abbrev Term.Ctx (wcard : Nat) (tcard : Nat) : Type :=
  Fin tcard → WidthExpr wcard

Term.BoolEnv

def Term.BoolEnv (bcard : Nat) : Type := Fin bcard → Bool

Term.IntEnv

def Term.IntEnv (icard : Nat) : Type := Fin icard → Nat

HTermFSMToBitStream

structure HTermFSMToBitStream {w : WidthExpr wcard}
  {tctx : Term.Ctx wcard tcard}
  {t : Term bcard ncard icard pcard tctx (.bv w)}
  (fsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm t)) : Prop where
  heq :
    ∀ {wenv : WidthExpr.Env wcard}
      (benv : Term.BoolEnv bcard)
      (nenv : Term.NatEnv ncard)
      (ienv : Term.IntEnv icard)
      (penv : Predicate.Env pcard) (tenv : tctx.Env wenv)
      (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
      (henv : HTermEnv fsmEnv tenv benv) →
        fsm.toFsmZext.eval fsmEnv =
        BitStream.ofBitVecZext (t.toBV benv nenv ienv penv tenv)

Predicate.Env

def Predicate.Env (pcard : Nat) : Type :=
  Fin pcard → Prop

TermKind

inductive TermKind (wcard : Nat) : Type
| bool
| bv (w : WidthExpr wcard)  : TermKind wcard
| prop
| nat
| int

HTermEnv

structure HTermEnv {wcard tcard bcard : Nat}
    {wenv : Fin wcard → Nat} {tctx : Term.Ctx wcard tcard}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (tenv : tctx.Env wenv)
    (benv : Term.BoolEnv bcard) : Prop
  extends HWidthEnv fsmEnv wenv where
    heq_term : ∀ (v : Fin tcard),
      fsmEnv (StateSpace.termVar v) = BitStream.ofBitVecZext (tenv v)
    heq_bool : ∀ (v : Fin bcard),
      fsmEnv (StateSpace.boolVar v) = BitStream.ofBool (benv v)

BitStream.ofBool

noncomputable def BitStream.ofBool (b : Bool) : BitStream := fun _i => b

Term.Ctx.Env

abbrev Term.Ctx.Env
  (tctx : Term.Ctx wcard tcard)
  (wenv : WidthExpr.Env wcard) :=
  (v : Fin tcard) → BitVec ((tctx v).toNat wenv)

ofBitVecZext

abbrev ofBitVecZext {w} (x : BitVec w) : BitStream :=
  fun i => x.getLsbD i

Term.NatEnv

def Term.NatEnv (ncard : Nat) : Type := Fin ncard → Nat

Term.toBV

def Term.toBV {wenv : WidthExpr.Env wcard}
    {tctx : Term.Ctx wcard tcard}
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (penv : Predicate.Env pcard)
    (tenv : tctx.Env wenv)
    (t : Term bcard ncard icard pcard tctx k) : k.denote wenv :=
match t with
| .ofNat w n => BitVec.ofNat (w.toNat wenv) n
| .boolConst b => b
| .var v => tenv.get v.1 v.2
| .add (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a + b
| .zext a v => (a.toBV benv nenv ienv penv tenv).zeroExtend (v.toNat wenv)
| .setWidth a v => (a.toBV benv nenv ienv penv tenv).zeroExtend (v.toNat wenv)
| .sext a v => (a.toBV benv nenv ienv penv tenv).signExtend (v.toNat wenv)
| .bor a b (w := w) =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a ||| b
| .band (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a &&& b
| .bxor (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a ^^^ b
| .bnot (w := w) a =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    ~~~ a
| .boolVar v => benv v
| .shiftl (w := w) a k =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    a <<< k
| .bvOfBool b => BitVec.ofBool (b.toBV benv nenv ienv penv tenv)

| .binWidthRel rel wa wb =>
  match rel with
  | .eq => wa.toNat wenv = wb.toNat wenv
  | .le => wa.toNat wenv ≤ wb.toNat wenv
| .binRel rel _w a b =>
  match rel with
  | .eq => a.toBV benv nenv ienv penv tenv = b.toBV benv nenv ienv penv tenv
  | .ne => a.toBV benv nenv ienv penv tenv ≠ b.toBV benv nenv ienv penv tenv
  | .ult => (a.toBV benv nenv ienv penv tenv).ult (b.toBV benv nenv ienv penv tenv) = true
  | .ule => (a.toBV benv nenv ienv penv tenv).ule (b.toBV benv nenv ienv penv tenv) = true
  | .slt => (a.toBV benv nenv ienv penv tenv).slt (b.toBV benv nenv ienv penv tenv) = true
  | .sle => (a.toBV benv nenv ienv penv tenv).sle (b.toBV benv nenv ienv penv tenv) = true
| .and p1 p2 => p1.toBV benv nenv ienv penv tenv  ∧ p2.toBV benv nenv ienv penv tenv
| .or p1 p2 => p1.toBV benv nenv ienv penv tenv ∨ p2.toBV benv nenv ienv penv tenv
| .boolBinRel rel a b =>
  match rel with

  | .eq => (a.toBV benv nenv ienv penv tenv) = (b.toBV benv nenv ienv penv tenv)
| .pvar v => penv v

Term

inductive Term {wcard tcard : Nat} (bcard : Nat) (ncard : Nat) (icard : Nat) (pcard : Nat)
  (tctx : Term.Ctx wcard tcard) : TermKind wcard → Type

| ofNat (w : WidthExpr wcard) (n : Nat) : Term bcard ncard icard pcard tctx (.bv w)

| var (v : Fin tcard) : Term bcard ncard icard pcard tctx (.bv (tctx v))

| add (a : Term bcard ncard icard pcard tctx (.bv w))
  (b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| shiftl (a : Term bcard ncard icard pcard tctx (.bv w)) (k : Nat) : Term bcard ncard icard pcard tctx (.bv w)

| bor (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| band (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| bxor (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| bnot (a : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| zext (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| setWidth (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| sext (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| bvOfBool (b : Term bcard ncard icard pcard tctx .bool) : Term bcard ncard icard pcard tctx (.bv (.const 1))

| boolConst (b : Bool) : Term bcard ncard icard pcard tctx .bool
| boolVar (v : Fin bcard) : Term bcard ncard icard pcard tctx .bool
| binWidthRel (k : WidthBinaryRelationKind) (wa wb : WidthExpr wcard) :
    Term bcard ncard icard pcard tctx .prop
| binRel
    (k : BinaryRelationKind)
    (w : WidthExpr wcard)
    (a : Term bcard ncard icard pcard tctx (.bv w))
    (b : Term bcard ncard icard pcard tctx (.bv w)) :
    Term bcard ncard icard pcard tctx .prop
| and (p1 p2 : Term bcard ncard icard pcard tctx (.prop)) : Term bcard ncard icard pcard tctx (.prop)
| or (p1 p2 : Term bcard ncard icard pcard tctx (.prop)) : Term bcard ncard icard pcard tctx (.prop)
| pvar (v : Fin pcard) : Term bcard ncard icard pcard tctx (.prop)

| boolBinRel
  (k : BoolBinaryRelationKind)
  (a b : Term bcard ncard icard pcard tctx .bool) :
  Term bcard ncard icard pcard tctx (.prop)

Term.Ctx.Env.get

def Term.Ctx.Env.get {tcard : Nat}
  {wcard : Nat} {wenv : Fin wcard → Nat}
  {tctx : Term.Ctx wcard tcard}
  (tenv : tctx.Env wenv) (i : Nat) (hi : i < tcard) :
  BitVec ((tctx ⟨i, hi⟩).toNat wenv) :=
  tenv ⟨i, hi⟩

BinaryRelationKind

inductive BinaryRelationKind
| eq
| ne
| ule
| slt
| sle
| ult
deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

WidthBinaryRelationKind

inductive WidthBinaryRelationKind
| eq
| le

deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

ofBitVecSext

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

zero

abbrev zero   : BitStream := fun _ => false

carry'

def carry' (initCarry : Bool) (x y : BitStream) : BitStream :=
  fun n =>
  match n with
  | 0 => initCarry
  | n + 1 => (addAux' initCarry x y n).2

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)

protected theorem carry'_eq_carry (x y : BitStream) (c : Bool)
    (x' y' : BitVec w)
    (hx : ∀ i, i < n → x'.getLsbD i = x i)
    (hy : ∀ i,  i < n → y'.getLsbD i = y i) :
    carry' c x y n = (BitVec.carry n x' y' c)

@[simp] theorem carry'_succ (initCarry : Bool) (x y : BitStream) :
    (carry' initCarry x y (i + 1)) =
    let out

MultiWidth.fsmTermSle

@[simp]
theorem eval_fsmMsb_eq {wcard bcard tcard : Nat}
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    {tctx : Term.Ctx wcard tcard}
    (tenv : Term.Ctx.Env tctx wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (w : WidthExpr wcard)
    (x : Term bcard ncard icard pcard tctx (.bv w))
    (xfsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm x))
    (hxfsm : HTermFSMToBitStream xfsm)
    (wfsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    (hwfsm : HNatFSMToBitstream wfsm)
    (htenv : HTermEnv fsmEnv tenv benv) :
    (fsmMsb xfsm.toFsmZext wfsm.toFsm).eval fsmEnv = (fun i =>
      BitStream.ofBitVecZext (x.toBV benv nenv ienv penv tenv) (min i (w.toNat wenv - 1)))

theorem eval_fsmMsb_eq_BitStream_ofBitVecSext {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    {tctx : Term.Ctx wcard tcard}
    (tenv : Term.Ctx.Env tctx wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (w : WidthExpr wcard)
    (x : Term bcard ncard icard pcard tctx (.bv w))
    (xfsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm x))
    (hxfsm : HTermFSMToBitStream xfsm)
    (wfsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    (hwfsm : HNatFSMToBitstream wfsm)
    (htenv : HTermEnv fsmEnv tenv benv) :
    (fsmMsb xfsm.toFsmZext wfsm.toFsm).eval fsmEnv =
      BitStream.ofBitVecSext (x.toBV benv nenv ienv penv tenv)

import Blase.Fast.FiniteStateMachine

import Blase.Vars

import Blase.MultiWidth.Defs

import Blase.KInduction.KInduction

import Lean

namespace MultiWidth

def fsmMsb (x w : FSM α) : FSM α :=
  composeBinaryAux'
    (FSM.latchImmediate false)
    (qfalse := x)
    (qtrue := w)

def fsmCarry'' (initialCarryVal : Bool): FSM Bool :=
  let outputCirc :=
    let carry := Circuit.var true (Sum.inl ())
    let a := Circuit.var true (Sum.inr true)
    let b := Circuit.var true (Sum.inr false)

    ((a &&& b) ||| (a &&& carry) ||| (b &&& carry))
  { α := Unit,

    initCarry := fun () => initialCarryVal,
    outputCirc := Circuit.var true (Sum.inl ()) ,
    nextStateCirc := fun () => outputCirc
  }

def fsmMsbEq (a : FSM α) (b : FSM α) : FSM α :=
  composeUnaryAux (FSM.ls false) <|
    composeBinaryAux' FSM.xor a b

def fsmTermSle
  {wcard tcard : Nat}
  {w : Nondep.WidthExpr}
  {a b : Nondep.Term}
  (wfsm : NatFSM wcard tcard bcard ncard icard pcard w)
  (afsm : TermFSM wcard tcard bcard ncard icard pcard a)
  (bfsm : TermFSM wcard tcard bcard ncard icard pcard b)
  : FSM (StateSpace wcard tcard bcard ncard icard pcard) :=
    let afsm := afsm.toFsmZext
    let bfsm := bfsm.toFsmZext
    let afsm := fsmMsb afsm wfsm.toFsm
    let bfsm := fsmMsb bfsm wfsm.toFsm
    let carryFsm :=
      ((composeBinaryAux' (fsmCarry'' true)  bfsm (~~~ afsm)))
    let xorFsm := fsmMsbEq afsm bfsm
    ~~~ ((~~~ xorFsm) ^^^ carryFsm)

theorem eval_fsmTermSle_eq_decide_sle {wcard tcard bcard : Nat}
    (tctx : Term.Ctx wcard tcard)
    {wenv : WidthExpr.Env wcard}
    (tenv : tctx.Env wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (w : WidthExpr wcard)
    (wfsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    (hwfsm : HNatFSMToBitstream wfsm)
    (a : Term bcard ncard icard pcard tctx (.bv w))
    (b : Term bcard ncard icard pcard tctx (.bv w))
    (afsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm a))
    (hafsm : HTermFSMToBitStream afsm)
    (bfsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm b))
    (hbfsm : HTermFSMToBitStream bfsm)
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (henv : HTermEnv fsmEnv tenv benv)
    :
    ((fsmTermSle
      wfsm
      afsm
      bfsm)).eval fsmEnv i =
       decide (((a.toBV benv nenv ienv penv tenv).signExtend i).sle
       ((b.toBV benv nenv ienv penv tenv).signExtend i)) :=

:= by
  have := BitVec.sle_eq_carry
    (x := (a.toBV benv nenv ienv penv tenv).signExtend i)
    (y := (b.toBV benv nenv ienv penv tenv).signExtend i)
  rw [this]
  clear this
  simp [fsmTermSle]
  simp [eval_fsmMsb_eq_BitStream_ofBitVecSext (hxfsm := hafsm) (hwfsm := hwfsm)
        (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv) (tenv := tenv) (htenv := henv)]
  simp [eval_fsmMsb_eq_BitStream_ofBitVecSext (hxfsm := hbfsm) (hwfsm := hwfsm)
    (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv) (tenv := tenv) (htenv := henv)]
  rw [BitStream.carry'_eq_carry
      (x' := BitVec.signExtend i (Term.toBV benv nenv ienv penv tenv b))
      (y' := ~~~ BitVec.signExtend i (Term.toBV benv nenv ienv penv tenv a))]
  simp [fsmMsbEq]
  · rcases i with rfl | i
    · simp
      simp [BitVec.of_length_zero]
    · simp [eval_fsmMsb_eq_BitStream_ofBitVecSext (hxfsm := hafsm) (hwfsm := hwfsm)
            (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv) (tenv := tenv) (htenv := henv)]
      simp [eval_fsmMsb_eq_BitStream_ofBitVecSext (hxfsm := hbfsm) (hwfsm := hwfsm)
        (benv := benv) (nenv := nenv) (ienv := ienv) (penv := penv) (tenv := tenv) (htenv := henv)]
      simp [BitStream.ofBitVecSext]
      by_cases hi : i < w.toNat wenv
      · simp [BitVec.msb_eq_getLsbD_last, BitVec.getElem_signExtend]
        simp [hi]
        grind [Bool]
      · simp [BitVec.msb_eq_getLsbD_last, BitVec.getElem_signExtend]
        simp [hi]
        grind [Bool]
  · intros j
    intros hj
    simp [hj]
    simp [BitVec.getElem_signExtend]
    simp [BitStream.ofBitVecSext]
    by_cases hw : j < w.toNat wenv
    · simp [hw]
    · simp [hw]
  · intros j
    simp
    intros hj
    simp [hj]
    simp [BitVec.getElem_signExtend]
    simp [BitStream.ofBitVecSext]
    by_cases hw : j < w.toNat wenv
    · simp [hw]
    · simp [hw]


syntax "min" : MLIR.Pretty.uniform_op

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

composeBinaryAux'

def composeBinaryAux'
    (p : FSM Bool)
    (qtrue : FSM α)
    (qfalse : FSM α) :
    FSM α :=
  p.compose (α)
    (λ _ => α)
    (λ _ i => i)
    (λ b => match b with
      | true => qtrue
      | false => qfalse)

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 }

compose

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

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

latchImmediate

def latchImmediate (initVal : Bool) : FSM Bool where
  α := Unit
  initCarry := fun _ => initVal
  outputCirc :=
    let xval := Circuit.var true (inr false)
    let control := Circuit.var true (inr true)
    let state := Circuit.var true (inl ())
    Circuit.ite control xval state
  nextStateCirc := fun () =>
    let xval := Circuit.var true (inr false)
    let control := Circuit.var true (inr true)
    let state := Circuit.var true (inl ())
    Circuit.ite control xval state

ite

def ite (cond t f : Circuit α) : Circuit α :=
  (cond &&& t) ||| (~~~ cond &&& f)

Term

inductive Term
| ofNat (w : WidthExpr) (n : Nat) : Term
| var (v : Nat) (w : WidthExpr) : Term
| add (w : WidthExpr) (a b : Term) : Term
| zext (a : Term) (wnew : WidthExpr) : Term
| setWidth (a : Term) (wnew : WidthExpr) : Term
| sext (a : Term) (wnew : WidthExpr) : Term
| bor (w : WidthExpr) (a b : Term) : Term
| band (w : WidthExpr) (a b : Term) : Term
| bxor (w : WidthExpr) (a b : Term) : Term
| bnot (w : WidthExpr)  (a : Term) : Term
| mul (w : WidthExpr) (a b : Term) : Term
| udiv (w : WidthExpr) (a b : Term) : Term
| umod (w : WidthExpr) (a b : Term) : Term
| boolVar (v : Nat) : Term
| boolConst (b : Bool) : Term
| shiftl (w : WidthExpr) (a : Term) (k : Nat) : Term
| junk (s : String)  : Term
deriving DecidableEq, Inhabited, Repr, Lean.ToExpr

TermFSM

structure TermFSM (wcard tcard bcard ncard icard pcard : Nat) (t : Nondep.Term) where
  toFsmZext : FSM (StateSpace wcard tcard bcard ncard icard pcard)
  width : NatFSM wcard tcard bcard ncard icard pcard t.width

NatFSM

structure NatFSM (wcard tcard bcard ncard icard pcard : Nat) (v : Nondep.WidthExpr) where
  toFsm : FSM (StateSpace wcard tcard bcard ncard icard pcard)

StateSpace

inductive StateSpace (wcard tcard bcard ncard icard pcard : Nat)
| widthVar (v : Fin wcard)
| termVar (v : Fin tcard)
| predVar (v : Fin pcard)
| boolVar (v : Fin bcard)
deriving DecidableEq, Repr, Hashable

Term

inductive Term
| ofNat (w : WidthExpr) (n : Nat) : Term
| var (v : Nat) (w : WidthExpr) : Term
| add (w : WidthExpr) (a b : Term) : Term
| zext (a : Term) (wnew : WidthExpr) : Term
| setWidth (a : Term) (wnew : WidthExpr) : Term
| sext (a : Term) (wnew : WidthExpr) : Term
| bor (w : WidthExpr) (a b : Term) : Term
| band (w : WidthExpr) (a b : Term) : Term
| bxor (w : WidthExpr) (a b : Term) : Term
| bnot (w : WidthExpr)  (a : Term) : Term
| boolVar (v : Nat) : Term
| boolConst (b : Bool) : Term
| shiftl (w : WidthExpr) (a : Term) (k : Nat) : Term
| bvOfBool (b : Term) : Term
| binWidthRel (k : WidthBinaryRelationKind) (wa wb : WidthExpr) : Term
| binRel (k : BinaryRelationKind) (w : WidthExpr)
    (a : Term) (b : Term) : Term
| or (p1 p2 : Term) : Term
| and (p1 p2 : Term) : Term
| pvar (v : Nat) : Term
| boolBinRel (k : BoolBinaryRelationKind)
    (a b : Term) : Term
deriving DecidableEq, Inhabited, Repr, Lean.ToExpr

WidthExpr

inductive WidthExpr where
| const : Nat → WidthExpr
| var : Nat → WidthExpr
| max : WidthExpr → WidthExpr → WidthExpr
| min : WidthExpr → WidthExpr → WidthExpr
| addK : WidthExpr → Nat → WidthExpr
| kadd : Nat → WidthExpr → WidthExpr
deriving Inhabited, Repr, Hashable, DecidableEq, Lean.ToExpr

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

WidthExpr

inductive WidthExpr where
| const : Nat → WidthExpr
| var : Nat → WidthExpr
| max : WidthExpr → WidthExpr → WidthExpr
| min : WidthExpr → WidthExpr → WidthExpr
| addK : WidthExpr → Nat → WidthExpr
| kadd : (k : Nat) → (v : WidthExpr) → WidthExpr
deriving Inhabited, Repr, Hashable, DecidableEq, Lean.ToExpr

HNatFSMToBitstream

structure HNatFSMToBitstream {wcard : Nat} {v : WidthExpr wcard} {tcard : Nat} {bcard : Nat} {pcard : Nat}
   (fsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep v)) : Prop where
  heq :
    ∀ (wenv : Fin wcard → Nat)
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
    (henv : HWidthEnv fsmEnv wenv) →
      fsm.toFsm.eval fsmEnv =
      BitStream.ofNatUnary (v.toNat wenv)

WidthExpr.toNat

def WidthExpr.toNat (e : WidthExpr wcard) (env : WidthExpr.Env wcard) : Nat :=
  match e with
  | .const n => n
  | .var v => env v
  | .min v w => Nat.min (v.toNat env) (w.toNat env)
  | .max v w => Nat.max (v.toNat env) (w.toNat env)
  | .addK v k => v.toNat env + k
  | .kadd k v => k + v.toNat env

WidthExpr

inductive WidthExpr (wcard : Nat) : Type
| const (n : Nat) :  WidthExpr wcard
| var : (v : Fin wcard) → WidthExpr wcard
| min : (v w : WidthExpr wcard) → WidthExpr wcard
| max : (v w : WidthExpr wcard) → WidthExpr wcard
| addK : (v : WidthExpr wcard) → (k : Nat) → WidthExpr wcard
| kadd : (k : Nat) → (v : WidthExpr wcard) → WidthExpr wcard

WidthExpr.Env

abbrev WidthExpr.Env (wcard : Nat) : Type :=
  Fin wcard → Nat

HWidthEnv

structure HWidthEnv {wcard tcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (wenv : Fin wcard → Nat) : Prop where
    heq_width : ∀ (v : Fin wcard),
      fsmEnv (StateSpace.widthVar v) = BitStream.ofNatUnary (wenv v)

HPredicateEnv

structure HPredicateEnv {wcard tcard bcard ncard icard pcard : Nat}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (penv : Fin pcard → Prop) : Prop where
    heq_width : ∀ (v : Fin pcard),
      fsmEnv (StateSpace.predVar v) = BitStream.ofProp (penv v)

BitStream

def BitStream : Type := Nat → Bool

ofNatUnary

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

HPredFSMToBitStream

structure HPredFSMToBitStream {pcard : Nat}
  {tctx : Term.Ctx wcard tcard}
  {p : Term bcard ncard icard pcard tctx .prop}
  (fsm : TermFSM wcard tcard bcard ncard icard pcard
    (.ofDepTerm p)) : Prop where
  heq :
    ∀ {wenv : WidthExpr.Env wcard}
      (benv : Term.BoolEnv bcard)
      (nenv : Term.NatEnv ncard)
      (ienv : Term.IntEnv icard)
      (penv : Predicate.Env pcard) (tenv : tctx.Env wenv)
      (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
      (htenv : HTermEnv fsmEnv tenv benv) →
      (hpenv : HPredicateEnv fsmEnv penv) →
        p.toBV benv nenv ienv penv tenv  ↔ (fsm.toFsmZext.eval fsmEnv = .negOne)

Term.Ctx

abbrev Term.Ctx (wcard : Nat) (tcard : Nat) : Type :=
  Fin tcard → WidthExpr wcard

Term.BoolEnv

def Term.BoolEnv (bcard : Nat) : Type := Fin bcard → Bool

Term.IntEnv

def Term.IntEnv (icard : Nat) : Type := Fin icard → Nat

ofBitVecZext

abbrev ofBitVecZext {w} (x : BitVec w) : BitStream :=
  fun i => x.getLsbD i

HTermFSMToBitStream

structure HTermFSMToBitStream {w : WidthExpr wcard}
  {tctx : Term.Ctx wcard tcard}
  {t : Term bcard ncard icard pcard tctx (.bv w)}
  (fsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm t)) : Prop where
  heq :
    ∀ {wenv : WidthExpr.Env wcard}
      (benv : Term.BoolEnv bcard)
      (nenv : Term.NatEnv ncard)
      (ienv : Term.IntEnv icard)
      (penv : Predicate.Env pcard) (tenv : tctx.Env wenv)
      (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream),
      (henv : HTermEnv fsmEnv tenv benv) →
        fsm.toFsmZext.eval fsmEnv =
        BitStream.ofBitVecZext (t.toBV benv nenv ienv penv tenv)

Predicate.Env

def Predicate.Env (pcard : Nat) : Type :=
  Fin pcard → Prop

TermKind

inductive TermKind (wcard : Nat) : Type
| bool
| bv (w : WidthExpr wcard)  : TermKind wcard
| prop
| nat
| int

HTermEnv

structure HTermEnv {wcard tcard bcard : Nat}
    {wenv : Fin wcard → Nat} {tctx : Term.Ctx wcard tcard}
    (fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream)
    (tenv : tctx.Env wenv)
    (benv : Term.BoolEnv bcard) : Prop
  extends HWidthEnv fsmEnv wenv where
    heq_term : ∀ (v : Fin tcard),
      fsmEnv (StateSpace.termVar v) = BitStream.ofBitVecZext (tenv v)
    heq_bool : ∀ (v : Fin bcard),
      fsmEnv (StateSpace.boolVar v) = BitStream.ofBool (benv v)

BitStream.ofBool

noncomputable def BitStream.ofBool (b : Bool) : BitStream := fun _i => b

Term.Ctx.Env

abbrev Term.Ctx.Env
  (tctx : Term.Ctx wcard tcard)
  (wenv : WidthExpr.Env wcard) :=
  (v : Fin tcard) → BitVec ((tctx v).toNat wenv)

Term.NatEnv

def Term.NatEnv (ncard : Nat) : Type := Fin ncard → Nat

Term.toBV

def Term.toBV {wenv : WidthExpr.Env wcard}
    {tctx : Term.Ctx wcard tcard}
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (penv : Predicate.Env pcard)
    (tenv : tctx.Env wenv)
    (t : Term bcard ncard icard pcard tctx k) : k.denote wenv :=
match t with
| .ofNat w n => BitVec.ofNat (w.toNat wenv) n
| .boolConst b => b
| .var v => tenv.get v.1 v.2
| .add (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a + b
| .zext a v => (a.toBV benv nenv ienv penv tenv).zeroExtend (v.toNat wenv)
| .setWidth a v => (a.toBV benv nenv ienv penv tenv).zeroExtend (v.toNat wenv)
| .sext a v => (a.toBV benv nenv ienv penv tenv).signExtend (v.toNat wenv)
| .bor a b (w := w) =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a ||| b
| .band (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a &&& b
| .bxor (w := w) a b =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    let b : BitVec (w.toNat wenv) := (b.toBV benv nenv ienv penv tenv)
    a ^^^ b
| .bnot (w := w) a =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    ~~~ a
| .boolVar v => benv v
| .shiftl (w := w) a k =>
    let a : BitVec (w.toNat wenv) := (a.toBV benv nenv ienv penv tenv)
    a <<< k
| .bvOfBool b => BitVec.ofBool (b.toBV benv nenv ienv penv tenv)

| .binWidthRel rel wa wb =>
  match rel with
  | .eq => wa.toNat wenv = wb.toNat wenv
  | .le => wa.toNat wenv ≤ wb.toNat wenv
| .binRel rel _w a b =>
  match rel with
  | .eq => a.toBV benv nenv ienv penv tenv = b.toBV benv nenv ienv penv tenv
  | .ne => a.toBV benv nenv ienv penv tenv ≠ b.toBV benv nenv ienv penv tenv
  | .ult => (a.toBV benv nenv ienv penv tenv).ult (b.toBV benv nenv ienv penv tenv) = true
  | .ule => (a.toBV benv nenv ienv penv tenv).ule (b.toBV benv nenv ienv penv tenv) = true
  | .slt => (a.toBV benv nenv ienv penv tenv).slt (b.toBV benv nenv ienv penv tenv) = true
  | .sle => (a.toBV benv nenv ienv penv tenv).sle (b.toBV benv nenv ienv penv tenv) = true
| .and p1 p2 => p1.toBV benv nenv ienv penv tenv  ∧ p2.toBV benv nenv ienv penv tenv
| .or p1 p2 => p1.toBV benv nenv ienv penv tenv ∨ p2.toBV benv nenv ienv penv tenv
| .boolBinRel rel a b =>
  match rel with

  | .eq => (a.toBV benv nenv ienv penv tenv) = (b.toBV benv nenv ienv penv tenv)
| .pvar v => penv v

Term

inductive Term {wcard tcard : Nat} (bcard : Nat) (ncard : Nat) (icard : Nat) (pcard : Nat)
  (tctx : Term.Ctx wcard tcard) : TermKind wcard → Type

| ofNat (w : WidthExpr wcard) (n : Nat) : Term bcard ncard icard pcard tctx (.bv w)

| var (v : Fin tcard) : Term bcard ncard icard pcard tctx (.bv (tctx v))

| add (a : Term bcard ncard icard pcard tctx (.bv w))
  (b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| shiftl (a : Term bcard ncard icard pcard tctx (.bv w)) (k : Nat) : Term bcard ncard icard pcard tctx (.bv w)

| bor (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| band (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| bxor (a b : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| bnot (a : Term bcard ncard icard pcard tctx (.bv w)) : Term bcard ncard icard pcard tctx (.bv w)

| zext (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| setWidth (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| sext (a : Term bcard ncard icard pcard tctx (.bv w)) (v : WidthExpr wcard) : Term bcard ncard icard pcard tctx (.bv v)

| bvOfBool (b : Term bcard ncard icard pcard tctx .bool) : Term bcard ncard icard pcard tctx (.bv (.const 1))

| boolConst (b : Bool) : Term bcard ncard icard pcard tctx .bool
| boolVar (v : Fin bcard) : Term bcard ncard icard pcard tctx .bool
| binWidthRel (k : WidthBinaryRelationKind) (wa wb : WidthExpr wcard) :
    Term bcard ncard icard pcard tctx .prop
| binRel
    (k : BinaryRelationKind)
    (w : WidthExpr wcard)
    (a : Term bcard ncard icard pcard tctx (.bv w))
    (b : Term bcard ncard icard pcard tctx (.bv w)) :
    Term bcard ncard icard pcard tctx .prop
| and (p1 p2 : Term bcard ncard icard pcard tctx (.prop)) : Term bcard ncard icard pcard tctx (.prop)
| or (p1 p2 : Term bcard ncard icard pcard tctx (.prop)) : Term bcard ncard icard pcard tctx (.prop)
| pvar (v : Fin pcard) : Term bcard ncard icard pcard tctx (.prop)

| boolBinRel
  (k : BoolBinaryRelationKind)
  (a b : Term bcard ncard icard pcard tctx .bool) :
  Term bcard ncard icard pcard tctx (.prop)

Term.Ctx.Env.get

def Term.Ctx.Env.get {tcard : Nat}
  {wcard : Nat} {wenv : Fin wcard → Nat}
  {tctx : Term.Ctx wcard tcard}
  (tenv : tctx.Env wenv) (i : Nat) (hi : i < tcard) :
  BitVec ((tctx ⟨i, hi⟩).toNat wenv) :=
  tenv ⟨i, hi⟩

BinaryRelationKind

inductive BinaryRelationKind
| eq
| ne
| ule
| slt
| sle
| ult
deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

WidthBinaryRelationKind

inductive WidthBinaryRelationKind
| eq
| le

deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

eval

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

nextBit

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

State

abbrev State : Type := p.α → Bool

carry

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

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)

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

and

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

@[simp] theorem and_eq : (x &&& y) i = (x i && y i)

@[simp]
theorem FSM.eval_and' (a b : FSM arity) : (a &&& b).eval env = a.eval env &&& b.eval env

theorem FSM.and_eq (a b : FSM arity) : (a &&& b) = composeBinaryAux' FSM.and a b

MultiWidth.fsmSext

@[simp]
theorem eval_fsmMsb_eq {wcard bcard tcard : Nat}
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    {tctx : Term.Ctx wcard tcard}
    (tenv : Term.Ctx.Env tctx wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (w : WidthExpr wcard)
    (x : Term bcard ncard icard pcard tctx (.bv w))
    (xfsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm x))
    (hxfsm : HTermFSMToBitStream xfsm)
    (wfsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep w))
    (hwfsm : HNatFSMToBitstream wfsm)
    (htenv : HTermEnv fsmEnv tenv benv) :
    (fsmMsb xfsm.toFsmZext wfsm.toFsm).eval fsmEnv = (fun i =>
      BitStream.ofBitVecZext (x.toBV benv nenv ienv penv tenv) (min i (w.toNat wenv - 1)))

import Blase.Fast.FiniteStateMachine

import Blase.Vars

import Blase.MultiWidth.Defs

import Blase.KInduction.KInduction

import Lean

namespace MultiWidth

def fsmMsb (x w : FSM α) : FSM α :=
  composeBinaryAux'
    (FSM.latchImmediate false)
    (qfalse := x)
    (qtrue := w)

def fsmSext (x wold wnew : FSM α) : FSM α :=
  (fsmMsb x wold) &&& wnew

theorem fsmSext_eval_eq
    (woldFsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep wold))
    (wnewFsm : NatFSM wcard tcard bcard ncard icard pcard (.ofDep wnew))
    {wenv : WidthExpr.Env wcard}
    {fsmEnv : StateSpace wcard tcard bcard ncard icard pcard → BitStream}
    (hwnew : HNatFSMToBitstream wnewFsm)
    (hwold : HNatFSMToBitstream woldFsm)
    {tctx : Term.Ctx wcard tcard}
    (tenv : Term.Ctx.Env tctx wenv)
    (benv : Term.BoolEnv bcard)
    (nenv : Term.NatEnv ncard)
    (ienv : Term.IntEnv icard)
    (t : Term bcard ncard icard pcard tctx (.bv wold))
    (tFsm : TermFSM wcard tcard bcard ncard icard pcard (.ofDepTerm t))
    (htfsm : HTermFSMToBitStream tFsm)
    (htenv : HTermEnv fsmEnv tenv benv) :
    (fsmSext tFsm.toFsmZext woldFsm.toFsm  wnewFsm.toFsm).eval fsmEnv = fun i =>
      ((BitStream.ofBitVecZext ((Term.sext t wnew).toBV benv nenv ienv penv tenv))) i :=

:= by
  ext i
  rw [fsmSext]
  simp [FSM.eval_and', BitStream.and_eq]
  rw [hwnew.heq (henv := htenv.toHWidthEnv)]
  rw [eval_fsmMsb_eq
        (xfsm := tFsm) (wfsm := woldFsm) (htenv := htenv)
        (hxfsm := htfsm) (hwfsm := hwold)]
  simp
  by_cases hwold : i < wold.toNat wenv
  · simp [hwold]
    by_cases hwnew : i < wnew.toNat wenv
    · simp [hwnew]
      simp [BitVec.getElem_signExtend]
      simp [hwold]
      congr; omega
    · simp [hwnew]
      simp [BitVec.getLsbD_signExtend]
      omega
  · by_cases hwnew : i < wnew.toNat wenv
    · simp [hwnew]
      simp at hwold
      rw [BitVec.getElem_signExtend]
      simp [show min i (wold.toNat wenv - 1) = wold.toNat wenv - 1 by omega]
      simp [show ¬ i < wold.toNat wenv by omega]
      rw [BitVec.msb_eq_getLsbD_last]
    · simp [hwnew]
      rw [BitVec.getLsbD_signExtend]
      simp; omega


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

R.repLength

theorem f_deg_eq : (f q n).degree = 2^n

theorem f_monic : Monic (f q n)

import Mathlib.RingTheory.Polynomial.Quotient

import Mathlib.RingTheory.Ideal.Defs

import Mathlib.RingTheory.Ideal.Basic

import Mathlib.Data.ZMod.Defs

import Mathlib.Data.ZMod.Basic

import Mathlib.Algebra.MonoidAlgebra.Basic

import Mathlib.Algebra.Polynomial.RingDivision

import Mathlib.Data.Finset.Sort

import Mathlib.Data.List.ToFinsupp

import Mathlib.Data.List.Basic

import Mathlib.Tactic.Cases

import LeanMLIR.Framework

open Polynomial -- for R[X] notation

section CommRing

variable (q t : Nat) [Fact (q > 1)] (n : Nat)

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

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

end CommRing

section Representative

variable (q t n : Nat)

private noncomputable def R.representative' :
    R q n → (ZMod q)[X] := Function.surjInv (R.surjective_fromPoly q n)

noncomputable def R.representative :
    R q n → (ZMod q)[X] := fun x => R.representative' q n x %ₘ (f q n)

variable [Fact (q > 1)]

end Representative

noncomputable def R.repLength {q n} (a : R q n) : Nat := match
  Polynomial.degree a.representative with
    | none => 0
    | some d => d + 1

theorem R.repLength_lt_n_plus_1 [Fact (q > 1)]: forall a : R q n, a.repLength < 2^n + 1 :=

:= by
  intro a
  simp only [R.repLength, representative]
  have : Polynomial.degree ( R.representative' q n a %ₘ f q n) < 2^n := by
    rw [← f_deg_eq q n]
    apply (Polynomial.degree_modByMonic_lt)
    apply f_monic
  /- simp only [LT.lt] at this -/
  /- let ⟨val, VAL, VAL_EQN⟩ := this -/
  rcases H : degree (R.representative' q n a %ₘ f q n) <;> simp
  case some val' =>
    rw [H] at this
    rw [WithBot.lt_def] at this
    simp at this
    obtain ⟨val'', twon, hval', heq1, heq2⟩ := this
    rw [Option.some_inj.mp heq1]
    have := Option.some_inj.mp heq2
    subst this
    simp_all


lemma varsFinset_assignVars [DecidableEq α] [DecidableEq β] :
    ∀ (c : Circuit α) (f : ∀ (a : α) (_ha : a ∈ c.vars), β ⊕ Bool),
      (c.assignVars f).varsFinset ⊆ c.varsFinset.biUnion
        (fun a => if ha : a ∈ c.vars
                  then
                    match f a ha with
                    | Sum.inl b => {b}
                    | Sum.inr _ => ∅
                  else ∅)
  | tru, _ => by simp [assignVars, varsFinset, vars]
  | fals, _ => by simp [vars, assignVars, varsFinset]
  | var c v, f => by
    intro x
    simp [assignVars, varsFinset, vars]
    split <;>
    simp [*, vars]
    split_ifs <;> simp [vars]
  | and c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_and _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
  | or c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_or _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
  | xor c₁ c₂, f => by
    intro x
    simp only [assignVars, Finset.mem_biUnion]
    intro hx
    replace hx := varsFinset_xor _ _ hx
    simp only [Finset.mem_union] at hx
    cases hx with
    | inl hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2
    | inr hx =>
      have := varsFinset_assignVars _ _ hx
      simp only [Finset.mem_biUnion] at this
      rcases this with ⟨a, ha⟩
      use a
      simp only [mem_varsFinset] at ha
      simpa [ha.1, mem_varsFinset, vars] using ha.2

import Mathlib.Data.Fin.Basic
import Mathlib.Data.Fintype.Basic
import Mathlib.Data.Finset.Union
import Mathlib.Data.Finset.Defs
import Mathlib.Data.Finset.Card
import Mathlib.Data.List.Pi
import Mathlib.Data.Finset.Basic

Circuit.assignVars

lemma mem_varsFinset [DecidableEq α] {c : Circuit α} :
    ∀ {x : α}, x ∈ c.varsFinset ↔ x ∈ c.vars

theorem varsFinset_and [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (c₁ &&& c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

theorem varsFinset_or [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (c₁ ||| c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

@[simp]
theorem simplifyNot_eq_complement (c : Circuit α) :
    simplifyNot c = ~~~ c

theorem vars_simplifyXor [DecidableEq α] (c₁ c₂ : Circuit α) :
    (vars (simplifyXor c₁ c₂)) ⊆ (vars c₁ ++ vars c₂).dedup

theorem varsFinset_simplifyXor [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (simplifyXor c₁ c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

theorem varsFinset_xor [DecidableEq α] (c₁ c₂ : Circuit α) :
    (varsFinset (c₁ ^^^ c₂)) ⊆ (varsFinset c₁ ∪ varsFinset c₂)

import Mathlib.Data.Finset.Card

import Mathlib.Data.List.Pi

import Mathlib.Data.Finset.Union

import Mathlib.Data.Fin.Basic

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Finset.Defs

import Mathlib.Data.Fintype.Basic

open Std Sat AIG

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

open Lean in

namespace Circuit

variable {α : Type u} {β : Type v}

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

def varsFinset [DecidableEq α] (c : Circuit α) : Finset α :=
  ⟨c.vars, nodup_vars c⟩

def simplifyAnd : Circuit α → Circuit α → Circuit α
  | tru, c => c
  | c, tru => c
  | fals, _ => fals
  | _, fals => fals
  | c₁, c₂ => and c₁ c₂

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

def simplifyNot : Circuit α → Circuit α
  | tru => fals
  | fals => tru
  | xor a b => xor (simplifyNot a) b
  | and a b => or (simplifyNot a) (simplifyNot b)
  | or a b => and (simplifyNot a) (simplifyNot b)
  | var b a => var (!b) a

@[simp]
def simplifyXor : Circuit α → Circuit α → Circuit α
  | fals, c => c
  | c, fals => c
  | tru, c => ~~~ c
  | c, tru => ~~~ c
  | c₁, c₂ => xor c₁ c₂

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

lemma varsFinset_assignVars [DecidableEq α] [DecidableEq β] :
    ∀ (c : Circuit α) (f : ∀ (a : α) (_ha : a ∈ c.vars), β ⊕ Bool),
      (c.assignVars f).varsFinset ⊆ c.varsFinset.biUnion
        (fun a => if ha : a ∈ c.vars
                  then
                    match f a ha with
                    | Sum.inl b => {b}
                    | Sum.inr _ => ∅
                  else ∅) :=


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

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

RelationOrdering

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

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

lemma BitVec.ule_iff_ult_or_eq {w : ℕ} (bv1 bv2 : BitVec w):
    (bv1.ule bv2) = true ↔ (bv1.ult bv2) = true ∨ bv1 = bv2

@[simp]
lemma ucmp_tricho {bv1 bv2 : BitVec w} : (bv2.ult bv1) = false → (bv1.ult bv2) = false → 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.unsignedCmpState : Type where
| eq | gt | lt
deriving Fintype, DecidableEq

def NFA.unsignedCmpStep (q : NFA.unsignedCmpState) (a : BitVec 2) : List NFA.unsignedCmpState :=
  match q, a with
  | .eq, 0 => [ .eq ] | .eq, 3 => [ .eq ] | .eq, 1 => [ .gt ] | .eq, 2 => [ .lt ]
  | .gt, 0 => [ .gt ] | .gt, 1 => [ .gt ] | .gt, 3 => [ .gt ] | .gt, 2 => [ .lt ]
  | .lt, 0 => [ .lt ] | .lt, 1 => [ .gt ] | .lt, 2 => [ .lt ] | .lt, 3 => [ .lt ]

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

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

def RelationOrdering.urel (cmp : RelationOrdering) : BVRel :=
  match cmp with
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .le => fun _ bv1 bv2 => bv1.ule bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1
  | .ge => fun _ bv1 bv2 => bv2.ule bv1

def NFA'.autUnsignedCmpSA (q : NFA.unsignedCmpState) : BVRel :=
  match q with
  | .eq => fun _ bv1 bv2 => bv1 = bv2
  | .lt => fun _ bv1 bv2 => bv1.ult bv2
  | .gt => fun _ bv1 bv2 => bv2.ult bv1

lemma NFA'.autUnsignedCmp_correct cmp : autUnsignedCmp cmp |>.correct2 autUnsignedCmpSA cmp.urel :=

:= by
  let getState {w} (bv1 bv2 : BitVec w) : NFA.unsignedCmpState :=
    if bv2.ult bv1 then .gt else if bv1.ult bv2 then .lt else .eq
  constructor <;> simp [NFA.autUnsignedCmp, autUnsignedCmp, autUnsignedCmpSA, RelationOrdering.urel]
  · rintro _ _ _; cases cmp <;> simp [BitVec.ule_iff_ult_or_eq]; tauto
  · rintro (_ | _ | _) <;> simp
  · rintro (_ | _ | _) a w bv1 bv2 <;> simp [NFA.stepSet, NFA.unsignedCmpStep]
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [instFinEnumBV]
      · rintro ⟨_, _⟩; use .eq; simp; fin_cases a <;> simp [instFinEnumBV] at * <;> tauto
    · constructor
      · rintro ⟨i, hi⟩; cases i <;> fin_cases a <;> simp_all [instFinEnumBV]
      · rintro _; fin_cases a <;> simp [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 [instFinEnumBV]
      · rintro _; fin_cases a <;> simp [instFinEnumBV] at *
        · use .lt; simp_all
        · use (getState bv1 bv2); simp [getState]; split_ifs <;> simp_all; apply ucmp_tricho <;> assumption
        · use .lt; simp_all


macro "⌜" p:term "⌝" : term => `(LE.pure $p)

MonadNonDet

class MonadNonDet (m : Type u → Type v) where
  pick : (τ : Type u) → [Inhabited τ] → m τ

  pickSuchThat : (τ : Type u) → (p : τ → Prop) → [Findable p] → m τ
  assume : (as : Prop) → [Decidable as] → m PUnit.{u+1}

  rep {α : Type u} : α → (α → m (ForInStep α)) → m α

NonDetT

inductive NonDetT (m : Type u -> Type v) : (α : Type u) -> Type _ where
  | pure {α} (ret : α) : NonDetT m α
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) : NonDetT m α
  | pickCont {α} (τ : Type u) (p : τ -> Prop) [Findable p] (f : τ → NonDetT m α) : NonDetT m α
  | repeatCont {α} {β} (init : β) (f : β -> NonDetT m (ForInStep β)) (cont : β -> NonDetT m α) : NonDetT m α

CCPOBot

class CCPOBot (m : Type u -> Type v)  where
  compBot {α} : m α

wp

def wp (c : m α) (post : α -> l) : l := liftM (n := Cont l) c post

MAlgDet

class MAlgDet (l : outParam (Type v)) [Monad m] [CompleteLattice l] [MAlgOrdered m l] where

  demonic {α ι : Type v} (c : m α) (p : ι -> α -> l) [Nonempty ι] :
    ⨅ i, MAlg.lift c (p i) ≤ MAlg.lift c (fun x => ⨅ i, p i x)

  angelic {α ι : Type v} (c : m α) (p : ι -> α -> l) [Nonempty ι] :
    ⨆ i, MAlg.lift c (p i) ≥ MAlg.lift c (fun x => ⨆ i, p i x)

MAlg.lift

abbrev MAlg.lift {m : Type u -> Type v} {l : Type u} [Monad m] [MAlg m l] :
  {α : Type u} -> m α -> Cont l α := fun x f => μ $ f <$> x

MAlg

class MAlg [Monad m] (l : outParam (Type v)) where
  μ : m l -> l
  pure : ∀ l, μ (pure l) = l
  bind : ∀ {α : Type v} (x : m α) (f g : α -> m l),
    μ ∘ f = μ ∘ g ->
    μ (x >>= f) = μ (x >>= g)

MAlgOrdered

class MAlgOrdered (l : outParam (Type v)) [Monad m] [CompleteLattice l] where
  μ : m l -> l
  μ_ord_pure : ∀ l, μ (pure l) = l
  μ_ord_bind {α : Type v} :
    ∀ (f g : α -> m l), μ ∘ f ≤ μ ∘ g ->
      ∀ x : m α, μ (x >>= f) ≤ μ (x >>= g)

CCPOBotLawful

class CCPOBotLawful (m : Type u -> Type v) [∀ α, Lean.Order.CCPO (m α)] [CCPOBot m] where
  prop {α} : CCPOBot.compBot (m := m) (α := α) = Lean.Order.bot

NonDetT.wp

def NonDetT.wp {l : Type u} {α : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m α -> Cont l α
  | .pure ret => pure ret
  | .vis x f => fun post => _root_.wp x fun a => wp (f a) post
  | .pickCont τ p f => fun post => let p : Set τ := p; ⨅ a ∈ (p : Set τ), wp (f a) post

NonDetT.μ

def NonDetT.μ  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] : NonDetT m l -> l := fun x => NonDetT.wp x id

NonDetT.bind

def NonDetT.bind (x : NonDetT m α) (f : α → NonDetT m β) : NonDetT m β :=
  match x with
  | pure ret => f ret
  | vis x f' => vis x fun y => bind (f' y) f
  | pickCont τ p f' => pickCont τ p fun t => bind (f' t) f
  | repeatCont init f' cont => repeatCont init f' fun t => bind (cont t) f

WPGen.bind

def WPGen.bind {x : m α} {f : α -> m β} (wpg : WPGen x) (wpgf : ∀ a, WPGen (f a)) :
  WPGen (x >>= f) where
  get := fun post => wpg.get (fun a => (wpgf a).get post)
  prop := by admit /- proof elided -/

private def _root_.Lean.SimpleScopedEnvExtension.get [Inhabited σ] (ext : SimpleScopedEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

Context

structure Context where
  ref         : Syntax

  m           : Syntax

  returnType  : Syntax
  mutableVars : VarSet := {}
  insideFor   : Bool := false

private def _root_.Lean.SimplePersistentEnvExtension.get [Inhabited σ] (ext : SimplePersistentEnvExtension α σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

WPGen

structure WPGen (x : m α) where
  get : Cont l α

  prop : ∀ post, get post <= wp x post

private def _root_.Lean.EnvExtension.get [Inhabited σ] (ext : EnvExtension σ)
  [Monad m] [MonadEnv m] : m σ := do
  return ext.getState (<- getEnv)

LogicLift

class LogicLift (l : outParam (Type u)) ( k : Type u) [CompleteLattice l] [CompleteLattice k] where
  [lift : MonadLift (Cont l) (Cont k)]
  lift_top {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊤) = ⊤
  lift_bot {α : Type u} :
    monadLift (m := Cont l) (n := Cont k) (fun (_ : α -> l) => ⊥) = ⊥

MAlgTotal

class MAlgTotal (m : Type u -> Type v) [Monad m] [∀ α, Lean.Order.CCPO (m α)]
  [CompleteLattice l] [MAlgOrdered m l] where
  bot_lift {α : Type u} (post : α -> l) :
    MAlg.lift (Lean.Order.bot : m α) post <= ⊥

lemma wp_pure (x : α) (post : α -> l) : wp (m := m) (pure x) post = post x

@[simp]
lemma NonDetT.wp_vis {β : Type u} (x : m β) (f : β → NonDetT m α) post :
  _root_.wp (NonDetT.vis x f) post = _root_.wp x fun a => _root_.wp (f a) post

lemma NonDetT.wp_eq_wp {α : Type u} (x : NonDetT m α) (post : α -> l) :
  _root_.wp x post = NonDetT.wp x post

lemma NonDetT.wp_bind  {l : Type u} [CompleteLattice l] [MAlgOrdered m l] [LawfulMonad m] {α β : Type u} (x : NonDetT m α) (f : α -> NonDetT m β)
  (post : β -> l):
  NonDetT.wp (x.bind f) post = NonDetT.wp x (fun x => NonDetT.wp (f x) post)

lemma wp_cons (x : m α) (post post' : α -> l) :
  (∀ y, post y ≤ post' y) ->
  wp x post ≤ wp x post'

@[simp]
lemma NonDetT.wp_pickCont {τ : Type u} p (f : τ → NonDetT m α) post :
  _root_.wp (NonDetT.pickCont τ p f) post = ⨅ a, ⌜p a⌝ ⇨ _root_.wp (f a) post

@[simp]
lemma NonDetT.wp_pickCont {τ : Type u} p (f : τ → NonDetT m α) post :
  _root_.wp (NonDetT.pickCont τ p f) post = ⨆ a, ⌜p a⌝ ⊓ _root_.wp (f a) post

@[simp]
lemma wp_bot :
  wp (bot : m α) = fun _ => (⊥ : l)

lemma wp_bind {β} (x : m α) (f : α -> m β) (post : β -> l) :
    wp (x >>= f) post = wp x (fun x => wp (f x) post)

NonDetT.extractWeak

import Mathlib.Logic.Function.Basic

import Mathlib.Order.CompleteBooleanAlgebra

import Mathlib.Order.Lattice

import Mathlib.Order.Basic

import Mathlib.Data.W.Basic

import Mathlib.Data.FinEnum

import Loom.MonadAlgebras.WP.Gen

import Loom.MonadAlgebras.WP.Liberal

import Loom.MonadAlgebras.NonDetT'.Basic

open Lean.Order

def findNat (p : Nat -> Prop) [DecidablePred p] : Option Nat :=
  let rec aux i :=
    if p i then
      some i
    else
      aux (i + 1)
  partial_fixpoint
  aux 0

def find [Encodable α] (p : α -> Prop) [DecidablePred p] : Option α :=
  findNat (fun x => (Encodable.decode x).any (p ·)) |>.bind Encodable.decode

class WeakFindable {α : Type u} (p : α -> Prop) where
  find : Unit -> Option α
  find_some_p : find () = some x -> p x

instance WeakFindable.of_Findable {α : Type u} (p : α -> Prop) [Findable p] : WeakFindable p where
  find := Findable.find p
  find_some_p := Findable.find_some_p

inductive ExtractNonDet (findable : {τ : Type u} -> (τ -> Prop) -> Type u) {m} : {α : Type u} -> NonDetT m α -> Type _ where
  | pure {α} : ∀ (x : α), ExtractNonDet findable (NonDetT.pure x)
  | vis {α} {β} (x : m β) (f : β → NonDetT m α) :
      (∀ y, ExtractNonDet findable (f y)) → ExtractNonDet findable (.vis x f)
  | pickSuchThat {α} (τ : Type u) (p : τ -> Prop) (f : τ → NonDetT m α)
    {_ : findable p}
     : (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont τ p f)
  | assume {α} (p : PUnit -> Prop) (f : PUnit → NonDetT m α) {_ : Decidable (p .unit)} :
    (∀ x, ExtractNonDet findable (f x)) → ExtractNonDet findable (.pickCont PUnit p f)

instance ExtractNonDet.pure' : ExtractNonDet findable (Pure.pure (f := NonDetT m) x) :=

instance ExtractNonDet.liftM (x : m α) :
  ExtractNonDet findable (liftM (n := NonDetT m) x) :=

instance ExtractNonDet.assume' {p : Prop} [Decidable p] : ExtractNonDet findable (MonadNonDet.assume (m :=  NonDetT m) p) :=

instance ExtractNonDet.pickSuchThat' {τ : Type u} (p : τ -> Prop) [Findable p] :
  ExtractNonDet Findable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

instance ExtractNonDet.pickSuchThat_weak {τ : Type u} (p : τ -> Prop) [WeakFindable p] :
  ExtractNonDet WeakFindable (MonadNonDet.pickSuchThat (m := NonDetT m) τ p) :=

instance ExtractNonDet.if {p : Prop} {dec : Decidable p} {x y : NonDetT m α}
  (_ : ExtractNonDet findable x) (_ : ExtractNonDet findable y) :
  ExtractNonDet findable (if p then x else y) :=

instance ExtractNonDet.ForIn_list {xs : List α} {init : β} {f : α → β → NonDetT m (ForInStep β)}
  (_ : ∀ a b, ExtractNonDet findable (f a b)) :
  ExtractNonDet findable (forIn xs init f) :=

variable [Monad m] [CCPOBot m] [CompleteBooleanAlgebra l] [MAlgOrdered m l] [MAlgDet m l] [LawfulMonad m]

@[simp, inline]
def NonDetT.extractGen {findable : {τ : Type u} -> (τ -> Prop) -> Type u} {α : Type u}
  (findOf : ∀ {τ : Type u} (p : τ -> Prop), findable p -> Unit -> Option τ)
  : (s : NonDetT m α) -> (ex : ExtractNonDet findable s := by admit /- proof elided -/
  ) -> m α
  | .pure x, _ => Pure.pure x
  | .vis x f, .vis _ _ _ => liftM x >>= (fun x => extractGen findOf (f x))
  | .pickCont _ p f, .pickSuchThat _ _ _ _ =>
    match findOf p ‹_› () with
    | none => CCPOBot.compBot
    | some x =>  extractGen findOf (f x)
  | .pickCont _ p f, .assume _ _ _ =>
    if p .unit then
      extractGen findOf (f .unit)
    else CCPOBot.compBot

def NonDetT.extractWeak {α : Type u} (s : NonDetT m α) (ex : ExtractNonDet WeakFindable s := by admit /- proof elided -/
) : m α :=
  NonDetT.extractGen WeakFindable.find s

namespace DemonicChoice

end DemonicChoice

namespace AngelicChoice

variable [∀ α, CCPO (m α)] [CCPOBotLawful m] [MAlgTotal m]

omit [MAlgDet m l] in
lemma ExtractNonDet.extract_refines_wp_weak (s : NonDetT m α) (inst : ExtractNonDet WeakFindable s) :
  wp s.extractWeak post <= wp s post :=

:= by
  unhygienic induction inst
  { simp [wp_pure, NonDetT.extractWeak] }
  { simp only [NonDetT.extractWeak, NonDetT.extractGen, monadLift_self, wp_bind, NonDetT.wp_vis];
    apply wp_cons; aesop (add norm inf_comm) }
  { simp only [NonDetT.extractWeak, NonDetT.extractGen, NonDetT.wp_pickCont]; split
    { simp [*, CCPOBotLawful.prop, TotalCorrectness.wp_bot] }
    apply le_iSup_of_le; simp; constructor; rotate_left
    apply a_ih; rename_i h; simp [x.find_some_p h] }
  simp only [NonDetT.extractWeak, NonDetT.extractGen, NonDetT.wp_pickCont]
  have : ∀ a : PUnit.{u_1 + 1}, a = .unit := by simp
  simp only [this, ge_iff_le]; split_ifs <;> simp [*, iSup_const, CCPOBotLawful.prop, TotalCorrectness.wp_bot]
  apply a_ih