(*** Montgomery Multiplication *)
(** This file implements the proofs for Montgomery Form, Montgomery
    Reduction, and Montgomery Multiplication on [Z].  We follow
    Wikipedia. *)
From Coq Require Import ZArith Lia Equalities.
Require Import Crypto.Arithmetic.MontgomeryReduction.Definition.
Require Import Crypto.Util.ZUtil.EquivModulo.
Require Import Crypto.Util.ZUtil.Tactics.PullPush.Modulo.
Require Import Crypto.Util.ZUtil.Tactics.LtbToLt.
Require Import Crypto.Util.Tactics.BreakMatch.
Require Import Crypto.Util.Tactics.SimplifyRepeatedIfs.
Require Import Crypto.Util.Notations.

Declare Module Nop : Nop.
Module Import ImportEquivModuloInstances := Z.EquivModuloInstances Nop.

Local Existing Instance eq_Reflexive. (* speed up setoid_rewrite as per https://coq.inria.fr/bugs/show_bug.cgi?id=4978 *)

Local Open Scope Z_scope.

Section montgomery.
  Context (N : Z)
          (N_reasonable : N <> 0)
          (R : Z)
          (R_good : Z.gcd N R = 1).
  Local Notation "x ≡ y" := (Z.equiv_modulo N x y) : type_scope.
  Local Notation "x ≡ᵣ y" := (Z.equiv_modulo R x y) : type_scope.
  Context (R' : Z)
          (R'_good : R * R' ≡ 1).

  Lemma R'_good' : R' * R ≡ 1.
  Proof using R'_good. rewrite <- R'_good; apply f_equal2; lia. Qed.

  Local Notation to_montgomery_naive := (to_montgomery_naive R) (only parsing).
  Local Notation from_montgomery_naive := (from_montgomery_naive R') (only parsing).

  Lemma to_from_montgomery_naive x : to_montgomery_naive (from_montgomery_naive x) ≡ x.
  Proof using R'_good.
    unfold to_montgomery_naive, from_montgomery_naive.
    rewrite <- Z.mul_assoc, R'_good'.
    autorewrite with zsimplify; reflexivity.
  Qed.
  Lemma from_to_montgomery_naive x : from_montgomery_naive (to_montgomery_naive x) ≡ x.
  Proof using R'_good.
    unfold to_montgomery_naive, from_montgomery_naive.
    rewrite <- Z.mul_assoc, R'_good.
    autorewrite with zsimplify; reflexivity.
  Qed.

  (** * Modular arithmetic and Montgomery form *)
  Section general.
    Local Infix "+" := add : montgomery_scope.
    Local Infix "-" := sub : montgomery_scope.
    Local Infix "*" := (mul_naive R') : montgomery_scope.

    Lemma add_correct_naive x y : from_montgomery_naive (x + y) = from_montgomery_naive x + from_montgomery_naive y.
    Proof using Type. unfold from_montgomery_naive, add; lia. Qed.
    Lemma add_correct_naive_to x y : to_montgomery_naive (x + y) = (to_montgomery_naive x + to_montgomery_naive y)%montgomery.
    Proof using Type. unfold to_montgomery_naive, add; autorewrite with push_Zmul; reflexivity. Qed.
    Lemma sub_correct_naive x y : from_montgomery_naive (x - y) = from_montgomery_naive x - from_montgomery_naive y.
    Proof using Type. unfold from_montgomery_naive, sub; lia. Qed.
    Lemma sub_correct_naive_to x y : to_montgomery_naive (x - y) = (to_montgomery_naive x - to_montgomery_naive y)%montgomery.
    Proof using Type. unfold to_montgomery_naive, sub; autorewrite with push_Zmul; reflexivity. Qed.

    Theorem mul_correct_naive x y : from_montgomery_naive (x * y) = from_montgomery_naive x * from_montgomery_naive y.
    Proof using Type. unfold from_montgomery_naive, mul_naive; lia. Qed.
    Theorem mul_correct_naive_to x y : to_montgomery_naive (x * y) ≡ (to_montgomery_naive x * to_montgomery_naive y)%montgomery.
    Proof using R'_good.
      unfold to_montgomery_naive, mul_naive.
      rewrite <- !Z.mul_assoc, R'_good.
      autorewrite with zsimplify; apply (f_equal2 Z.modulo); lia.
    Qed.
  End general.

  (** * The REDC algorithm *)
  Section redc.
    Context (N' : Z)
            (N'_in_range : 0 <= N' < R)
            (N'_good : N * N' ≡ᵣ -1).

    Lemma N'_good' : N' * N ≡ᵣ -1.
    Proof using N'_good. rewrite <- N'_good; apply f_equal2; lia. Qed.

    Lemma N'_good'_alt x : (((x mod R) * (N' mod R)) mod R) * (N mod R) ≡ᵣ x * -1.
    Proof using N'_good.
      rewrite <- N'_good', Z.mul_assoc.
      unfold Z.equiv_modulo; push_Zmod.
      reflexivity.
    Qed.

    Section redc.
      Context (T : Z).

      Local Notation m := (((T mod R) * N') mod R).
      Local Notation prereduce := (prereduce N R N').

      Local Ltac t_fin_correct :=
        unfold Z.equiv_modulo; push_Zmod; autorewrite with zsimplify; reflexivity.

      Lemma prereduce_correct : prereduce T ≡ T * R'.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        transitivity ((T + m * N) * R').
        { unfold prereduce.
          autorewrite with zstrip_div; push_Zmod.
          rewrite N'_good'_alt.
          autorewrite with zsimplify pull_Zmod.
          reflexivity. }
        t_fin_correct.
      Qed.

      Lemma reduce_correct : reduce N R N' T ≡ T * R'.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        unfold reduce.
        break_match; rewrite prereduce_correct; t_fin_correct.
      Qed.

      Lemma partial_reduce_correct : partial_reduce N R N' T ≡ T * R'.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        unfold partial_reduce.
        break_match; rewrite prereduce_correct; t_fin_correct.
      Qed.

      Lemma reduce_via_partial_correct : reduce_via_partial N R N' T ≡ T * R'.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        unfold reduce_via_partial.
        break_match; rewrite partial_reduce_correct; t_fin_correct.
      Qed.

      Let m_small : 0 <= m < R. Proof. auto with zarith. Qed.

      Section generic.
        Lemma prereduce_in_range_gen B
        : 0 <= N
          -> 0 <= T <= R * B
          -> 0 <= prereduce T < B + N.
        Proof using N_reasonable m_small. unfold prereduce; auto with zarith nia. Qed.
      End generic.

      Section N_very_small.
        Context (N_very_small : 0 <= 4 * N < R).

        Lemma prereduce_in_range_very_small
          : 0 <= T <= (2 * N - 1) * (2 * N - 1)
            -> 0 <= prereduce T < 2 * N.
        Proof using N_reasonable N_very_small m_small. pose proof (prereduce_in_range_gen N); nia. Qed.
      End N_very_small.

      Section N_small.
        Context (N_small : 0 <= 2 * N < R).

        Lemma prereduce_in_range_small
          : 0 <= T <= (2 * N - 1) * (N - 1)
            -> 0 <= prereduce T < 2 * N.
        Proof using N_reasonable N_small m_small. pose proof (prereduce_in_range_gen N); nia. Qed.

        Lemma prereduce_in_range_small_fully_reduced
          : 0 <= T <= 2 * N
            -> 0 <= prereduce T <= N.
        Proof using N_reasonable N_small m_small. pose proof (prereduce_in_range_gen 1); nia. Qed.
      End N_small.

      Section N_small_enough.
        Context (N_small_enough : 0 <= N < R).

        Lemma prereduce_in_range_small_enough
          : 0 <= T <= R * R
            -> 0 <= prereduce T < R + N.
        Proof using N_reasonable N_small_enough m_small. pose proof (prereduce_in_range_gen R); nia. Qed.

        Lemma reduce_in_range_R
          : 0 <= T <= R * R
            -> 0 <= reduce N R N' T < R.
        Proof using N_reasonable N_small_enough m_small.
          intro H; pose proof (prereduce_in_range_small_enough H).
          unfold reduce, prereduce in *; break_match; Z.ltb_to_lt; nia.
        Qed.

        Lemma partial_reduce_in_range_R
          : 0 <= T <= R * R
            -> 0 <= partial_reduce N R N' T < R.
        Proof using N_reasonable N_small_enough m_small.
          intro H; pose proof (prereduce_in_range_small_enough H).
          unfold partial_reduce, prereduce in *; break_match; Z.ltb_to_lt; nia.
        Qed.

        Lemma reduce_via_partial_in_range_R
          : 0 <= T <= R * R
            -> 0 <= reduce_via_partial N R N' T < R.
        Proof using N_reasonable N_small_enough m_small.
          intro H; pose proof (prereduce_in_range_small_enough H).
          unfold reduce_via_partial, partial_reduce, prereduce in *; break_match; Z.ltb_to_lt; nia.
        Qed.
      End N_small_enough.

      Section unconstrained.
        Lemma prereduce_in_range
          : 0 <= T <= R * N
            -> 0 <= prereduce T < 2 * N.
        Proof using N_reasonable m_small. pose proof (prereduce_in_range_gen N); nia. Qed.

        Lemma reduce_in_range
        : 0 <= T <= R * N
          -> 0 <= reduce N R N' T < N.
        Proof using N_reasonable m_small.
          intro H; pose proof (prereduce_in_range H).
          unfold reduce, prereduce in *; break_match; Z.ltb_to_lt; nia.
        Qed.

        Lemma partial_reduce_in_range
        : 0 <= T <= R * N
          -> Z.min 0 (R - N) <= partial_reduce N R N' T < 2 * N.
        Proof using N_reasonable m_small.
          intro H; pose proof (prereduce_in_range H).
          unfold partial_reduce, prereduce in *; break_match; Z.ltb_to_lt;
            apply Z.min_case_strong; nia.
        Qed.

        Lemma reduce_via_partial_in_range
        : 0 <= T <= R * N
          -> Z.min 0 (R - N) <= reduce_via_partial N R N' T < N.
        Proof using N_reasonable m_small.
          intro H; pose proof (partial_reduce_in_range H).
          unfold reduce_via_partial in *; break_match; Z.ltb_to_lt; lia.
        Qed.
      End unconstrained.

      Section alt.
        Context (N_in_range : 0 <= N < R)
                (T_representable : 0 <= T < R * R).
        Lemma partial_reduce_alt_eq : partial_reduce_alt N R N' T = partial_reduce N R N' T.
        Proof using N_in_range N_reasonable T_representable m_small.
          assert (0 <= T + m * N < 2 * (R * R)) by nia.
          assert (0 <= T + m * N < R * (R + N)) by nia.
          assert (0 <= (T + m * N) / R < R + N) by auto with zarith.
          assert ((T + m * N) / R - N < R) by lia.
          assert (R * R <= T + m * N -> R <= (T + m * N) / R) by auto with zarith.
          assert (T + m * N < R * R -> (T + m * N) / R < R) by auto with zarith.
          assert (H' : (T + m * N) mod (R * R) = if R * R <=? T + m * N then T + m * N - R * R else T + m * N)
            by (break_match; Z.ltb_to_lt; autorewrite with zsimplify; lia).
          unfold partial_reduce, partial_reduce_alt, prereduce.
          rewrite H'; clear H'.
          simplify_repeated_ifs.
          set (m' := m) in *.
          autorewrite with zsimplify; push_Zmod; autorewrite with zsimplify; pull_Zmod.
          break_match; Z.ltb_to_lt; autorewrite with zsimplify; try reflexivity; lia.
        Qed.

        Lemma reduce_via_partial_alt_eq : reduce_via_partial_alt N R N' T = reduce_via_partial N R N' T.
        Proof.
            cbv [reduce_via_partial_alt reduce_via_partial].
            rewrite partial_reduce_alt_eq by lia. reflexivity.
        Qed.
      End alt.
    End redc.

    (** * Arithmetic in Montgomery form *)
    Section arithmetic.
      Local Infix "*" := (mul N R N') : montgomery_scope.

      Local Notation to_montgomery := (to_montgomery N R N').
      Local Notation from_montgomery := (from_montgomery N R N').
      Lemma to_from_montgomery a : to_montgomery (from_montgomery a) ≡ a.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        unfold to_montgomery, from_montgomery.
        transitivity ((a * 1) * 1); [ | apply f_equal2; lia ].
        rewrite <- !R'_good, !reduce_correct.
        unfold Z.equiv_modulo; push_Zmod; pull_Zmod.
        apply f_equal2; lia.
      Qed.
      Lemma from_to_montgomery a : from_montgomery (to_montgomery a) ≡ a.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        unfold to_montgomery, from_montgomery.
        rewrite !reduce_correct.
        transitivity (a * ((R * (R * R' mod N) * R') mod N)).
        { unfold Z.equiv_modulo; push_Zmod; pull_Zmod.
          apply f_equal2; lia. }
        { repeat first [ rewrite R'_good
                       | reflexivity
                       | push_Zmod; pull_Zmod; progress autorewrite with zsimplify
                       | progress unfold Z.equiv_modulo ]. }
      Qed.

      Theorem mul_correct x y : from_montgomery (x * y) ≡ from_montgomery x * from_montgomery y.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        unfold from_montgomery, mul.
        rewrite !reduce_correct; apply f_equal2; lia.
      Qed.
      Theorem mul_correct_to x y : to_montgomery (x * y) ≡ (to_montgomery x * to_montgomery y)%montgomery.
      Proof using N'_good N'_in_range N_reasonable R'_good.
        unfold to_montgomery, mul.
        rewrite !reduce_correct.
        transitivity (x * y * R * 1 * 1 * 1);
          [ rewrite <- R'_good at 1
          | rewrite <- R'_good at 1 2 3 ];
          autorewrite with zsimplify;
          unfold Z.equiv_modulo; push_Zmod; pull_Zmod.
        { apply f_equal2; lia. }
        { apply f_equal2; lia. }
      Qed.
    End arithmetic.
  End redc.
End montgomery.

Module Import LocalizeEquivModuloInstances := Z.RemoveEquivModuloInstances Nop.
