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