Require Import Coq.ZArith.ZArith.
Require Import Coq.micromega.Lia.
Require Import Coq.micromega.Psatz.
Require Import Coq.Bool.Bool.

Require Import biteq.Lib.

Require Import bbv.Word.

#[local] Open Scope Z_scope.

Record bv := Zbv { size: N; val: Z; }.

Definition bvZ (b: bv) := (val b) mod (2 ^ (Z.of_N (size b))).

Definition bvadd (sz: N) (a b: bv) :=
  Zbv sz (bvZ a + bvZ b).

Compute bvZ (bvadd 4%N (Zbv 3%N 5) (Zbv 3%N 5)).

Lemma pow_2_div : forall p q, 0 <= p <= q -> (2 ^ p | 2 ^ q).
Proof.
  intros.
  exploit Z.le_exists_sub.
  assert (p <= q) by lia. apply H0.
  intros. inversion H0. inversion H1. rewrite H2.
  rewrite Z.pow_add_r; try lia.
  apply Z.divide_mul_r. apply Z.divide_refl.
Qed.

Lemma mod_mod_pow_2 :
  forall a p q,
    0 <= p <= q -> (a mod 2 ^ q) mod (2 ^ p) = a mod (2 ^ p).
Proof.
  intros. rewrite <- Znumtheory.Zmod_div_mod; try lia. now apply pow_2_div.
Qed.

Lemma bounded_add_assoc :
  forall (p q r s v u: N) (a b c: Z),
    (v <= q)%N -> (v <= u)%N ->
    bvZ (bvadd v (bvadd u (Zbv p a) (Zbv r b)) (Zbv s c))
    = bvZ (bvadd v (Zbv p a) (bvadd q (Zbv r b) (Zbv s c))).
Proof.
  intros.
  (* Zplus_mod *)
  unfold bvadd, bvZ; simpl in *.
  rewrite Zplus_mod by eauto. rewrite mod_mod_pow_2 by lia. symmetry.
  rewrite Zplus_mod by eauto.
  replace (((b mod 2 ^ Z.of_N r + c mod 2 ^ Z.of_N s) mod 2 ^ Z.of_N q) mod 2 ^ Z.of_N v)
    with ((b mod 2 ^ Z.of_N r + c mod 2 ^ Z.of_N s) mod 2 ^ Z.of_N v)
    by (symmetry; apply mod_mod_pow_2; lia).
  rewrite <- Zplus_mod. symmetry.
  rewrite <- Zplus_mod. now rewrite Z.add_assoc.
Qed.