[Refactor] Adds some tests and contains some documentation suggestions for the fix/cast-lossy-field-to-group PR

circuit/types/field/src/square_root.rs

@@ -64,22 +64,22 @@ impl<E: Environment> Field<E> {
 }
 
 impl<E: Environment> Field<E> {
-    /// Returns both square roots of `self` (hence the plural 'roots' in the name of the function),
-    /// along with a boolean error flag, which is set iff `self` is not a square.
+    /// Returns both square roots of `self` and a `Boolean` flag, which is set iff `self` is not a square.
     ///
-    /// In the console computation:
-    /// if `self` is a non-zero square,
-    /// the first field result is the positive root (i.e. closer to 0)
-    /// and the second field result is the negative root (i.e. closer to the prime);
-    /// if `self` is 0, both field results are 0;
-    /// if `self` is not a square, both field results are 0, but immaterial.
+    /// If `self` is a non-zero square,
+    ///  - the first field result is the positive root (i.e. closer to 0)
+    ///  - the second field result is the negative root (i.e. closer to the prime)
+    ///  - the flag is 0
     ///
-    /// The 'nondeterministic' part of the function name refers to the synthesized circuit,
-    /// whose represented computation, unlike the console computation just described,
-    /// returns the two roots (if `self` is a non-zero square) in no specified order.
-    /// This nondeterminism saves constraints, but generally this circuit should be only used
-    /// as part of larger circuits for which the nondeterminism in the order of the two roots does not matter,
-    /// and where the larger circuits represent deterministic computations despite this internal nondeterminism.
+    /// If `self` is 0,
+    ///  - both field results are 0
+    ///  - the flag is 0
+    ///
+    /// If `self` is not a square,
+    ///  - both field results are 0
+    ///  - the flag is 1
+    ///
+    /// Note that there is no ordering on the two roots returned by this function.
     pub fn square_roots_flagged_nondeterministic(&self) -> (Self, Self, Boolean<E>) {
         // Obtain (p-1)/2, as a constant field element.
         let modulus_minus_one_div_two = match E::BaseField::from_bigint(E::BaseField::modulus_minus_one_div_two()) {
@@ -92,30 +92,24 @@ impl<E: Environment> Field<E> {
         let is_nonzero_square = euler.is_one();
 
         // Calculate the witness for the first square result.
-        // The called function square_root returns the square root closer to 0.
+        // Note that the function `square_root` returns the square root closer to 0.
         let root_witness = match self.eject_value().square_root() {
             Ok(root) => root,
             Err(_) => console::Field::zero(),
         };
 
-        // In order to avoid actually calculating the square root in the circuit,
-        // we would like to generate a constraint saying that squaring the root yields self.
-        // But this constraint would have no solutions if self is not a square.
-        // So we introduce a new variable that is either self (if square) or 0 (otherwise):
-        // either way, this new variable is a square.
+        // Initialize the square element, which is either `self` or 0, depending on whether `self` is a square.
         let square = Self::ternary(&is_nonzero_square, self, &Field::zero());
 
-        // We introduce a variable for the first root we return,
-        // and constrain it to yield, when squared, the square introduced just above.
-        // Thus, if self is a square this is a square root of self; otherwise it is 0, because only 0 yields 0 when squared.
-        // The variable is actually a constant if self is constant, otherwise it is private (even if self is public).
+        // Initialize a new variable for the first root.
         let mode = if self.eject_mode() == Mode::Constant { Mode::Constant } else { Mode::Private };
         let first_root = Field::new(mode, root_witness);
+
+        // Enforce that the first root squared is equal to the square.
+        // Note that if `self` is not a square, then `first_root` and `square` are both zero and the constraint is satisfied.
         E::enforce(|| (&first_root, &first_root, &square));
 
-        // The second root returned by this function is the negation of the first one.
-        // So if self is a non-zero square, this is always different from the first root,
-        // but in the circuit it can be either positive (and the other negative) or vice versa.
+        // Initialize the second root as the negation of the first root.
         let second_root = first_root.clone().neg();
 
         // The error flag is set iff self is a non-square, i.e. it is neither zero nor a non-zero square.

circuit/types/group/src/helpers/from_x_coordinate.rs

@@ -28,8 +28,8 @@ impl<E: Environment> Group<E> {
     }
 
     /// Initializes an affine group element from a given x-coordinate field element.
-    /// Also returns an error flag, set if there is no group element with the given x-coordinate;
-    /// in that case, the returned point is `(0, 0)`, but immaterial.
+    /// Returns an error flag, indicating if there is a group element with the given x-coordinate;
+    /// If the error flag is set, the returned point is `(0, 0)`.
     pub fn from_x_coordinate_flagged(x: Field<E>) -> (Self, Boolean<E>) {
         // Obtain the A and D coefficients of the elliptic curve.
         let a = Field::constant(console::Field::new(E::EDWARDS_A));
@@ -48,50 +48,31 @@ impl<E: Environment> Group<E> {
         let yy: Field<E> = witness!(|a_xx_minus_1, d_xx_minus_1| { a_xx_minus_1 / d_xx_minus_1 });
         E::enforce(|| (&yy, &d_xx_minus_1, &a_xx_minus_1));
 
-        // Compute both square roots of y^2, in no specified order, with a flag saying whether y^2 is a square or not.
-        // That is, finish solving the curve equation for y.
-        // If the x-coordinate line does not intersect the elliptic curve, this returns (1, 0, 0).
+        // Compute both square roots of y^2, with a flag indicating whether y^2 is a square or not.
+        // Note that there is **no** ordering on the square roots.
+        // Note that if the x-coordinate line does not intersect the elliptic curve, this returns (0, 0, false).
         let (y1, y2, yy_is_not_square) = yy.square_roots_flagged_nondeterministic();
 
-        // Form the two points, which are on the curve if yy_is_not_square is false.
-        // Note that the Group<E> type is not restricted to the points in the subgroup or even on the curve;
-        // it includes all possible points, i.e. all possible pairs of field elements.
+        // Construct the two points.
+        // Note that if `yy_is_not_square` is `false`, the points are guaranteed to be on the curve.
+        // Note that the two points are **not** necessarily in the subgroup.
         let point1 = Self { x: x.clone(), y: y1.clone() };
         let point2 = Self { x: x.clone(), y: y2.clone() };
 
-        // We need to check whether either of the two points is in the subgroup.
-        // There may be at most one, but in a circuit we need to always represent both computation paths.
-        // In fact, we represent this computation also when yy_is_not_square is true,
-        // in which case the results of checking whether either point is in the subgroup are meaningless,
-        // but ignored in the final selection of the results returned below.
-        // The criterion for membership in the subgroup is that
-        // multiplying the point by the subgroup order yields the zero point (0, 1).
-        // The group operation that we use here is for the type `Group<E>` of the subgroup,
-        // which as mentioned above it can be performed on points outside the subgroup as well.
-        // We turn the subgroup order into big endian bits,
-        // to get around the issue that the subgroup order is not of Scalar<E> type.
-        let order = E::ScalarField::modulus();
-        let order_bits_be = order.to_bits_be();
-        let mut order_bits_be_constants = Vec::with_capacity(order_bits_be.len());
-        for bit in order_bits_be.iter() {
-            order_bits_be_constants.push(Boolean::constant(*bit));
-        }
-        let point1_times_order = order_bits_be_constants.mul(point1);
-        let point2_times_order = order_bits_be_constants.mul(point2);
-        let point1_is_in_subgroup = point1_times_order.is_zero();
-        let point2_is_in_subgroup = point2_times_order.is_zero();
-
-        // We select y1 if (x, y1) is in the subgroup (which implies that (x, y2) is not in the subgroup),
-        // or y2 if (x, y2) is in the subgroup (which implies that (x, y1) is not in the subgroup),
-        // or 0 if neither is in the subgroup, or x does not even intersect the elliptic curve.
-        // Since at most one of the two points can be in the subgroup, the order of y1 and y2 returned by square root is immaterial:
-        // that nondeterminism (in the circuit) is resolved, and the circuit for from_x_coordinate_flagged is deterministic.
-        let y2_or_zero = Field::ternary(&point2_is_in_subgroup, &y2, &Field::zero());
-        let y1_or_y2_or_zero = Field::ternary(&point1_is_in_subgroup, &y1, &y2_or_zero);
+        // Determine if either of the two points is in the subgroup.
+        // Note that at most **one** of the points can be in the subgroup.
+        let point1_is_in_group = point1.is_in_group();
+        let point2_is_in_group = point2.is_in_group();
+
+        // Select y1 if (x, y1) is in the subgroup.
+        // Otherwise, select y2 if (x, y2) is in the subgroup.
+        // Otherwise, use the zero field element.
+        let y2_or_zero = Field::ternary(&point2_is_in_group, &y2, &Field::zero());
+        let y1_or_y2_or_zero = Field::ternary(&point1_is_in_group, &y1, &y2_or_zero);
         let y = Field::ternary(&yy_is_not_square, &Field::zero(), &y1_or_y2_or_zero);
 
         // The error flag is set iff x does not intersect the elliptic curve or neither intersection point is in the subgroup.
-        let neither_in_subgroup = point1_is_in_subgroup.not().bitand(point2_is_in_subgroup.not());
+        let neither_in_subgroup = point1_is_in_group.not().bitand(point2_is_in_group.not());
         let error_flag = yy_is_not_square.bitor(&neither_in_subgroup);
 
         (Self { x, y }, error_flag)

circuit/types/group/src/lib.rs

@@ -129,6 +129,18 @@ impl<E: Environment> Group<E> {
         // i.e. that it is 4 (= cofactor) times the postulated point on the curve.
         double_point.enforce_double(self);
     }
+
+    /// Returns a `Boolean` indicating if `self` is in the largest prime-order subgroup.
+    pub fn is_in_group(&self) -> Boolean<E> {
+        let order = E::ScalarField::modulus();
+        let order_bits_be = order.to_bits_be();
+        let mut order_bits_be_constants = Vec::with_capacity(order_bits_be.len());
+        for bit in order_bits_be.iter() {
+            order_bits_be_constants.push(Boolean::constant(*bit));
+        }
+        let self_times_order = order_bits_be_constants.mul(self);
+        self_times_order.is_zero()
+    }
 }
 
 #[cfg(console)]
@@ -382,4 +394,32 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_is_in_group() {
+        fn check_is_in_group(mode: Mode, num_constants: u64, num_public: u64, num_private: u64, num_constraints: u64) {
+            let mut rng = TestRng::default();
+
+            for i in 0..ITERATIONS {
+                // Sample a random element.
+                let point: console::Group<<Circuit as Environment>::Network> = Uniform::rand(&mut rng);
+
+                // Inject the x-coordinate.
+                let x_coordinate = Field::new(mode, point.to_x_coordinate());
+
+                // Initialize the group element.
+                let element = Group::<Circuit>::from_x_coordinate(x_coordinate);
+
+                Circuit::scope(format!("{mode} {i}"), || {
+                    let is_in_group = element.is_in_group();
+                    assert!(is_in_group.eject_value());
+                    assert_scope!(num_constants, num_public, num_private, num_constraints);
+                });
+                Circuit::reset();
+            }
+        }
+        check_is_in_group(Mode::Constant, 1752, 0, 0, 0);
+        check_is_in_group(Mode::Public, 750, 0, 2755, 2755);
+        check_is_in_group(Mode::Private, 750, 0, 2755, 2755);
+    }
 }