-- Clean/Gadgets/Rotation32/Theorems.lean

import Clean.Utils.Field

import Clean.Utils.Bitwise

import Clean.Utils.Rotation

import Clean.Types.U32

import Clean.Gadgets.ByteDecomposition.ByteDecomposition

variable {p : ℕ} [Fact p.Prime]

variable [p_large_enough: Fact (p > 2^16 + 2^8)]

namespace Gadgets.Rotation32.Theorems

open Gadgets.ByteDecomposition.Theorems (byteDecomposition_lift)

open Utils.Rotation

def rotRight32_u32 : U32 ℕ → ℕ → U32 ℕ
  | ⟨ x0, x1, x2, x3 ⟩, o => ⟨
    (x0 / 2^o) + (x1 % 2^o) * 2^(8-o),
    (x1 / 2^o) + (x2 % 2^o) * 2^(8-o),
    (x2 / 2^o) + (x3 % 2^o) * 2^(8-o),
    (x3 / 2^o) + (x0 % 2^o) * 2^(8-o),
  ⟩

-- Repository definitions used:
content: structure U32 (T : Type) where
  x0 : T
  x1 : T
  x2 : T
  x3 : T
deriving DecidableEq

content: def rotRight32 (x : ℕ) (offset : ℕ) : ℕ :=
  let offset := offset % 32
  let low := x % (2^offset)
  let high := x / (2^offset)
  low * (2^(32-offset)) + high

content: def valueNat (x : U32 ℕ) :=
  x.x0 + x.x1 * 256 + x.x2 * 256^2 + x.x3 * 256^3

content: def valueNat (x : U64 ℕ) :=
  x.x0 + x.x1 * 256 + x.x2 * 256^2 + x.x3 * 256^3 +
  x.x4 * 256^4 + x.x5 * 256^5 + x.x6 * 256^6 + x.x7 * 256^7

content: structure U64 (T : Type) where
  x0 : T
  x1 : T
  x2 : T
  x3 : T
  x4 : T
  x5 : T
  x6 : T
  x7 : T
deriving DecidableEq

-- Repository lemmas used:
content: lemma mul_div_256_off {offset : ℕ} (ho : offset < 8) {x : ℕ} (i : ℕ) (h : i > 0):
    (x * 256^i) / 2^offset = x * 256^(i-1) * 2^(8-offset)

content: lemma div_256_two_power {offset : ℕ} (ho : offset < 8) {i : ℕ} (h : i > 0):
    256^i / 2^offset = 256^(i-1) * 2^(8-offset)

content: lemma divides_256_two_power {offset : ℕ} (ho : offset < 8) {x i : ℕ} (h : i > 0):
    (2^offset) ∣ x * (256 ^ i)

content: lemma mul_mod256_off {offset : ℕ} (ho : offset < 8) (x i : ℕ) (h : i > 0):
    (x * 256^i) % 2^offset = 0

content: lemma shifted_decomposition_eq {offset : ℕ} (ho : offset < 8) {x1 x2 : ℕ} :
    (x1 / 2 ^ offset + x2 % 2 ^ offset * 2 ^ (8-offset)) * 256 =
    (2^offset * (x1 / 2^offset) + (x2 % 2^offset) * 256) * 2^(8-offset)

content: lemma soundness_simp {offset : ℕ} {x y : ℕ} :
    x % 2 ^ offset * 2 ^ (8-offset) * y + 2 ^ offset * (x / 2 ^ offset) * 2 ^ (8-offset) * y =
    x * y * 2^ (8-offset)

content: lemma soundness_simp' {offset : ℕ} {x : ℕ} :
    x % 2 ^ offset * 2 ^ (8-offset) + 2 ^ offset * (x / 2 ^ offset) * 2 ^ (8-offset) =
    x * 2^ (8-offset)

content: lemma shifted_decomposition_eq'' {offset : ℕ} (ho : offset < 8) {x1 x2 i : ℕ} (hi : i > 0) :
    (x1 / 2 ^ offset + x2 % 2 ^ offset * 2 ^ (8-offset)) * 256^i =
    (2^offset * (x1 / 2^offset) * 2^(8-offset) * 256^(i-1) +
    (x2 % 2^offset) * 2^(8-offset) * 256^i)

content: lemma Nat.pow_minus_one_mul {x y : ℕ} (hy : y > 0) : x ^ y = x * x ^ (y - 1)

content: lemma shifted_decomposition_eq' {offset : ℕ} (ho : offset < 8) {x1 x2 i : ℕ} (hi : i > 0) :
    (x1 / 2 ^ offset + x2 % 2 ^ offset * 2 ^ (8-offset)) * 256^i =
    (2^offset * (x1 / 2^offset) + (x2 % 2^offset) * 256) * 2^(8-offset) * 256^(i-1)

-- Local definitions used:
content: def rotRight32_u32 : U32 ℕ → ℕ → U32 ℕ
  | ⟨ x0, x1, x2, x3 ⟩, o => ⟨
    (x0 / 2^o) + (x1 % 2^o) * 2^(8-o),
    (x1 / 2^o) + (x2 % 2^o) * 2^(8-o),
    (x2 / 2^o) + (x3 % 2^o) * 2^(8-o),
    (x3 / 2^o) + (x0 % 2^o) * 2^(8-o),
  ⟩

-- Local lemmas used:
content: lemma h_mod32 {o : ℕ} (ho : o < 8) {x0 x1 x2 x3 : ℕ} :
    (x0 + x1 * 256 + x2 * 256^2 + x3 * 256^3) % 2^o = x0 % 2^o

content: lemma h_div32 {o : ℕ} (ho : o < 8) {x0 x1 x2 x3: ℕ} :
    (x0 + x1 * 256 + x2 * 256^2 + x3 * 256^3) / 2^o
    = x0 / 2^o + x1 * 2^(8-o) + x2 * 256 * 2^(8-o) + x3 * 256^2 * 2^(8-o)

content: lemma h_x0_const32 {o : ℕ} (ho : o < 8) :
    2^(8-o) * 256^3 = 2^(32-o)

-- Target theorem and proof:
theorem rotation32_bits_soundness {o : ℕ} (ho : o < 8) {x : U32 ℕ} :
    (rotRight32_u32 x o).valueNat = rotRight32 x.valueNat o
:= by
  -- simplify the goal
  simp only [rotRight32, rotRight32_u32, U32.valueNat]

  have offset_mod_32 : o % 32 = o := Nat.mod_eq_of_lt (by linarith)
  simp only [offset_mod_32]
  rw [h_mod32 ho, h_div32 ho]

  -- proof technique: we care about only what happens to x0, all "internal" terms remain
  -- the same, and are just divided by 2^o
  rw [shifted_decomposition_eq ho]
  repeat rw [shifted_decomposition_eq'' ho (by simp only [Nat.ofNat_pos])]
  simp only [Nat.add_one_sub_one, pow_one, add_mul, add_assoc]
  rw [←add_assoc _ _ (_ * 256 ^ 3), soundness_simp]
  nth_rw 4 [←add_assoc]
  rw [Nat.mul_right_comm _ 256, soundness_simp]
  nth_rw 2 [←add_assoc]
  rw [Nat.mul_right_comm _ 256, soundness_simp']
  rw [mul_assoc (x.x0 % 2 ^ o), h_x0_const32 ho]
  ac_rfl