-- Blase/Blase/Fast/Defs.lean

import Mathlib.Data.Bool.Basic

import Mathlib.Data.Fin.Basic

import Blase.Fast.BitStream

import Blase.SingleWidth.Defs

open Term

open BitStream in

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

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

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

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

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

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

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

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

open BitStream in

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

section Predicate

end Predicate

-- Repository definitions used:
content: 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 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: 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 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 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: abbrev falseIffLe (n : Nat) : BitStream := fun i => decide (i > 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: def Term.denote (w : Nat) (t : Term) (vars : List (BitVec w)) : BitVec w :=
  match t with
  | ofNat n => BitVec.ofNat w n
  | var n => vars.getD n default
  | zero => 0#w
  | negOne => -1#w
  | one  => 1#w
  | and a b => (a.denote w vars) &&& (b.denote w vars)
  | or a b => (a.denote w vars) ||| (b.denote w vars)
  | xor a b => (a.denote w vars) ^^^ (b.denote w vars)
  | not a => ~~~ (a.denote w vars)
  | add a b => (a.denote w vars) + (b.denote w vars)
  | sub a b => (a.denote w vars) - (b.denote w vars)
  | neg a => - (a.denote w vars)
  
  
  | shiftL a n => (a.denote w vars) <<< n

content: def Predicate.denote (p : Predicate) (w : Nat) (vars : List (BitVec w)) : Prop :=
  match p with
  | .width .ge k => k ≤ w 
  | .width .gt k => k < w 
  | .width .le k => w ≤ k
  | .width .lt k => w < k
  | .width .neq k => w ≠ k
  | .width .eq k => w = k
  | .binary .eq t₁ t₂ => t₁.denote w vars = t₂.denote w vars
  | .binary .neq t₁ t₂ => t₁.denote w vars ≠ t₂.denote w vars
  | .binary .sle  t₁ t₂ => ((t₁.denote w vars).sle (t₂.denote w vars)) = true
  | .binary .slt  t₁ t₂ => ((t₁.denote w vars).slt (t₂.denote w vars)) = true
  | .binary .ule  t₁ t₂ => ((t₁.denote w vars).ule (t₂.denote w vars)) = true
  | .binary .ult  t₁ t₂ => (t₁.denote w vars).ult (t₂.denote w vars) = true
  | .land  p q => p.denote w vars ∧ q.denote w vars
  | .lor  p q => p.denote w vars ∨ q.denote w vars

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: inductive WidthPredicate
| eq
| neq
| lt
| le
| gt
| ge
deriving Repr, Inhabited, Lean.ToExpr

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

content: def toBitVec (w : Nat) (x : BitStream) : BitVec w :=
  match w with
  | 0   => 0#0
  | w+1 => (x.toBitVec w).cons (x w)

-- Repository lemmas used:
content: @[simp] theorem subAux_eq_BitVec_carry (a b : BitStream) (w i : Nat) (hi : i < w) :
    (a.subAux b i).2 = !(BitVec.carry (i + 1) (a.toBitVec w) ((~~~b).toBitVec w) true)

-- Local definitions used:
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

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

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

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

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

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

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

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

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

-- Local lemmas used:
content: private theorem BitVec.lt_eq_decide_ult {x y : BitVec w} : (x < y) = decide (x.ult y)

content: theorem Predicate.evalUlt_denote_false_iff {w : Nat} (a b : Term) (vars : List (BitVec w)) :
    evalUlt (a.eval (List.map .ofBitVecSext vars)) (b.eval (List.map .ofBitVecSext vars)) w = false ↔
    (Term.denote w a vars < Term.denote w b vars)

-- Target theorem and proof:
theorem Predicate.evalUlt_denote_true_iff {w : Nat} (a b : Term) (vars : List (BitVec w)) :
    evalUlt (a.eval (List.map .ofBitVecSext vars)) (b.eval (List.map .ofBitVecSext vars)) w = true ↔
      (Term.denote w b vars) ≤ (Term.denote w a vars)
:= by
  obtain ⟨h₁, h₂⟩ := evalUlt_denote_false_iff a b vars
  constructor
  · intros h
    by_contra h'
    simp at h'
    specialize (h₂ h')
    simp [h₂] at h
  · intros h
    by_contra h'
    simp at h'
    specialize (h₁ h')
    bv_omega