-- Blase/Blase/Fast/Lemmas.lean

import Mathlib.Data.Fintype.Card

import Mathlib.Data.Fintype.Sum

import Mathlib.Data.Fintype.Sigma

import Mathlib.Data.Fintype.BigOperators

import Mathlib.Tactic.Zify

import Mathlib.Tactic.Ring

import Blase.Fast.Defs

import Blase.Fast.BitStream

open Term

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

syntax "slt" : 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 Predicate.evalUlt (t₁ t₂ : BitStream) : BitStream := (~~~ (t₁.borrow t₂)).concat true

content: def BitStream : Type := Nat → Bool

content: def borrow (x y : BitStream) : BitStream :=
  fun n => (subAux x y n).2

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

content: inductive Predicate : Type where
 
| width (wp : WidthPredicate) (n : Nat) : Predicate
| binary (p : BinaryPredicate) (t₁ t₂ : Term)
| land  (p q : Predicate) : Predicate
| lor (p q : Predicate) : Predicate
deriving Repr, Inhabited, Lean.ToExpr

content: @[simp] def Predicate.evalFin (p : Predicate) (vars : Fin (arity p) → BitStream) : BitStream :=
match p with
| .width .eq n => BitStream.falseIffEq n
| .width .neq n => BitStream.falseIffNeq n
| .width .lt n => BitStream.falseIffLt n
| .width .le n => BitStream.falseIffLe n
| .width .gt n => BitStream.falseIffGt n
| .width .ge n => BitStream.falseIffGe n
| .binary .eq t₁ t₂ =>
    let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    Predicate.evalEq x₁ x₂
| .binary .neq t₁ t₂  =>
    let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
    ) i))
    Predicate.evalNeq x₁ x₂
| .land p q =>
  
  
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLand x₁ x₂
| .lor p q =>
  
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLor x₁ x₂
| .binary .slt p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalSlt x₁ x₂
| .binary .sle p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLor (Predicate.evalSlt x₁ x₂) (Predicate.evalEq x₁ x₂)
| .binary .ult p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  (Predicate.evalUlt x₁ x₂)
| .binary .ule p q =>
  let x₁ := p.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  let x₂ := q.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
  ) i))
  Predicate.evalLor (Predicate.evalUlt x₁ x₂) (Predicate.evalEq x₁ x₂)

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

content: @[simp] def Term.evalFin (t : Term) (vars : Fin (arity t) → BitStream) : BitStream :=
  match t with
  | var n => vars (Fin.last n)
  | zero    => BitStream.zero
  | one     => BitStream.one
  | negOne  => BitStream.negOne
  | ofNat n => BitStream.ofNat n
  | and t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ &&& x₂
  | or t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ||| x₂
  | xor t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ ^^^ x₂
  | not t     => ~~~(t.evalFin vars)
  | add t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ + x₂
  | sub t₁ t₂ =>
      let x₁ := t₁.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      let x₂ := t₂.evalFin (fun i => vars (Fin.castLE (by admit /- proof elided -/
      ) i))
      x₁ - x₂
  | neg t       => -(Term.evalFin t vars)
 
 
  | shiftL t n  => BitStream.shiftLeft (Term.evalFin t vars) n

content: abbrev negOne : BitStream := fun _ => true

content: def shiftLeft (x : BitStream) (k : Nat) : BitStream :=
  fun i => if i < k then false else x (i - k) 

content: def ofNat (x : Nat) : BitStream :=
  Nat.testBit x

content: abbrev one    : BitStream := (· == 0)

content: abbrev zero   : BitStream := fun _ => false

content: @[simp] def Term.arity : Term → Nat
| (var n) => n+1
| zero => 0
| one => 0
| negOne => 0
| ofNat _ => 0
| Term.and t₁ t₂ => max (arity t₁) (arity t₂)
| Term.or t₁ t₂ => max (arity t₁) (arity t₂)
| Term.xor t₁ t₂ => max (arity t₁) (arity t₂)
| Term.not t => arity t
| add t₁ t₂ => max (arity t₁) (arity t₂)
| sub t₁ t₂ => max (arity t₁) (arity t₂)
| neg t => arity t


| shiftL t .. => arity t

content: inductive Term : Type
| var : Nat → Term
 
| zero : Term
 
| negOne : Term
 
| one : Term
 
| ofNat (n : Nat) : Term
 
| and : Term → Term → Term
 
| or : Term → Term → Term
 
| xor : Term → Term → Term
 
| not : Term → Term
 
| add : Term → Term → Term
 
| sub : Term → Term → Term
 
| neg : Term → Term


 
| shiftL : Term → Nat → Term


deriving Repr, Inhabited, Lean.ToExpr

content: def Predicate.evalSlt (t₁ t₂ : BitStream) : BitStream :=
    (((Predicate.evalUlt t₁ t₂)) ^^^ (Predicate.evalMsbEq t₁ t₂))

content: def Predicate.evalMsbEq (t₁ t₂ : BitStream) : BitStream :=
  (t₁ ^^^ t₂).concat false

content: def Predicate.evalLand (t₁ t₂ : BitStream) : BitStream := (t₁ ||| t₂)

content: def Predicate.evalNeq (t₁ t₂ : BitStream) : BitStream := (t₁.nxor t₂).concat true |>.scanAnd

content: def nxor (a b : BitStream) : BitStream := fun i => a i == b i

content: def scanAnd (s : BitStream) : BitStream := scanl true Bool.and s

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

content: def Predicate.evalEq (t₁ t₂ : BitStream) : BitStream := (t₁ ^^^ t₂).concat false |>.scanOr

content: def scanOr (s : BitStream) : BitStream := scanl false Bool.or s

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

content: abbrev falseIffLt (n : Nat) : BitStream := fun i => decide (i ≥ n)

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

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

content: abbrev falseIffEq (n : Nat) : BitStream := fun i => decide (i != n)

content: abbrev falseIffGt (n : Nat) : BitStream := fun i => decide (i ≤ n)

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

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

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

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

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

-- Repository lemmas used:
content: theorem and_comm (x y : BitStream) : x &&& y = y &&& x

-- Local definitions used:

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

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