-- Clean/Gadgets/Rotation64/Theorems.lean

import Clean.Utils.Field

import Clean.Utils.Bitwise

import Clean.Utils.Rotation

import Clean.Types.U64

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

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

namespace Gadgets.Rotation64.Theorems

open Utils.Rotation

def rotRight64_u64 : U64 ℕ → ℕ → U64 ℕ
  | ⟨ x0, x1, x2, x3, x4, x5, x6, x7 ⟩, 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) + (x4 % 2^o) * 2^(8-o),
    (x4 / 2^o) + (x5 % 2^o) * 2^(8-o),
    (x5 / 2^o) + (x6 % 2^o) * 2^(8-o),
    (x6 / 2^o) + (x7 % 2^o) * 2^(8-o),
    (x7 / 2^o) + (x0 % 2^o) * 2^(8-o),
  ⟩

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

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: def rotRight64 (x : ℕ) (offset : ℕ) : ℕ :=
  let offset := offset % 64
  let low := x % (2^offset)
  let high := x / (2^offset)
  low * (2^(64-offset)) + high

-- 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 rotRight64_u64 : U64 ℕ → ℕ → U64 ℕ
  | ⟨ x0, x1, x2, x3, x4, x5, x6, x7 ⟩, 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) + (x4 % 2^o) * 2^(8-o),
    (x4 / 2^o) + (x5 % 2^o) * 2^(8-o),
    (x5 / 2^o) + (x6 % 2^o) * 2^(8-o),
    (x6 / 2^o) + (x7 % 2^o) * 2^(8-o),
    (x7 / 2^o) + (x0 % 2^o) * 2^(8-o),
  ⟩

-- Local lemmas used:
content: lemma h_mod {o : ℕ} (ho : o < 8) {x0 x1 x2 x3 x4 x5 x6 x7 : ℕ} :
    (x0 + x1 * 256 + x2 * 256 ^ 2 + x3 * 256 ^ 3 + x4 * 256 ^ 4 + x5 * 256 ^ 5 + x6 * 256 ^ 6 + x7 * 256 ^ 7) %
      2^o = x0 % 2^o

content: lemma h_div {o : ℕ} (ho : o < 8) {x0 x1 x2 x3 x4 x5 x6 x7 : ℕ} :
    (x0 + x1 * 256 + x2 * 256^2 + x3 * 256^3 + x4 * 256^4 + x5 * 256^5 + x6 * 256^6 + x7 * 256^7) / 2 ^ o
    = x0 / 2^o + x1 * 2^(8-o) + x2 * 256 * 2^(8-o) + x3 * 256^2 * 2^(8-o) + x4 * 256^3 * 2^(8-o) +
    x5 * 256^4 * 2^(8-o) + x6 * 256^5 * 2^(8-o) + x7 * 256^6 * 2^(8-o)

content: lemma h_x0_const {o : ℕ} (ho : o < 8) :
    2^(8-o) * 256^7 = 2^(64-o)

-- Target theorem and proof:
theorem rotation64_bits_soundness {o : ℕ} (ho : o < 8) {x : U64 ℕ} :
    (rotRight64_u64 x o).valueNat = rotRight64 x.valueNat o
:= by
  -- simplify the goal
  simp only [rotRight64, rotRight64_u64, U64.valueNat]

  have offset_mod_64 : o % 64 = o := Nat.mod_eq_of_lt (by linarith)
  simp only [offset_mod_64]
  rw [h_mod ho, h_div 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]

  -- we do a bit of expression juggling here
  rw [←add_assoc _ _ (_ * 256 ^ 7), soundness_simp]
  nth_rw 12 [←add_assoc]
  rw [soundness_simp]
  nth_rw 10 [←add_assoc]
  rw [soundness_simp]
  nth_rw 8 [←add_assoc]
  rw [soundness_simp]
  nth_rw 6 [←add_assoc]
  rw [soundness_simp]
  nth_rw 4 [←add_assoc]
  nth_rw 2 [Nat.mul_right_comm]
  rw [soundness_simp]
  nth_rw 2 [←add_assoc]
  rw [soundness_simp']
  nth_rw 7 [mul_assoc]

  -- at the end the terms are equal
  rw [h_x0_const ho]
  ac_rfl