From 880bc87bb13d1efa19506566b88438b7ce3e677b Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Mon, 26 Sep 2022 10:43:18 -0700
Subject: [PATCH] sqrt

---
 field/src/goldilocks_field.rs   | 55 +++++++++++++++++++++++++++++++++
 plonky2/examples/square_root.rs | 30 +++++++++++-------
 2 files changed, 73 insertions(+), 12 deletions(-)

diff --git a/field/src/goldilocks_field.rs b/field/src/goldilocks_field.rs
index c1bb60b0..21235e63 100644
--- a/field/src/goldilocks_field.rs
+++ b/field/src/goldilocks_field.rs
@@ -280,6 +280,61 @@ impl DivAssign for GoldilocksField {
     }
 }
 
+impl GoldilocksField {
+    pub fn is_quadratic_residue(&self) -> bool {
+        if self.is_zero() {
+            return true;
+        }
+        // This is based on Euler's criterion.
+        let power = Self::NEG_ONE.to_canonical_biguint() / 2u8;
+        let exp = self.exp_biguint(&power);
+        if exp == Self::ONE {
+            return true;
+        }
+        if exp == Self::NEG_ONE {
+            return false;
+        }
+        panic!("Unreachable")
+    }
+
+    pub fn sqrt(&self) -> Option<Self> {
+        if self.is_zero() {
+            Some(*self)
+        } else if self.is_quadratic_residue() {
+            let t = (Self::order() - BigUint::from(1u32)) / (BigUint::from(2u32).pow(Self::TWO_ADICITY as u32));
+            let mut z = Self::POWER_OF_TWO_GENERATOR.exp_biguint(&t);
+            let mut w = self.exp_biguint(&((t - BigUint::from(1u32)) / BigUint::from(2u32)));
+            let mut x = w * *self;
+            let mut b = x * w;
+
+            let mut v = Self::TWO_ADICITY as usize;
+
+            while !b.is_one() {
+                let mut k = 0usize;
+                let mut b2k = b;
+                while !b2k.is_one() {
+                    b2k = b2k * b2k;
+                    k += 1;
+                }
+                let j = v - k - 1;
+                w = z;
+                for _ in 0..j {
+                    w = w * w;
+                }
+
+                z = w * w;
+                b = b * z;
+                x = x * w;
+                v = k;
+            }
+            Some(x)
+        } else {
+            None
+        }
+    }
+
+}
+
 /// Fast addition modulo ORDER for x86-64.
 /// This function is marked unsafe for the following reasons:
 ///   - It is only correct if x + y < 2**64 + ORDER = 0x1ffffffff00000001.
diff --git a/plonky2/examples/square_root.rs b/plonky2/examples/square_root.rs
index 0aab8283..5cbdf02e 100644
--- a/plonky2/examples/square_root.rs
+++ b/plonky2/examples/square_root.rs
@@ -10,6 +10,7 @@ use plonky2::plonk::circuit_builder::CircuitBuilder;
 use plonky2::plonk::circuit_data::CircuitConfig;
 use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
 use plonky2_field::extension::Extendable;
+use plonky2_field::goldilocks_field::GoldilocksField;
 
 #[derive(Debug)]
 struct SquareRootGenerator<F: RichField + Extendable<D>, const D: usize> {
@@ -18,15 +19,14 @@ struct SquareRootGenerator<F: RichField + Extendable<D>, const D: usize> {
     _phantom: PhantomData<F>,
 }
 
-impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F> for SquareRootGenerator<F, D> {
+impl SimpleGenerator<GoldilocksField> for SquareRootGenerator<GoldilocksField, 2> {
     fn dependencies(&self) -> Vec<Target> {
         vec![self.x_squared]
     }
 
-    fn run_once(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) {
-        let x_squared = witness.get_target(self.x);
-        let s = F::from_canonical_u32(4294967295);
-        let x = 
+    fn run_once(&self, witness: &PartitionWitness<GoldilocksField>, out_buffer: &mut GeneratedValues<GoldilocksField>) {
+        let x_squared = witness.get_target(self.x_squared);
+        let x = x_squared.sqrt().unwrap();
 
         out_buffer.set_target(self.x, x);
     }
@@ -48,16 +48,22 @@ fn main() -> Result<()> {
     builder.register_public_input(x);
     builder.register_public_input(x_squared);
 
-    // builder.add_simple_generator(SquareGenerator::<F, D> {
-    //     x,
-    //     x_squared,
-    //     _phantom: PhantomData,
-    // });
+    builder.add_simple_generator(SquareRootGenerator::<F, D> {
+        x,
+        x_squared,
+        _phantom: PhantomData,
+    });
 
-    let x_value = F::rand();
+    let x_squared_value = {
+        let mut val = F::rand();
+        while !val.is_quadratic_residue() {
+            val = F::rand();
+        }
+        val
+    };
 
     let mut pw = PartialWitness::new();
-    pw.set_target(x, x_value);
+    pw.set_target(x_squared, x_squared_value);
 
     let data = builder.build::<C>();
     let proof = data.prove(pw)?;