Montgomery.XZ, Loops: montladder proof scaffolding

2 files changed, 53 insertions, 25 deletions

diff --git a/src/Curves/Montgomery/XZ.v b/src/Curves/Montgomery/XZ.v
index 87a53b7fe..336ee6b95 100644
--- a/src/Curves/Montgomery/XZ.v
+++ b/src/Curves/Montgomery/XZ.v
@@ -2,7 +2,7 @@ Require Import Crypto.Algebra.Field.
 Require Import Crypto.Util.GlobalSettings Crypto.Util.Notations.
 Require Import Crypto.Util.Sum Crypto.Util.Prod Crypto.Util.LetIn.
 Require Import Crypto.Util.Decidable.
-Require Import Crypto.Util.ForLoop.
+Require Import Crypto.Experiments.Loops.
 Require Import Crypto.Spec.MontgomeryCurve Crypto.Curves.Montgomery.Affine.
 
 Module M.
@@ -110,26 +110,27 @@ Module M.
         ((x2, z2), (x3, z3))%core
       end.
  
-    Context {cswap:bool->F*F->F*F->(F*F)*(F*F)}.
-
+    Context {cswap:bool->F->F->F*F}.
     Local Notation xor := Coq.Init.Datatypes.xorb.
-
-    (* Ideally, we would verify that this corresponds to x coordinate
-       multiplication *)
     Local Open Scope core_scope.
-    Definition montladder (bound : positive) (testbit:Z->bool) (u:F) :=
-      let '(P1, P2, swap) :=
-          for (int i = BinInt.Z.pos bound; i >= 0; i--)
-              updating ('(P1, P2, swap) = ((1%F, 0%F), (u, 1%F), false)) {{
-            dlet s_i := testbit i in
-            dlet swap := xor swap s_i in
-            let '(P1, P2) := cswap swap P1 P2 in
-            dlet swap := s_i in
-            let '(P1, P2) := xzladderstep u P1 P2 in
-            (P1, P2, swap)
-      }} in
-      let '((x, z), _) := cswap swap P1 P2 in
-      x * Finv z.
 
+    (* TODO: make a nice notations for loops like here *)
+    Definition montladder (scalarbits : Z) (testbit:Z->bool) (x1:F) : F :=
+      let '(x2, z2, x3, z3, swap, _) := (* names of variables as used after the loop *)
+          (while (fun '(_, i) => BinInt.Z.geb i 0) (* the test of the loop *)
+          (fun '(x2, z2, x3, z3, swap, i) => (* names of variables as used in the loop; we should reuse the same names as for after the loop *)
+            dlet b := testbit i in (* the body... *)
+            dlet swap := xor swap b in
+            let (x2, x3) := cswap swap x2 x3 in
+            let (z2, z3) := cswap swap z2 z3 in
+            dlet swap := b in
+            let '((x2, z2), (x3, z3)) := xzladderstep x1 (x2, z2) (x3, z3) in
+            let i := BinInt.Z.pred i in (* the third "increment" component of a for loop; either between the test and body or just inlined into the body like here *)
+            (x2, z2, x3, z3, swap, i)) (* the "return value" of the body is always the exact same variable names as in the beginning of the body because we shadow the original binders, but I think for now this will be unavoidable boilerplate. *)
+          (BinInt.Z.to_nat scalarbits) (* bound on number of loop iterations, should come between test and body *)
+          (1%F, 0%F, x1, 1%F, false, BinInt.Z.pred scalarbits)) in (* initial values, these should come before the test and body *)
+      let (x2, x3) := cswap swap x2 x3 in
+      let (z2, z3) := cswap swap z2 z3 in
+      x2 * Finv z2.
   End MontgomeryCurve.
 End M.
diff --git a/src/Curves/Montgomery/XZProofs.v b/src/Curves/Montgomery/XZProofs.v
index fd2a7ee49..1e584f35f 100644
--- a/src/Curves/Montgomery/XZProofs.v
+++ b/src/Curves/Montgomery/XZProofs.v
@@ -138,12 +138,17 @@ Module M.
     Proof.
       cbv [ladder_invariant] in *.
       pose proof difference_preserved Q Q' as Hrw.
-      (* TODO: rewrite with in match argument with [sumwise (fieldwise (n:=2) Feq) (fun _ _ => True)] *)
-      match type of Hrw with
-        M.eq ?X ?Y => replace X with Y by admit
-      end.
-      assumption.
-    Admitted.
+      (* FIXME: what we actually want to do here is to rewrite with in
+      match argument with
+      [sumwise (fieldwise (n:=2) Feq) (fun _ _ => True)] *)
+      cbv [M.eq] in *; break_match; break_match_hyps;
+        destruct_head' @and; repeat split; subst;
+          try solve [intuition congruence].
+      congruence (* congruence failed, idk WHY *)
+      || (match goal with
+            [H:?f1 = ?x1, G:?f = ?f1 |- ?f = ?x1] => rewrite G; exact H
+          end).
+    Qed.
 
     Lemma to_xz_add x1 xz x'z' Q Q'
           (Hxz : projective xz) (Hx'z' : projective x'z')
@@ -179,5 +184,27 @@ Module M.
 
     Lemma projective_to_xz Q : projective (to_xz Q).
     Proof. t. Qed.
+
+    Local Notation montladder := (M.montladder(a24:=a24)(Fadd:=Fadd)(Fsub:=Fsub)(Fmul:=Fmul)(Fzero:=Fzero)(Fone:=Fone)(Finv:=Finv)(cswap:=fun b x y => if b then pair y x else pair x y)).
+    Import Crypto.Experiments.Loops.
+    Lemma montladder_correct P scalarbits testbit point : P (montladder scalarbits testbit point).
+    Proof.
+      cbv beta delta [M.montladder].
+      lazymatch goal with
+      | |- context [while ?t ?b ?l ?i] => eassert (_ (while t b l i) : Prop)
+      end.
+      lazymatch goal with
+      | |- ?e ?x
+        => eapply (while.by_invariant
+                     (* loop invariant *) _
+                     (* decreasing measure *) (fun s => BinInt.Z.to_nat (snd s))
+                     e)
+      end.
+      { (* invariant start *) admit. }
+      { (* invariant preservation *) admit. }
+      { (* measure start *) admit. }
+      { (* measure decreases *) admit. }
+      { (* invariant implies postcondition *) admit. }
+    Abort.
   End MontgomeryCurve.
 End M.