-- Blase/Blase/MultiWidth/GoodFSM.lean

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)

-- Repository definitions used:
content: syntax "min" : MLIR.Pretty.uniform_op

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

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

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

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

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

content: structure FSM (arity : Type) : Type 1 where
   
  ( α  : Type )
  [ i : FinEnum α ]
  [ h : Hashable α ]
  [ dec_eq : DecidableEq α ]
   
  ( initCarry : α → Bool )
   
  outputCirc : Circuit (α ⊕ arity)
  nextStateCirc : α → Circuit (α ⊕ arity)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: def BitStream : Type := Nat → Bool

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: inductive WidthBinaryRelationKind
| eq
| le


deriving DecidableEq, Repr, Inhabited, Lean.ToExpr

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

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

-- Local definitions used:
content: def fsmMsb (x w : FSM α) : FSM α :=
  composeBinaryAux'
    (FSM.latchImmediate false)
    (qfalse := x)
    (qtrue := w)

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

-- Target theorem and proof:
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)
:= by
  rw [eval_fsmMsb_eq (wfsm := wfsm) (hwfsm := hwfsm) (hxfsm := hxfsm)
    (tenv := tenv) (htenv := htenv)]
  ext i
  simp [BitStream.ofBitVecSext]
  by_cases hi : i < w.toNat wenv
  · simp [hi]
    congr
    omega
  · simp [hi, BitVec.msb_eq_getLsbD_last]
    congr
    omega