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