fix WWMM partial evaluation

6 files changed, 21 insertions, 9 deletions

diff --git a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Definition.v b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Definition.v
index 9a6b089da..2ea623b0b 100644
--- a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Definition.v
+++ b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Definition.v
@@ -7,6 +7,7 @@
 Require Import Coq.ZArith.ZArith.
 Require Import Crypto.Util.Notations.
 Require Import Crypto.Util.LetIn.
+Require Import Crypto.Util.ZUtil.Definitions.
 
 Local Open Scope Z_scope.
 
@@ -33,7 +34,7 @@ Section WordByWordMontgomery.
     Local Definition A_a := dlet p := divmod A in p. Local Definition A' := fst A_a. Local Definition a := snd A_a.
     Local Definition S1 := add S (scmul a B).
     Local Definition s := snd (divmod S1).
-    Local Definition q := s * k mod r.
+    Local Definition q := fst (Z.mul_split r s k).
     Local Definition S2 := add S1 (scmul q N).
     Local Definition S3 := fst (divmod S2).
     Local Definition S4 := drop_high S3.
diff --git a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Definition.v b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Definition.v
index 71c8ea117..11b73b6a8 100644
--- a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Definition.v
+++ b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Definition.v
@@ -7,6 +7,7 @@
 Require Import Coq.ZArith.ZArith.
 Require Import Crypto.Util.Notations.
 Require Import Crypto.Util.LetIn.
+Require Import Crypto.Util.ZUtil.Definitions.
 
 Local Open Scope Z_scope.
 
@@ -35,7 +36,7 @@ Section WordByWordMontgomery.
     Local Definition A_a := dlet p := divmod _ A in p. Local Definition A' := fst A_a. Local Definition a := snd A_a.
     Local Definition S1 := add _ S (scmul _ a B).
     Local Definition s := snd (divmod _ S1).
-    Local Definition q := s * k mod r.
+    Local Definition q := fst (Z.mul_split r s k).
     Local Definition S2 := add _ S1 (scmul _ q N).
     Local Definition S3 := fst (divmod _ S2).
     Local Definition S4 := drop_high S3.
diff --git a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Proofs.v b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Proofs.v
index e09add277..253b56b40 100644
--- a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Proofs.v
+++ b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Dependent/Proofs.v
@@ -8,6 +8,7 @@ Require Import Crypto.Util.ZUtil.
 Require Import Crypto.Arithmetic.ModularArithmeticTheorems Crypto.Spec.ModularArithmetic.
 Require Import Crypto.Arithmetic.MontgomeryReduction.WordByWord.Abstract.Dependent.Definition.
 Require Import Crypto.Algebra.Ring.
+Require Import Crypto.Util.ZUtil.MulSplit.
 Require Import Crypto.Util.Sigma.
 Require Import Crypto.Util.Tactics.SetEvars.
 Require Import Crypto.Util.Tactics.SubstEvars.
@@ -100,7 +101,7 @@ Section WordByWordMontgomery.
       { transitivity ((N+B-1 + (r-1)*B + (r-1)*N) / r);
           [ | set_evars; ring_simplify_subterms; subst_evars; reflexivity ].
         Z.peel_le; repeat apply Z.add_le_mono; repeat apply Z.mul_le_mono_nonneg; try lia;
-          repeat autounfold with word_by_word_montgomery;
+          repeat autounfold with word_by_word_montgomery; rewrite ?Z.mul_split_mod;
           autorewrite with push_eval;
             try Z.zero_bounds;
             auto with lia. }
@@ -121,7 +122,7 @@ Section WordByWordMontgomery.
 
     Lemma S3_nonneg : 0 <= eval S3.
     Proof.
-      repeat autounfold with word_by_word_montgomery;
+      repeat autounfold with word_by_word_montgomery; rewrite ?Z.mul_split_mod;
         autorewrite with push_eval; [].
       rewrite ?Npos_correct; Z.zero_bounds; lia.
     Qed.
@@ -163,6 +164,7 @@ Section WordByWordMontgomery.
       autorewrite with pull_Zmod.
       replace 0 with (0 mod r) by apply Zmod_0_l.
       eapply F.eq_of_Z_iff.
+      rewrite Z.mul_split_mod.
       repeat rewrite ?F.of_Z_add, ?F.of_Z_mul, <-?F.of_Z_mod.
       rewrite <-Algebra.Hierarchy.associative.
       replace ((F.of_Z r k * F.of_Z r (eval N))%F) with (F.opp (m:=r) F.one).
diff --git a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Proofs.v b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Proofs.v
index 056b816a2..5ce3e197c 100644
--- a/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Proofs.v
+++ b/src/Arithmetic/MontgomeryReduction/WordByWord/Abstract/Proofs.v
@@ -8,6 +8,7 @@ Require Import Crypto.Util.ZUtil.
 Require Import Crypto.Arithmetic.ModularArithmeticTheorems Crypto.Spec.ModularArithmetic.
 Require Import Crypto.Arithmetic.MontgomeryReduction.WordByWord.Abstract.Definition.
 Require Import Crypto.Algebra.Ring.
+Require Import Crypto.Util.ZUtil.MulSplit.
 Require Import Crypto.Util.Sigma.
 Require Import Crypto.Util.Tactics.SetEvars.
 Require Import Crypto.Util.Tactics.SubstEvars.
@@ -110,7 +111,7 @@ Section WordByWordMontgomery.
       { transitivity ((N+B-1 + (r-1)*B + (r-1)*N) / r);
           [ | set_evars; ring_simplify_subterms; subst_evars; reflexivity ].
         Z.peel_le; repeat apply Z.add_le_mono; repeat apply Z.mul_le_mono_nonneg; try lia;
-          repeat autounfold with word_by_word_montgomery;
+          repeat autounfold with word_by_word_montgomery; rewrite ?Z.mul_split_mod;
           autorewrite with push_eval;
             try Z.zero_bounds;
             auto with lia. }
@@ -131,7 +132,7 @@ Section WordByWordMontgomery.
 
     Lemma S3_nonneg : 0 <= eval S3.
     Proof.
-      repeat autounfold with word_by_word_montgomery;
+      repeat autounfold with word_by_word_montgomery; rewrite Z.mul_split_mod;
         autorewrite with push_eval; [].
       rewrite ?Npos_correct; Z.zero_bounds; lia.
     Qed.
@@ -153,7 +154,7 @@ Section WordByWordMontgomery.
     Lemma numlimbs_S4 : numlimbs S4 = min (max (1 + numlimbs S) (1 + max (1 + numlimbs B) (numlimbs N))) (1 + R_numlimbs).
     Proof.
       cbn [plus].
-      repeat autounfold with word_by_word_montgomery.
+      repeat autounfold with word_by_word_montgomery; rewrite Z.mul_split_mod.
       repeat autorewrite with push_numlimbs.
       change Init.Nat.max with Nat.max.
       rewrite <- ?(Max.max_assoc (numlimbs S)).
@@ -182,6 +183,7 @@ Section WordByWordMontgomery.
       autorewrite with pull_Zmod.
       replace 0 with (0 mod r) by apply Zmod_0_l.
       eapply F.eq_of_Z_iff.
+      rewrite Z.mul_split_mod.
       repeat rewrite ?F.of_Z_add, ?F.of_Z_mul, <-?F.of_Z_mod.
       rewrite <-Algebra.Hierarchy.associative.
       replace ((F.of_Z r k * F.of_Z r (eval N))%F) with (F.opp (m:=r) F.one).
diff --git a/src/Arithmetic/MontgomeryReduction/WordByWord/Definition.v b/src/Arithmetic/MontgomeryReduction/WordByWord/Definition.v
index 55724b940..0b4666856 100644
--- a/src/Arithmetic/MontgomeryReduction/WordByWord/Definition.v
+++ b/src/Arithmetic/MontgomeryReduction/WordByWord/Definition.v
@@ -9,6 +9,7 @@ Require Import Crypto.Arithmetic.Saturated.
 Require Import Crypto.Arithmetic.MontgomeryReduction.WordByWord.Abstract.Dependent.Definition.
 Require Import Crypto.Util.Notations.
 Require Import Crypto.Util.LetIn.
+Require Import Crypto.Util.ZUtil.Definitions.
 
 Local Open Scope Z_scope.
 
@@ -42,7 +43,7 @@ Section WordByWordMontgomery.
     := divmod_cps A (fun '(A, a) =>
        @scmul_cps r _ a B _ (fun aB => @add_cps r _ S' aB _ (fun S1 =>
        divmod_cps S1 (fun '(_, s) =>
-       dlet q := s * k mod r in
+       dlet q := fst (Z.mul_split r s k) in
        @scmul_cps r _ q N _ (fun qN => @add_cps r _ S1 qN _ (fun S2 =>
        divmod_cps S2 (fun '(S3, _) =>
        @drop_high_cps (S R_numlimbs) S3 _ (fun S4 => rest (A, S4))))))))).
diff --git a/src/Specific/MontgomreyP256.v b/src/Specific/MontgomreyP256.v
index fb3716bbe..f0a156574 100644
--- a/src/Specific/MontgomreyP256.v
+++ b/src/Specific/MontgomreyP256.v
@@ -20,6 +20,8 @@ Definition mulmod_256 : { f:Tuple.tuple Z 4 -> Tuple.tuple Z 4 -> Tuple.tuple Z
 Proof.
   eapply (lift2_sig (fun A B c => c = (redc (r:=r)(R_numlimbs:=4) p256 A B 1)
                            )); eexists.
+  cbv -[Definitions.Z.add_get_carry_full Definitions.Z.mul_split runtime_add runtime_mul Let_In].
+  (*
   cbv [
       r wt sz p256
         CPSUtil.Tuple.tl_cps CPSUtil.Tuple.hd_cps
@@ -78,8 +80,11 @@ Nat.max
 Positional.to_associational_cps
 Z.of_nat *)
     ].
+  Unset Printing Notations.
 
   (* cbv -[runtime_add runtime_mul LetIn.Let_In Definitions.Z.add_get_carry_full Definitions.Z.mul_split]. *)
 
   (* basesystem_partial_evaluation_RHS. *)
-Abort.
+   *)
+  reflexivity.
+Defined.