snarkvm_circuit_types_field/
square_root.rs

1// Copyright (c) 2019-2025 Provable Inc.
2// This file is part of the snarkVM library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use super::*;
17
18impl<E: Environment> SquareRoot for Field<E> {
19    type Output = Self;
20
21    /// Returns the square root of `self`.
22    /// If there are two square roots, the bitwise lesser one is returned.
23    /// If there are no square roots, zero is returned.
24    fn square_root(&self) -> Self::Output {
25        let square_root: Field<E> = witness!(|self| match self.square_root() {
26            Ok(square_root) => square_root,
27            _ => console::Field::zero(),
28        });
29
30        // Ensure `square_root` * `square_root` == `self`.
31        E::enforce(|| (&square_root, &square_root, self));
32
33        // Define the MODULUS_MINUS_ONE_DIV_TWO as a constant.
34        let modulus_minus_one_div_two = match E::BaseField::from_bigint(E::BaseField::modulus_minus_one_div_two()) {
35            Some(modulus_minus_one_div_two) => Field::constant(console::Field::new(modulus_minus_one_div_two)),
36            None => E::halt("Failed to initialize MODULUS_MINUS_ONE_DIV_TWO as a constant"),
37        };
38
39        // Ensure that `square_root` is less than or equal to (MODULUS - 1) / 2.
40        // This ensures that the resulting square root is unique.
41        let is_less_than_or_equal = square_root.is_less_than_or_equal(&modulus_minus_one_div_two);
42        E::assert(is_less_than_or_equal);
43
44        square_root
45    }
46}
47
48impl<E: Environment> Field<E> {
49    /// Returns the square root of `self`, where the least significant bit of the square root is zero.
50    pub fn even_square_root(&self) -> Self {
51        let square_root: Field<E> = witness!(|self| match self.even_square_root() {
52            Ok(square_root) => square_root,
53            _ => console::Field::zero(),
54        });
55
56        // Ensure `square_root` * `square_root` == `self`.
57        E::enforce(|| (&square_root, &square_root, self));
58
59        // Ensure that the LSB of the square root is zero.
60        // Note that this unwrap is safe since the number of bits is always greater than zero.
61        E::assert(!square_root.to_bits_be().last().unwrap());
62
63        square_root
64    }
65}
66
67impl<E: Environment> Field<E> {
68    /// Returns both square roots of `self` and a `Boolean` flag, which is set iff `self` is not a square.
69    ///
70    /// If `self` is a non-zero square,
71    ///  - the first field result is the positive root (i.e. closer to 0)
72    ///  - the second field result is the negative root (i.e. closer to the prime)
73    ///  - the flag is 0
74    ///
75    /// If `self` is 0,
76    ///  - both field results are 0
77    ///  - the flag is 0
78    ///
79    /// If `self` is not a square,
80    ///  - both field results are 0
81    ///  - the flag is 1
82    ///
83    /// Note that the constraints do **not** impose an ordering on the two roots returned by this function;
84    /// this is what the `nondeterministic` part of this function name refers to.
85    pub fn square_roots_flagged_nondeterministic(&self) -> (Self, Self, Boolean<E>) {
86        // Obtain (p-1)/2, as a constant field element.
87        let modulus_minus_one_div_two = match E::BaseField::from_bigint(E::BaseField::modulus_minus_one_div_two()) {
88            Some(modulus_minus_one_div_two) => Field::constant(console::Field::new(modulus_minus_one_div_two)),
89            None => E::halt("Failed to initialize (modulus - 1) / 2"),
90        };
91
92        // Use Euler's criterion: self is a non-zero square iff self^((p-1)/2) is 1.
93        let euler = self.pow(modulus_minus_one_div_two);
94        let is_nonzero_square = euler.is_one();
95
96        // Calculate the witness for the first square result.
97        // Note that the **console** function `square_root` returns the square root closer to 0.
98        let root_witness = match self.eject_value().square_root() {
99            Ok(root) => root,
100            Err(_) => console::Field::zero(),
101        };
102
103        // Initialize the square element, which is either `self` or 0, depending on whether `self` is a square.
104        // This is done to ensure that the below constraint is satisfied even if `self` is not a square.
105        let square = Self::ternary(&is_nonzero_square, self, &Field::zero());
106
107        // Initialize a new variable for the first root.
108        let mode = if self.eject_mode() == Mode::Constant { Mode::Constant } else { Mode::Private };
109        let first_root = Field::new(mode, root_witness);
110
111        // Enforce that the first root squared is equal to the square.
112        // Note that if `self` is not a square, then `first_root` and `square` are both zero and the constraint is satisfied.
113        E::enforce(|| (&first_root, &first_root, &square));
114
115        // Initialize the second root as the negation of the first root.
116        let second_root = first_root.clone().neg();
117
118        // The error flag is set iff self is a non-square, i.e. it is neither zero nor a non-zero square.
119        let is_nonzero = !self.is_zero();
120        let error_flag = is_nonzero.bitand(is_nonzero_square.not());
121
122        (first_root, second_root, error_flag)
123    }
124}
125
126impl<E: Environment> Metrics<dyn SquareRoot<Output = Field<E>>> for Field<E> {
127    type Case = Mode;
128
129    fn count(case: &Self::Case) -> Count {
130        match case.is_constant() {
131            true => Count::is(2, 0, 0, 0),
132            false => Count::is(1, 0, 758, 761),
133        }
134    }
135}
136
137impl<E: Environment> OutputMode<dyn SquareRoot<Output = Field<E>>> for Field<E> {
138    type Case = Mode;
139
140    fn output_mode(case: &Self::Case) -> Mode {
141        match case.is_constant() {
142            true => Mode::Constant,
143            false => Mode::Private,
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use snarkvm_circuit_environment::Circuit;
152
153    const ITERATIONS: u64 = 1_000;
154
155    fn check_square_root(name: &str, mode: Mode, rng: &mut TestRng) {
156        for _ in 0..ITERATIONS {
157            // Sample a random element.
158            let given: console::Field<<Circuit as Environment>::Network> = Uniform::rand(rng);
159            // Compute its square root, or skip this iteration if it does not natively exist.
160            if let Ok(expected) = given.square_root() {
161                let input = Field::<Circuit>::new(mode, given);
162
163                Circuit::scope(name, || {
164                    let candidate = input.square_root();
165                    assert_eq!(expected, candidate.eject_value());
166                    assert_count!(SquareRoot(Field) => Field, &mode);
167                    assert_output_mode!(SquareRoot(Field) => Field, &mode, candidate);
168                });
169                Circuit::reset();
170            }
171        }
172    }
173
174    fn check_even_square_root(
175        name: &str,
176        rng: &mut TestRng,
177        mode: Mode,
178        num_constants: u64,
179        num_public: u64,
180        num_private: u64,
181        num_constraints: u64,
182    ) {
183        for _ in 0..ITERATIONS {
184            // Sample a random element.
185            let given: console::Field<<Circuit as Environment>::Network> = Uniform::rand(rng);
186            // Compute its square root, or skip this iteration if it does not natively exist.
187            if let Ok(expected) = given.even_square_root() {
188                let input = Field::<Circuit>::new(mode, given);
189
190                Circuit::scope(name, || {
191                    let candidate = input.even_square_root();
192                    assert_eq!(expected, candidate.eject_value());
193                    assert_scope!(num_constants, num_public, num_private, num_constraints);
194                });
195                Circuit::reset();
196            }
197        }
198    }
199
200    fn check_square_roots_flagged_nondeterministic(
201        name: &str,
202        mode: Mode,
203        rng: &mut TestRng,
204        num_constants: u64,
205        num_public: u64,
206        num_private: u64,
207        num_constraints: u64,
208    ) {
209        for _ in 0..ITERATIONS {
210            // Sample a random element.
211            let given: console::Field<<Circuit as Environment>::Network> = Uniform::rand(rng);
212            // Compute square roots and error flag in console-land.
213            let (expected_error_flag, expected_positive_root, expected_negative_root) = match given.square_root() {
214                Ok(root) => (false, root, -root),
215                Err(_) => (true, console::Field::zero(), console::Field::zero()),
216            };
217            // Compute square roots and error flag in circuit-land.
218            let input = Field::<Circuit>::new(mode, given);
219            Circuit::scope(name, || {
220                let (candidate_first_root, candidate_second_root, candidate_error_flag) =
221                    input.square_roots_flagged_nondeterministic();
222                // Although the order of the roots is unspecified in the circuit,
223                // the witness values are in a fixed order (first positive, then negative).
224                assert_eq!(expected_error_flag, candidate_error_flag.eject_value());
225                assert_eq!(expected_positive_root, candidate_first_root.eject_value());
226                assert_eq!(expected_negative_root, candidate_second_root.eject_value());
227                assert_scope!(num_constants, num_public, num_private, num_constraints);
228            });
229            Circuit::reset();
230        }
231    }
232
233    #[test]
234    fn test_square_root() {
235        let mut rng = TestRng::default();
236
237        check_square_root("Constant", Mode::Constant, &mut rng);
238        check_square_root("Public", Mode::Public, &mut rng);
239        check_square_root("Private", Mode::Private, &mut rng);
240    }
241
242    #[test]
243    fn test_even_square_root() {
244        let mut rng = TestRng::default();
245
246        check_even_square_root("Constant", &mut rng, Mode::Constant, 254, 0, 0, 0);
247        check_even_square_root("Public", &mut rng, Mode::Public, 0, 0, 506, 509);
248        check_even_square_root("Private", &mut rng, Mode::Private, 0, 0, 506, 509);
249    }
250
251    #[test]
252    fn test_square_roots_flagged_nondeterministic() {
253        let mut rng = TestRng::default();
254
255        check_square_roots_flagged_nondeterministic("Constant", Mode::Constant, &mut rng, 257, 0, 0, 0);
256        check_square_roots_flagged_nondeterministic("Public", Mode::Public, &mut rng, 254, 0, 344, 344);
257        check_square_roots_flagged_nondeterministic("Private", Mode::Private, &mut rng, 254, 0, 344, 344);
258    }
259}