1use anyhow::Result;
  2use plonky2::iop::generator::{GeneratedValues, SimpleGenerator};
  3use plonky2::iop::target::{BoolTarget, Target};
  4use plonky2::iop::witness::{PartialWitness, PartitionWitness, Witness, WitnessWrite};
  5use plonky2::field::extension::{Extendable, FieldExtension};
  6use plonky2::field::types::{Field, PrimeField64};
  7use plonky2::field::goldilocks_field::GoldilocksField;
  8use plonky2::hash::hash_types::RichField;
  9use plonky2::hash::poseidon::PoseidonHash;
 10use plonky2::plonk::circuit_builder::CircuitBuilder;
 11use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, VerifierCircuitData, VerifierOnlyCircuitData};
 12use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
 13use plonky2::plonk::proof::ProofWithPublicInputs;
 14use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
 15
 16use crate::schnorr::{SchnorrPublicKey, SchnorrSignature};
 17
 18type GoldF = GoldilocksField;
 19
 20pub struct MessageTarget {
 21    msg: Vec<Target>,
 22}
 23
 24impl MessageTarget {
 25    fn new_with_size(builder: &mut CircuitBuilder<GoldF, 2>, n: usize) -> Self {
 26        Self {
 27            msg: builder.add_virtual_targets(n),
 28        }
 29    }
 30
 31    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, msg: &Vec<GoldF>) -> Result<()> {
 32        assert!(msg.len() == self.msg.len());
 33        for (&t, &x) in self.msg.iter().zip(msg.iter()) {
 34            pw.set_target(t, x)?;
 35        }
 36
 37        Ok(())
 38    }
 39}
 40
 41pub struct SchnorrSignatureTarget {
 42    s: Target,
 43    e: Target,
 44}
 45
 46impl SchnorrSignatureTarget {
 47    fn new_virtual(builder: &mut CircuitBuilder<GoldF, 2>) -> Self {
 48        let s = builder.add_virtual_target();
 49        let e = builder.add_virtual_target();
 50        Self{ s, e }
 51    }
 52
 53    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, sig: &SchnorrSignature) -> Result<()> {
 54        pw.set_target(self.s, GoldilocksField::from_canonical_u64(sig.s))?;
 55        pw.set_target(self.e, GoldilocksField::from_canonical_u64(sig.e))?;
 56        Ok(())
 57    }
 58}
 59
 60pub struct SchnorrPublicKeyTarget {
 61    pk: Target,
 62}
 63
 64impl SchnorrPublicKeyTarget {
 65    fn new_virtual(builder: &mut CircuitBuilder<GoldF, 2>) -> Self {
 66        Self{ pk: builder.add_virtual_target() }
 67    }
 68
 69    fn set_witness(&self, pw: &mut PartialWitness<GoldF>, pk: &SchnorrPublicKey) -> Result<()> {
 70        pw.set_target(self.pk, pk.pk)?;
 71        Ok(())
 72    }
 73}
 74
 75#[derive(Debug, Default)]
 76pub struct Mod65537Generator {
 77    a: Target,
 78    q: Target,
 79    r: Target,
 80}
 81
 82impl SimpleGenerator<GoldF, 2> for Mod65537Generator {
 83    fn id(&self) -> String {
 84        "Mod65537Generator".to_string()
 85    }
 86    fn dependencies(&self) -> Vec<Target> {
 87        vec![self.a]
 88    }
 89
 90    fn run_once(
 91        &self,
 92        witness: &PartitionWitness<GoldF>,
 93        out_buffer: &mut GeneratedValues<GoldF>,
 94    ) -> Result<()> {
 95        let a = witness.get_target(self.a);
 96        let a64 = a.to_canonical_u64();
 97        let q64 = a64 / 65537;
 98        let r64 = a64 % 65537;
 99
100        out_buffer.set_target(self.q, GoldF::from_canonical_u64(q64));
101        out_buffer.set_target(self.r, GoldF::from_canonical_u64(r64));
102
103        Ok(())
104    }
105
106    fn serialize(&self, dst: &mut Vec<u8>, common_data: &CommonCircuitData<GoldF, 2>) -> IoResult<()> {
107        println!("SERIALIZATION!  What is this good for?");
108        dst.write_target(self.a)?;
109        dst.write_target(self.q)?;
110        dst.write_target(self.r)?;
111        Ok(())
112    }
113
114    fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData<GoldF, 2>) -> IoResult<Self>
115    where
116        Self: Sized 
117    {
118            println!("DESERIALIZATION!  What is this good for?");
119            let a = src.read_target()?;
120            let q = src.read_target()?;
121            let r = src.read_target()?;
122            Ok(Self { a, q, r })
123    }
124}
125
126pub struct SchnorrBuilder {}
127
128impl SchnorrBuilder {
129    // Reduce a modulo the constant 65537
130    // where a is the canonical representative for an element of the field
131    // (meaning: 0 \leq a < p)
132
133    // To verify this, write
134    // a = 65537 * q + r, and do range checks to check that:
135    // 0 <= q <= floor(p / 65537)
136    // 0 <= r < 65537
137    // (these first two checks guarantee that a lies in the range [0, p + 65536])
138    // if q = floor(p / 65537) then r = 0
139    // (note that p % 65537 == 1 so this is the only possibility)
140    pub(crate) fn mod_65537 <
141        //C: GenericConfig<2, F = GoldF>,
142    > (
143        builder: &mut CircuitBuilder::<GoldF, 2>,
144        a: Target,
145    ) -> Target {
146        let q = builder.add_virtual_target();
147        let r = builder.add_virtual_target();
148
149        // the Mod65537Generator will assign values to q and r later
150        builder.add_simple_generator( Mod65537Generator { a, q, r } );
151
152        // impose four constraints
153        // 1. a = 65537 * q + r
154        let t65537 = builder.constant(GoldF::from_canonical_u64(65537));
155        let a_copy = builder.mul_add(t65537, q, r);
156        builder.connect(a, a_copy);
157
158        // 2. 0 <= q <= floor(p / 65537)
159        // max_q is 281470681743360 = floor(p / 65537) = (p-1) / 65537 = 2^48 - 2^32
160        let max_q = builder.constant(GoldF::from_canonical_u64(281470681743360));
161        builder.range_check(q, 48);
162        let diff_q = builder.sub(max_q, q);
163        builder.range_check(diff_q, 48);
164
165        // 3. 0 <= r < 65537
166        let max_r = builder.constant(GoldF::from_canonical_u64(65537));
167        builder.range_check(r, 17);
168        let diff_r = builder.sub(max_r, r);
169        builder.range_check(diff_r, 17);
170
171        // 4. if q = floor(p / 65537) then r = 0
172        let q_equals_max = builder.is_equal(q, max_q);
173        let prod_temp = builder.mul(q_equals_max.target, r);
174        let zero_temp = builder.zero();
175        builder.connect(prod_temp, zero_temp);
176
177        // throw in the Generator to tell builder how to compute r
178        builder.add_simple_generator( Mod65537Generator {a, q, r} );
179
180        r
181    }
182
183    pub fn constrain_sig <
184        C: GenericConfig<2, F = GoldF>,
185    > (
186        &self,
187        builder: &mut CircuitBuilder::<GoldF, 2>,
188        sig: &SchnorrSignatureTarget,
189        msg: &MessageTarget,
190        pk: &SchnorrPublicKeyTarget,
191    ) -> () {
192        println!("WARNING constrain_sig() is not done yet DONT USE IT");
193
194        let PRIME_GROUP_GEN: Target = builder.constant(GoldF::from_canonical_u64(6612579038192137166));
195        let PRIME_GROUP_ORDER: Target = builder.constant(GoldF::from_canonical_u64(65537));
196        const num_bits_exp: usize = 32;
197
198        /*
199        let r: GoldF = Self::pow(self.PRIME_GROUP_GEN, sig.s)
200            * Self::pow(pk.pk, sig.e);
201        let e_v: u64 = self.hash_insecure(&r, msg);
202        e_v == sig.e   */
203
204        let gs: Target = builder.exp(PRIME_GROUP_GEN, sig.s, num_bits_exp);
205        let pe: Target = builder.exp(pk.pk, sig.e, num_bits_exp);
206        let r: Target = builder.mul(gs, pe);
207
208        // compute hash
209        // note that it's safe to clone Targets since they just contain indices
210        let hash_input: Vec<Target> = std::iter::once(r)
211            .chain(msg.msg.iter().cloned()) 
212            .collect();
213        let hash_output: Target = builder.hash_n_to_hash_no_pad::<PoseidonHash>(
214            hash_input,
215        ).elements[0]; // whoops have to take mod group order;
216
217        let e: Target = Self::mod_65537(builder, hash_output);
218
219        // enforce equality
220        builder.connect(e, sig.e);
221    }
222}
223
224#[cfg(test)]
225mod tests{
226    use crate::schnorr::{SchnorrPublicKey, SchnorrSecretKey, SchnorrSigner, SchnorrSignature};
227    use crate::schnorr_prover::{MessageTarget, SchnorrBuilder, SchnorrPublicKeyTarget, SchnorrSignatureTarget};
228    use plonky2::hash::poseidon::Poseidon;
229    use plonky2::iop::target::Target;
230    use plonky2::iop::witness::{PartialWitness, PartitionWitness, Witness, WitnessWrite};
231    use plonky2::plonk::circuit_builder::CircuitBuilder;
232    use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, VerifierCircuitData, VerifierOnlyCircuitData};
233    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
234    use plonky2::field::goldilocks_field::GoldilocksField;
235    use plonky2::field::types::Field;
236    use rand;
237
238    #[test]
239    fn test_mod65537() -> () {
240        const D: usize = 2;
241        const p: u64 = 18446744069414584321;  // the Goldilocks prime
242        type C = PoseidonGoldilocksConfig;
243        type F = <C as GenericConfig<D>>::F;
244
245        let config = CircuitConfig::standard_recursion_config();
246        let mut builder = CircuitBuilder::<F, D>::new(config);
247
248        let a64: Vec<u64> = vec![0, 1, 2, 65535, 65536, 65537, p - 4, p - 3, p - 2, p - 1];
249
250        let a: Vec<Target> = a64
251            .iter()
252            .map(|x| builder.constant(GoldilocksField::from_canonical_u64(*x)))
253            .collect();
254
255        let r: Vec<Target> = a.iter()
256            .map(|targ| SchnorrBuilder::mod_65537(&mut builder, *targ))
257            .collect();
258
259        // check that the outputs are correct,
260        // obviously you don't need this in your own code
261        let r_expected64: Vec<u64> = a64.iter().map(|x| x % 65537).collect();
262        println!("Expected residues mod 64: {:?}", r_expected64);
263        let r_expected: Vec<Target> = r_expected64.iter()
264            .map(|x| builder.constant(GoldilocksField::from_canonical_u64(*x)))
265            .collect();
266
267        r.iter().zip(r_expected.iter())
268            .for_each(|(x, y)| builder.connect(*x, *y));
269        
270        let mut pw: PartialWitness<F> = PartialWitness::new();
271        let data = builder.build::<C>();
272        let proof = data.prove(pw).unwrap();
273
274        ()
275    }
276
277    #[test]
278    fn test_schnorr() -> () {
279        const D: usize = 2;
280        type C = PoseidonGoldilocksConfig;
281        type F = <C as GenericConfig<D>>::F;
282
283        let mut rng: rand::rngs::ThreadRng = rand::thread_rng();
284
285        let config = CircuitConfig::standard_recursion_config();
286        let mut builder = CircuitBuilder::<F, D>::new(config);
287
288        let sb: SchnorrBuilder = SchnorrBuilder{};
289
290        // create keypair, message, signature
291        let sk: SchnorrSecretKey = SchnorrSecretKey{ sk: 133 };
292        let ss = SchnorrSigner::new();
293        let pk: SchnorrPublicKey = ss.keygen(&sk);
294        let msg: Vec<GoldilocksField> = ss.u64_into_goldilocks_vec(
295            vec![1500, 1600, 0, 0, 0]
296        );
297        let msg_size: usize = msg.len();
298        let sig: SchnorrSignature = ss.sign(&msg, &sk, &mut rng);
299
300        let pk_targ = SchnorrPublicKeyTarget::new_virtual(&mut builder);
301        let sig_targ = SchnorrSignatureTarget::new_virtual(&mut builder);
302        let msg_targ = MessageTarget::new_with_size(&mut builder, msg_size);
303
304        
305        // instead of verifying we're going to prove the verification
306        sb.constrain_sig::<PoseidonGoldilocksConfig> (
307            &mut builder,
308            &sig_targ,
309            &msg_targ,
310            &pk_targ
311        );
312
313        // assign witnesses for verification
314        let mut pw: PartialWitness<F> = PartialWitness::new();
315        pk_targ.set_witness(&mut pw, &pk);
316        sig_targ.set_witness(&mut pw, &sig);
317        msg_targ.set_witness(&mut pw, &msg);
318
319
320        let data = builder.build::<C>();
321        let proof = data.prove(pw).unwrap();
322    }
323}