-- Blase/Blase/KInduction/KInduction.lean

import Mathlib.Data.Bool.Basic

import Mathlib.Data.Fin.Basic

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Finset.Defs

import Mathlib.Data.Multiset.FinsetOps

import Blase.Fast.BitStream

import Blase.Fast.Defs

import Blase.Fast.FiniteStateMachine

import Blase.Fast.Decide

import Blase.SingleWidth.Syntax

import Lean.Meta.ForEachExpr

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

import Blase.Fast.ForLean

import Blase.Vars

import Blase.EnvBitstream

import Lean

open Fin.NatCast

namespace ReflectVerif

open Lean Meta Elab Tactic

namespace BvDecide

open Std Sat AIG

structure KInductionCircuits {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity] (fsm : FSM arity) (n : Nat) where
  
  cInitCarryAssignCirc : Circuit (Vars fsm.α arity 0)
  
  cSuccCarryAssignCirc : Circuit (Vars fsm.α arity (n+2))
  
  cOutAssignCirc : Circuit (Vars fsm.α arity (n + 2))
  
  cStatesUniqueCirc : Circuit (Vars fsm.α arity n)

namespace KInductionCircuits

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

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

def envBitstream_set (x : arity → BitStream) (n : Nat) (v : arity → Bool) :
    arity → BitStream :=
  fun a j => if j = n then v a else x a j

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

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

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

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

syntax "return" : MLIR.Pretty.uniform_op

content: macro_rules
  
  | `(mlir_op| $v:mlir_op_operand = mod_arith.constant $x:neg_num : $t) =>
    `(mlir_op| $v:mlir_op_operand = "mod_arith.constant" () {value = $x:neg_num} : () -> ($t))

  
  | `(mlir_op| $v:mlir_op_operand = mod_arith.constant ${ $x:term } : $t) => do
      let ctor := mkIdent ``MLIR.AST.AttrValue.int
      let x ← `($ctor $x [mlir_type| i64])
      `(mlir_op| $v:mlir_op_operand = "mod_arith.constant" () {value = $$($x)} : () -> ($t))

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

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

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

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

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

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

content: macro_rules
  | `(mlir_op| $v:mlir_op_operand = arith.const $x:neg_num : $t) =>
      `(mlir_op| $v:mlir_op_operand = "arith.const" () {value = $x:neg_num } : () -> ($t))
  | `(mlir_op| $v:mlir_op_operand = arith.const ${ $x:term } : $t) => do
      let ctor := mkIdent ``MLIR.AST.AttrValue.int
      let x ← `($ctor $x [mlir_type| i64])
      
      
      
      
      `(mlir_op| $v:mlir_op_operand = "arith.const" () {value = $$($x) } : () -> ($t))

content: macro_rules
  | `(mlir_op| $v:mlir_op_operand = poly.const $x:neg_num : $t) =>
      `(mlir_op| $v:mlir_op_operand = "poly.const" () {value = $x:neg_num } : () -> ($t))
  | `(mlir_op| $v:mlir_op_operand = poly.const ${ $x:term } : $t) => do
      let ctor := mkIdent ``MLIR.AST.AttrValue.int
      let x ← `($ctor $x [mlir_type| i64])
      
      
      
      
      `(mlir_op| $v:mlir_op_operand = "poly.const" () {value = $$($x) } : () -> ($t))

content: macro_rules
  | `(mlir_op| $v:mlir_op_operand = poly.monomial $xs,* : ($ts,*) -> $t) =>
      `(mlir_op| $v:mlir_op_operand = "poly.monomial" ($xs,*) : ($ts,*) -> $t)

content: def BitStream : Type := Nat → Bool

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 carryWith (p : FSM arity) (carryState : p.α → Bool) (x : arity → BitStream) (n : Nat) : p.α → Bool := fun a =>
  (p.changeInitCarry carryState).carry x n a

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

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

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

content: abbrev State : Type := p.α → Bool

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

content: def rhs := [poly q, n, hq | {
^bb0(%a : !R):
  return %a : !R

}]

content: inductive AttrValue where
  | symbol: String -> AttrValue 
  | str : String -> AttrValue
  | int : Int -> MLIRType φ -> AttrValue
  | nat: Nat -> AttrValue
  | bool : Bool -> AttrValue
  | float : Float -> MLIRType φ -> AttrValue
  | type : MLIRType φ -> AttrValue
  | affine: AffineMap -> AttrValue
  | permutation: List Nat -> AttrValue 
  | list: List AttrValue -> AttrValue
  | nestedsymbol: AttrValue -> AttrValue -> AttrValue
  | alias: String -> AttrValue
  | dict: AttrDict -> AttrValue
  | opaque_: (dialect: String) -> (value: String) -> AttrValue
  | opaqueElements: (dialect: String) -> (value: String) -> (type: MLIRType φ) -> AttrValue
  | unit: AttrValue

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

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

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

-- Repository lemmas used:
content: theorem carryWith_congrEnv {p : FSM arity}
    {carryState : p.α → Bool} {x y : arity → BitStream} {n : Nat}
    (h : ∀ a i, i < n → x a i = y a i) :
    p.carryWith carryState x n = p.carryWith carryState y n

content: theorem carry_congrEnv {p : FSM arity}
    {x y : arity → BitStream} {n : Nat} (h : ∀ a i, i < n → x a i = y a i) :
    p.carry x n = p.carry y n

-- Local definitions used:
content: structure KInductionCircuits {arity : Type _}
  [DecidableEq arity] [Fintype arity] [Hashable arity] (fsm : FSM arity) (n : Nat) where
  
  cInitCarryAssignCirc : Circuit (Vars fsm.α arity 0)
  
  cSuccCarryAssignCirc : Circuit (Vars fsm.α arity (n+2))
  
  cOutAssignCirc : Circuit (Vars fsm.α arity (n + 2))
  
  cStatesUniqueCirc : Circuit (Vars fsm.α arity n)

content: def envBitstream_set (x : arity → BitStream) (n : Nat) (v : arity → Bool) :
    arity → BitStream :=
  fun a j => if j = n then v a else x a j

-- Local lemmas used:

-- Target theorem and proof:
theorem FSM.carryWith_congrEnv_envBitstream_set_of_le (fsm : FSM arity)
    (s0 : fsm.α → Bool) (env : arity → BitStream) (n : Nat) (v : arity → Bool)
    (k : Nat) (hk : k ≤ n) :
  fsm.carryWith s0 (envBitstream_set env n v) k =
  fsm.carryWith s0 env k
:= by
  apply FSM.carryWith_congrEnv
  intros a l hl
  simp [envBitstream_set]
  intros hcontra; omega