import Hax.Tactic.Init
import Hax.rust_primitives.USize64
import Hax.Tactic.SpecSet
import Hax.MissingLean
import Hax.rust_primitives.RustM
open Std.Do
open Std.Tactic

open Std.Do
set_option mvcgen.warning false

/-
  Integer types are represented as the corresponding type in Lean
-/
abbrev u8 := UInt8
abbrev u16 := UInt16
abbrev u32 := UInt32
abbrev u64 := UInt64
abbrev u128 := UInt128
abbrev usize := USize64
abbrev i8 := Int8
abbrev i16 := Int16
abbrev i32 := Int32
abbrev i64 := Int64
abbrev i128 := Int128
abbrev isize := ISize

abbrev f32 := Float32
abbrev f64 := Float

/-- Class of objects that can be transformed into Nat -/
class ToNat (α: Type) where
  toNat : α -> Nat

attribute [grind] ToNat.toNat

@[simp, grind]
instance : ToNat usize where
  toNat x := x.toNat
@[simp, grind]
instance : ToNat u128 where
  toNat x := x.toNat
@[simp, grind]
instance : ToNat u64 where
  toNat x := x.toNat
@[simp, grind]
instance : ToNat u32 where
  toNat x := x.toNat
@[simp, grind]
instance : ToNat u16 where
  toNat x := x.toNat
@[simp, grind]
instance : ToNat u8 where
  toNat x := x.toNat

infixl:58 " ^^^? " => fun a b => pure (HXor.hXor a b)
infixl:60 " &&&? " => fun a b => pure (HAnd.hAnd a b)
infixl:60 " |||? " => fun a b => pure (HOr.hOr a b)
prefix:75 "~?"     => fun a => pure (~~~a)

/-

## Boolean comparisons

Boolean comparisons that are prettyfied for the integer and boolean types.

-/

namespace rust_primitives.cmp

def eq {α : Type} [BEq α] (a b : α) : RustM Bool := pure (a == b)
def ne {α : Type} [BEq α] (a b : α) : RustM Bool := pure (a != b)
def lt {α : Type} [LT α] [DecidableLT α] (a b : α) : RustM Bool :=
  pure (decide (a < b))
def le {α : Type} [LE α] [DecidableLE α] (a b : α) : RustM Bool :=
  pure (decide (a <= b))
def gt {α : Type} [LT α] [DecidableLT α] (a b : α) : RustM Bool :=
  pure (decide (a > b))
def ge {α : Type} [LE α] [DecidableLE α] (a b : α) : RustM Bool :=
  pure (decide (a >= b))

infixl:80 " ==? "  => rust_primitives.cmp.eq
infixl:80 " !=? "  => rust_primitives.cmp.ne
infixl:80 " <? "   => rust_primitives.cmp.lt
infixl:80 " <=? "  => rust_primitives.cmp.le
infixl:80 " >? "   => rust_primitives.cmp.gt
infixl:80 " >=? "  => rust_primitives.cmp.ge

attribute [spec 100, specset bv, hax_bv_decide]
rust_primitives.cmp.eq
rust_primitives.cmp.ne
rust_primitives.cmp.lt
rust_primitives.cmp.le
rust_primitives.cmp.gt
rust_primitives.cmp.ge

open Lean in
set_option hygiene false in
macro "declare_comparison_specs" s:(&"signed" <|> &"unsigned") typeName:ident width:term : command => do

  let signed ← match s.raw[0].getKind with
  | `signed => pure true
  | `unsigned => pure false
  | _ => throw .unsupportedSyntax

  if signed then
    return ← `(
      namespace $typeName

      @[specset int]
      def eq_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ eq x y ⦃ ⇓ r => ⌜ r = (x.toInt == y.toInt) ⌝ ⦄ := by
        mvcgen [eq]; rw [← @Bool.coe_iff_coe]; simp [x.toInt_inj]

      @[specset int]
      def ne_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ ne x y ⦃ ⇓ r => ⌜ r = (x.toInt != y.toInt) ⌝ ⦄ := by
        mvcgen [ne]; rw [← @Bool.coe_iff_coe]; simp [x.toInt_inj]

      @[specset int]
      def lt_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ lt x y ⦃ ⇓ r => ⌜ r = decide (x.toInt < y.toInt) ⌝ ⦄ := by
        mvcgen [lt]; simp [x.lt_iff_toInt_lt]

      @[specset int]
      def le_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ le x y ⦃ ⇓ r => ⌜ r = decide (x.toInt ≤ y.toInt) ⌝ ⦄ := by
        mvcgen [le]; simp [x.le_iff_toInt_le]

      @[specset int]
      def gt_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ gt x y ⦃ ⇓ r => ⌜ r = decide (x.toInt > y.toInt ) ⌝ ⦄ := by
        mvcgen [gt]; simp [y.lt_iff_toInt_lt]

      @[specset int]
      def ge_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ ge x y ⦃ ⇓ r => ⌜ r = decide (x.toInt ≥ y.toInt) ⌝ ⦄ := by
        mvcgen [ge]; simp [y.le_iff_toInt_le]

      end $typeName
    )
  else return ← `(
      namespace $typeName

      @[specset int]
      def eq_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ eq x y ⦃ ⇓ r => ⌜ r = (x.toNat == y.toNat) ⌝ ⦄ := by
        mvcgen [eq]; rw [← @Bool.coe_iff_coe]; simp [x.toNat_inj]

      @[specset int]
      def ne_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ ne x y ⦃ ⇓ r => ⌜ r = (x.toNat != y.toNat) ⌝ ⦄ := by
        mvcgen [ne]; rw [← @Bool.coe_iff_coe]; simp [x.toNat_inj]

      @[specset int]
      def lt_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ lt x y ⦃ ⇓ r => ⌜ r = decide (x.toNat < y.toNat) ⌝ ⦄ := by
        mvcgen [lt]

      @[specset int]
      def le_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ le x y ⦃ ⇓ r => ⌜ r = decide (x.toNat ≤ y.toNat) ⌝ ⦄ := by
        mvcgen [le]

      @[specset int]
      def gt_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ gt x y ⦃ ⇓ r => ⌜ r = decide (x.toNat > y.toNat ) ⌝ ⦄ := by
        mvcgen [gt]

      @[specset int]
      def ge_spec (x y : $typeName) : ⦃ ⌜ True ⌝ ⦄ ge x y ⦃ ⇓ r => ⌜ r = decide (x.toNat ≥ y.toNat) ⌝ ⦄ := by
        mvcgen [ge]

      end $typeName
  )

declare_comparison_specs signed Int8 8
declare_comparison_specs signed Int16 16
declare_comparison_specs signed Int32 32
declare_comparison_specs signed Int64 64
declare_comparison_specs signed ISize System.Platform.numBits
declare_comparison_specs unsigned UInt8 8
declare_comparison_specs unsigned UInt16 16
declare_comparison_specs unsigned UInt32 32
declare_comparison_specs unsigned UInt64 64
declare_comparison_specs unsigned USize64 64

end rust_primitives.cmp

set_option linter.unusedVariables false in

/-

## Arithmetic operations

The Rust arithmetic operations have their own notations, using a `?`. They
return a `RustM`, that is `.fail` when arithmetic overflows occur.

-/

class rust_primitives.ops.arith.Add (α : Type) where
  add : α → α → RustM α
class rust_primitives.ops.arith.Sub (α : Type) where
  sub : α → α → RustM α
class rust_primitives.ops.arith.Mul (α : Type) where
  mul : α → α → RustM α
class rust_primitives.ops.arith.Rem (α : Type) where
  rem : α → α → RustM α
class rust_primitives.ops.arith.Div (α : Type) where
  div : α → α → RustM α
class rust_primitives.ops.arith.Neg (α : Type) where
  neg : α → RustM α
class rust_primitives.ops.bit.Shr (α : Type) (β : Type) where
  shr : α → β → RustM α
class rust_primitives.ops.bit.Shl (α : Type) (β : Type) where
  shl : α → β → RustM α

infixl:65 " +? "   => rust_primitives.ops.arith.Add.add
infixl:65 " -? "   => rust_primitives.ops.arith.Sub.sub
infixl:70 " *? "   => rust_primitives.ops.arith.Mul.mul
infixl:75 " >>>? " => rust_primitives.ops.bit.Shr.shr
infixl:75 " <<<? " => rust_primitives.ops.bit.Shl.shl
infixl:70 " %? "   => rust_primitives.ops.arith.Rem.rem
infixl:70 " /? "   => rust_primitives.ops.arith.Div.div
prefix:75 "-?"   => rust_primitives.ops.arith.Neg.neg

attribute [specset bv, hax_bv_decide]
  rust_primitives.ops.arith.Add.add
  rust_primitives.ops.arith.Sub.sub
  rust_primitives.ops.arith.Mul.mul
  rust_primitives.ops.bit.Shr.shr
  rust_primitives.ops.bit.Shl.shl
  rust_primitives.ops.arith.Rem.rem
  rust_primitives.ops.arith.Div.div
  rust_primitives.ops.arith.Neg.neg

open Lean in
macro "declare_Hax_int_ops" s:(&"signed" <|> &"unsigned") typeName:ident width:term : command => do

  let signed ← match s.raw[0].getKind with
  | `signed => pure true
  | `unsigned => pure false
  | _ => throw .unsupportedSyntax

  let mut cmds ← Syntax.getArgs <$> `(

    /-- Addition on Rust integers. Panics on overflow. -/
    instance : rust_primitives.ops.arith.Add $typeName where
      add x y :=
        if ($(mkIdent (if signed then `BitVec.saddOverflow else `BitVec.uaddOverflow)) x.toBitVec y.toBitVec) then
          .fail .integerOverflow
        else pure (x + y)

    /-- Subtraction on Rust integers. Panics on overflow. -/
    instance : rust_primitives.ops.arith.Sub $typeName where
      sub x y :=
        if ($(mkIdent (if signed then `BitVec.ssubOverflow else `BitVec.usubOverflow)) x.toBitVec y.toBitVec) then
          .fail .integerOverflow
        else pure (x - y)

    /-- Multiplication on Rust integers. Panics on overflow. -/
    instance : rust_primitives.ops.arith.Mul $typeName where
      mul x y :=
        if ($(mkIdent (if signed then `BitVec.smulOverflow else `BitVec.umulOverflow)) x.toBitVec y.toBitVec) then
          .fail .integerOverflow
        else pure (x * y)
  )
  if signed then
    cmds := cmds.append $ ← Syntax.getArgs <$> `(
      /-- Division of signed Rust integers. Panics on overflow (when x is IntMin and `y = -1`)
        and when dividing by zero. -/
      instance : rust_primitives.ops.arith.Div $typeName where
        div x y :=
          if x = $(mkIdent (typeName.getId ++ `minValue)) && y = -1 then .fail .integerOverflow
          else if y = 0 then .fail .divisionByZero
          else pure (x / y)

      /-- Remainder of signed Rust integers. Panics on overflow (when x is IntMin and `y = -1`)
        and when the modulus is zero. -/
      instance : rust_primitives.ops.arith.Rem $typeName where
        rem x y :=
          if x = $(mkIdent (typeName.getId ++ `minValue)) && y = -1 then .fail .integerOverflow
          else if y = 0 then .fail .divisionByZero
          else pure (x % y)

      /-- Negation on signed integers. Panics on overflow (when `x` is `minValue`). -/
      instance : rust_primitives.ops.arith.Neg $typeName where
        neg x :=
          if x = $(mkIdent (typeName.getId ++ `minValue))
          then .fail .integerOverflow
          else pure (- x)
    )
  else -- unsigned
    cmds := cmds.append $ ← Syntax.getArgs <$> `(
      /-- Division on unsigned Rust integers. Panics when dividing by zero.  -/
      instance : rust_primitives.ops.arith.Div $typeName where
        div x y :=
          if y = 0 then .fail .divisionByZero
          else pure (x / y)

      /-- Division on unsigned Rust integers. Panics when the modulus is zero. -/
      instance : rust_primitives.ops.arith.Rem $typeName where
        rem x y :=
          if y = 0 then .fail .divisionByZero
          else pure (x % y)
    )
  return ⟨mkNullNode cmds⟩

declare_Hax_int_ops unsigned UInt8 8
declare_Hax_int_ops unsigned UInt16 16
declare_Hax_int_ops unsigned UInt32 32
declare_Hax_int_ops unsigned UInt64 64
declare_Hax_int_ops unsigned UInt128 128
declare_Hax_int_ops unsigned USize64 64
declare_Hax_int_ops signed Int8 8
declare_Hax_int_ops signed Int16 16
declare_Hax_int_ops signed Int32 32
declare_Hax_int_ops signed Int64 64
declare_Hax_int_ops signed Int128 128
declare_Hax_int_ops signed ISize System.Platform.numBits



open Lean in
set_option hygiene false in
macro "declare_Hax_shift_ops" : command => do
  let mut cmds := #[]
  let tys := [
    ("UInt8", ← `(term| 8)),
    ("UInt16", ← `(term| 16)),
    ("UInt32", ← `(term| 32)),
    ("UInt64", ← `(term| 64)),
    ("UInt128", ← `(term| 128)),
    ("USize64", ← `(term| 64)),
    ("Int8", ← `(term| 8)),
    ("Int16", ← `(term| 16)),
    ("Int32", ← `(term| 32)),
    ("Int64", ← `(term| 64)),
    ("Int128", ← `(term| 128)),
    ("ISize", ← `(term| OfNat.ofNat System.Platform.numBits))
  ]
  for (ty1, width1) in tys do
    for (ty2, _width2) in tys do

      let ty1Ident := mkIdent ty1.toName
      let ty2Ident := mkIdent ty2.toName
      let toTy1 := mkIdent ("to" ++ ty1).toName
      let ty2Signed := ty2.startsWith "I"
      let ty2ToNat := mkIdent (if ty2Signed then `toNatClampNeg else `toNat)
      let yConverted ← if ty1 == ty2 then `(y) else `(y.$ty2ToNat.$toTy1)

      cmds := cmds.push $ ← `(
        /-- Shift right for Rust integers. Panics when shifting by a negative number or
          by the bitsize or more. -/
        instance : rust_primitives.ops.bit.Shr $ty1Ident $ty2Ident where
          shr x y :=
            if 0 ≤ y && y < $width1
            then pure (x >>> $yConverted)
            else .fail .integerOverflow

        /-- Left shifting on signed integers. Panics when shifting by a negative number,
          or when shifting by more than the size. -/
        instance : rust_primitives.ops.bit.Shl $ty1Ident $ty2Ident where
          shl x y :=
            if 0 ≤ y && y < $width1
            then pure (x <<< $yConverted)
            else
              .fail .integerOverflow
      )
  return ⟨mkNullNode cmds⟩

declare_Hax_shift_ops


/-
## Specifications for integer operations
-/

open Lean in
set_option hygiene false in
macro "declare_Hax_int_ops_spec" s:(&"signed" <|> &"unsigned") typeName:ident width:term : command => do

  let signed ← match s.raw[0].getKind with
  | `signed => pure true
  | `unsigned => pure false
  | _ => throw .unsupportedSyntax

  let toX := if signed then mkIdent `toInt else mkIdent `toNat
  let minValue := mkIdent (typeName.getId ++ `minValue)
  let grind : TSyntax `tactic ←
    if signed then `(tactic| grind)
    else `(tactic| grind [toNat_add_of_lt, toNat_sub_of_le', toNat_mul_of_lt])

  let mut cmds ← Syntax.getArgs <$> `(
    namespace $typeName

      /-- Specification for rust addition -/
      @[specset int]
      theorem haxAdd_spec {x y : $typeName}
          (h : ¬ $(mkIdent (typeName.getId ++ `addOverflow)) x y) :
          ⦃ ⌜ True ⌝ ⦄ (x +? y) ⦃ ⇓ r => ⌜ r.$toX = x.$toX + y.$toX ⌝ ⦄ := by
        mvcgen [rust_primitives.ops.arith.Add.add]; $grind

      /-- Specification for rust subtraction -/
      @[specset int]
      theorem haxSub_spec {x y : $typeName}
          (h : ¬ $(mkIdent (typeName.getId ++ `subOverflow)) x y) :
          ⦃ ⌜ True ⌝ ⦄ (x -? y) ⦃ ⇓ r => ⌜ r.$toX = x.$toX - y.$toX ⌝ ⦄ := by
        mvcgen [rust_primitives.ops.arith.Sub.sub]; $grind

      /-- Specification for rust multiplication -/
      @[specset int]
      theorem haxMul_spec {x y : $typeName}
          (h : ¬ $(mkIdent (typeName.getId ++ `mulOverflow)) x y) :
          ⦃ ⌜ True ⌝ ⦄ (x *? y) ⦃ ⇓ r => ⌜ r.$toX = x.$toX * y.$toX ⌝ ⦄ := by
        mvcgen [rust_primitives.ops.arith.Mul.mul]; $grind
  )
  if signed then
    cmds := cmds.append $ ← Syntax.getArgs <$> `(
      /-- Specification for rust negation for signed integers-/
      @[specset int]
      theorem haxNeg_spec {x : $typeName} (hx : x ≠ $minValue) :
          ⦃ ⌜ True ⌝ ⦄ (-? x) ⦃ ⇓ r => ⌜ r.toInt = - x.toInt ⌝ ⦄ := by
        mvcgen [rust_primitives.ops.arith.Neg.neg]
        rw [toInt_neg_of_ne_intMin hx]

      /-- Specification for rust multiplication for signed integers-/
      @[specset int]
      theorem haxDiv_spec {x y : $typeName}
          (hx : x ≠ $minValue ∨ y ≠ -1) (hy : ¬ y = 0) :
          ⦃ ⌜ True ⌝ ⦄ (x /? y) ⦃ ⇓ r => ⌜ r.toInt = x.toInt.tdiv y.toInt ⌝ ⦄ := by
        have : ¬ (x = $minValue && y = -1) := by grind
        mvcgen [rust_primitives.ops.arith.Div.div]
        cases hx with
        | inl hx => apply toInt_div_of_ne_left x y hx
        | inr hx => apply toInt_div_of_ne_right x y hx

      /-- Specification for rust remainder for signed integers -/
      @[specset int]
      theorem haxRem_spec (x y : $typeName)
          (hx : x ≠ $minValue ∨ y ≠ -1) (hy : ¬ y = 0) :
          ⦃ ⌜ True ⌝ ⦄ (x %? y) ⦃ ⇓ r => ⌜ r.toInt = x.toInt.tmod y.toInt ⌝ ⦄ :=  by
        have : ¬ (x = $minValue && y = -1) := by grind
        mvcgen [rust_primitives.ops.arith.Rem.rem]
        apply toInt_mod
    )
  else -- unsigned
    cmds := cmds.append $ ← Syntax.getArgs <$> `(
      /-- Specification for rust multiplication for unsigned integers -/
      @[specset int]
      theorem haxDiv_spec (x y : $typeName) (h : ¬ y = 0) :
          ⦃ ⌜ True ⌝ ⦄ (x /? y) ⦃ ⇓ r => ⌜ r.toNat = x.toNat / y.toNat ⌝ ⦄ := by
        mvcgen [rust_primitives.ops.arith.Div.div]

      /-- Specification for rust remainder for unsigned integers -/
      @[specset int]
      theorem haxRem_spec (x y : $typeName) (h : ¬ y = 0) :
          ⦃ ⌜ True ⌝ ⦄ (x %? y) ⦃ ⇓ r => ⌜ r.toNat = x.toNat % y.toNat ⌝ ⦄ := by
        mvcgen [rust_primitives.ops.arith.Rem.rem]
    )
  cmds := cmds.push $ ← `(
    end $typeName
  )
  return ⟨mkNullNode cmds⟩

declare_Hax_int_ops_spec unsigned UInt8 8
declare_Hax_int_ops_spec unsigned UInt16 16
declare_Hax_int_ops_spec unsigned UInt32 32
declare_Hax_int_ops_spec unsigned UInt64 64
declare_Hax_int_ops_spec unsigned UInt128 128
declare_Hax_int_ops_spec unsigned USize64 64
declare_Hax_int_ops_spec signed Int8 8
declare_Hax_int_ops_spec signed Int16 16
declare_Hax_int_ops_spec signed Int32 32
declare_Hax_int_ops_spec signed Int64 64
declare_Hax_int_ops_spec signed Int128 128
declare_Hax_int_ops_spec signed ISize System.Platform.numBits

open Lean in
macro "declare_Hax_shift_ops_spec" : command => do
  let mut cmds := #[]
  let tys := [
    ("UInt8", ← `(term| 8)),
    ("UInt16", ← `(term| 16)),
    ("UInt32", ← `(term| 32)),
    ("UInt64", ← `(term| 64)),
    -- ("UInt128", ← `(term| 128)),
    ("Int8", ← `(term| 8)),
    ("Int16", ← `(term| 16)),
    ("Int32", ← `(term| 32)),
    ("Int64", ← `(term| 64)),
    -- ("Int128", ← `(term| 128)),
  ]
  for (ty1, width1) in tys do
    for (ty2, _width2) in tys do

      let ty1Ident := mkIdent ty1.toName
      let ty2Ident := mkIdent ty2.toName
      let toTy1 := mkIdent ("to" ++ ty1).toName
      let ty2Signed := ty2.startsWith "I"
      let ty2ToNat := mkIdent (if ty2Signed then `toNatClampNeg else `toNat)
      let yConverted ← if ty1 == ty2 then `(y) else `(y.$ty2ToNat.$toTy1)
      let haxShiftRight_spec := mkIdent ("haxShiftRight_" ++ ty2 ++ "_spec").toName
      let haxShiftLeft_spec := mkIdent ("haxShiftLeft_" ++ ty2 ++ "_spec").toName

      cmds := cmds.push $ ← `(
        namespace $ty1Ident
          /-- Bitvec-based specification for rust remainder on unsigned integers -/
          @[spec]
          theorem $haxShiftRight_spec (x : $ty1Ident) (y : $ty2Ident) :
            0 ≤ y →
            y.$ty2ToNat < $width1 →
            ⦃ ⌜ True ⌝ ⦄ (x >>>? y) ⦃ ⇓ r => ⌜ r = x >>> $yConverted ⌝ ⦄
            := by intros; mvcgen [rust_primitives.ops.bit.Shr.shr]; grind

          /-- Bitvec-based specification for rust remainder on unsigned integers -/
          @[spec]
          theorem $haxShiftLeft_spec (x : $ty1Ident) (y : $ty2Ident) :
            0 ≤ y →
            y.$ty2ToNat < $width1 →
            ⦃ ⌜ True ⌝ ⦄ (x <<<? y) ⦃ ⇓ r => ⌜ r = x <<< $yConverted ⌝ ⦄
            := by intros; mvcgen [rust_primitives.ops.bit.Shl.shl]; grind
        end $ty1Ident
      )
  return ⟨mkNullNode cmds⟩

declare_Hax_shift_ops_spec
