From be82bd8e0347e090037ff1d30a22a9d614db8c9f Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter@wuille.net>
Date: Mon, 31 Jan 2022 18:19:45 -0500
Subject: [PATCH] Improve comments/checks for fe_sqrt

---
 src/field.h      | 14 ++++++++------
 src/field_impl.h | 17 +++++++++++++++--
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/src/field.h b/src/field.h
index 597c9fe3..eb6c586b 100644
--- a/src/field.h
+++ b/src/field.h
@@ -246,12 +246,14 @@ static void secp256k1_fe_mul(secp256k1_fe *r, const secp256k1_fe *a, const secp2
  */
 static void secp256k1_fe_sqr(secp256k1_fe *r, const secp256k1_fe *a);
 
-/** If a has a square root, it is computed in r and 1 is returned. If a does not
- *  have a square root, the root of its negation is computed and 0 is returned.
- *  The input's magnitude can be at most 8. The output magnitude is 1 (but not
- *  guaranteed to be normalized). The result in r will always be a square
- *  itself. */
-static int secp256k1_fe_sqrt(secp256k1_fe *r, const secp256k1_fe *a);
+/** Compute a square root of a field element.
+ *
+ * On input, a must be a valid field element with magnitude<=8; r need not be initialized.
+ * Performs {r = sqrt(a)} or {r = sqrt(-a)}, whichever exists. The resulting value
+ * represented by r will be a square itself. Variables r and a must not point to the same object.
+ * On output, r will have magnitude 1 but will not be normalized.
+ */
+static int secp256k1_fe_sqrt(secp256k1_fe * SECP256K1_RESTRICT r, const secp256k1_fe * SECP256K1_RESTRICT a);
 
 /** Sets a field element to be the (modular) inverse of another. Requires the input's magnitude to be
  *  at most 8. The output magnitude is 1 (but not guaranteed to be normalized). */
diff --git a/src/field_impl.h b/src/field_impl.h
index ff4c8eca..89ffc76c 100644
--- a/src/field_impl.h
+++ b/src/field_impl.h
@@ -55,9 +55,13 @@ static int secp256k1_fe_sqrt(secp256k1_fe *r, const secp256k1_fe *a) {
      *  itself always a square (a ** ((p+1)/4) is the square of a ** ((p+1)/8)).
      */
     secp256k1_fe x2, x3, x6, x9, x11, x22, x44, x88, x176, x220, x223, t1;
-    int j;
+    int j, ret;
 
+#ifdef VERIFY
     VERIFY_CHECK(r != a);
+    secp256k1_fe_verify(a);
+    VERIFY_CHECK(a->magnitude <= 8);
+#endif
 
     /** The binary representation of (p + 1)/4 has 3 blocks of 1s, with lengths in
      *  { 2, 22, 223 }. Use an addition chain to calculate 2^n - 1 for each block:
@@ -141,7 +145,16 @@ static int secp256k1_fe_sqrt(secp256k1_fe *r, const secp256k1_fe *a) {
     /* Check that a square root was actually calculated */
 
     secp256k1_fe_sqr(&t1, r);
-    return secp256k1_fe_equal(&t1, a);
+    ret = secp256k1_fe_equal(&t1, a);
+
+#ifdef VERIFY
+    if (!ret) {
+        secp256k1_fe_negate(&t1, &t1, 1);
+        secp256k1_fe_normalize_var(&t1);
+        VERIFY_CHECK(secp256k1_fe_equal_var(&t1, a));
+    }
+#endif
+    return ret;
 }
 
 #ifndef VERIFY