1use anyhow::Result;
  2
  3use plonky2::iop::{
  4    generator::{GeneratedValues, SimpleGenerator},
  5    target::Target,
  6    witness::{PartitionWitness, Witness, WitnessWrite},
  7};
  8use plonky2::field::{
  9    goldilocks_field::GoldilocksField,
 10    types::{Field, PrimeField64},
 11};
 12use plonky2::plonk::{
 13    circuit_builder::CircuitBuilder,
 14    circuit_data::CommonCircuitData,
 15};
 16use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
 17
 18use crate::schnorr::{SchnorrPublicKey, SchnorrSignature};
 19
 20type GoldF = GoldilocksField;
 21
 22
 23
 24#[derive(Debug, Default)]
 25pub struct Mod65537Generator {
 26    a: Target,
 27    q: Target,
 28    r: Target,
 29}
 30
 31impl SimpleGenerator<GoldF, 2> for Mod65537Generator {
 32    fn id(&self) -> String {
 33        "Mod65537Generator".to_string()
 34    }
 35    fn dependencies(&self) -> Vec<Target> {
 36        vec![self.a]
 37    }
 38
 39    fn run_once(
 40        &self,
 41        witness: &PartitionWitness<GoldF>,
 42        out_buffer: &mut GeneratedValues<GoldF>,
 43    ) -> Result<()> {
 44        let a = witness.get_target(self.a);
 45        let a64 = a.to_canonical_u64();
 46        let q64 = a64 / 65537;
 47        let r64 = a64 % 65537;
 48
 49        out_buffer.set_target(self.q, GoldF::from_canonical_u64(q64));
 50        out_buffer.set_target(self.r, GoldF::from_canonical_u64(r64));
 51
 52        Ok(())
 53    }
 54
 55    fn serialize(&self, dst: &mut Vec<u8>, common_data: &CommonCircuitData<GoldF, 2>) -> IoResult<()> {
 56        dst.write_target(self.a)?;
 57        dst.write_target(self.q)?;
 58        dst.write_target(self.r)?;
 59        Ok(())
 60    }
 61
 62    fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData<GoldF, 2>) -> IoResult<Self>
 63    where
 64        Self: Sized 
 65    {
 66            let a = src.read_target()?;
 67            let q = src.read_target()?;
 68            let r = src.read_target()?;
 69            Ok(Self { a, q, r })
 70    }
 71}
 72
 73pub struct Mod65537Builder {}
 74
 75impl Mod65537Builder {
 76    // Reduce a modulo the constant 65537
 77    // where a is the canonical representative for an element of the field
 78    // (meaning: 0 \leq a < p)
 79
 80    // To prove this, write
 81    // a = 65537 * q + r, and do range checks to check that:
 82    // 0 <= q <= floor(p / 65537)
 83    // 0 <= r < 65537
 84    // (these first two checks guarantee that a lies in the range [0, p + 65536])
 85    // if q = floor(p / 65537) then r = 0
 86    // (note that p % 65537 == 1 so this is the only possibility)
 87    pub(crate) fn mod_65537 (
 88        builder: &mut CircuitBuilder::<GoldF, 2>,
 89        a: Target,
 90    ) -> Target {
 91        let q = builder.add_virtual_target();
 92        let r = builder.add_virtual_target();
 93
 94        // the Mod65537Generator will assign values to q and r later
 95        builder.add_simple_generator( Mod65537Generator { a, q, r } );
 96
 97        // impose four constraints
 98        // 1. a = 65537 * q + r
 99        let t65537 = builder.constant(GoldF::from_canonical_u64(65537));
100        let a_copy = builder.mul_add(t65537, q, r);
101        builder.connect(a, a_copy);
102
103        // 2. 0 <= q <= floor(p / 65537)
104        // max_q is 281470681743360 = floor(p / 65537) = (p-1) / 65537 = 2^48 - 2^32
105        let max_q = builder.constant(GoldF::from_canonical_u64(281470681743360));
106        builder.range_check(q, 48);
107        let diff_q = builder.sub(max_q, q);
108        builder.range_check(diff_q, 48);
109
110        // 3. 0 <= r < 65537
111        let max_r = builder.constant(GoldF::from_canonical_u64(65537));
112        builder.range_check(r, 17);
113        let diff_r = builder.sub(max_r, r);
114        builder.range_check(diff_r, 17);
115
116        // 4. if q = floor(p / 65537) then r = 0
117        let q_equals_max = builder.is_equal(q, max_q);
118        let prod_temp = builder.mul(q_equals_max.target, r);
119        let zero_temp = builder.zero();
120        builder.connect(prod_temp, zero_temp);
121
122        // throw in the Generator to tell builder how to compute r
123        builder.add_simple_generator( Mod65537Generator {a, q, r} );
124
125        r
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use crate::mod65537::Mod65537Builder;
132    use plonky2::iop::{
133        target::Target,
134        witness::PartialWitness,
135    };
136    use plonky2::plonk::{
137        circuit_builder::CircuitBuilder,
138        circuit_data::CircuitConfig,
139        config::{GenericConfig, PoseidonGoldilocksConfig},
140    };
141    use plonky2::field::{
142        goldilocks_field::GoldilocksField,
143        types::Field,
144    };
145
146    #[test]
147    fn test_mod65537() -> () {
148        const D: usize = 2;
149        const p: u64 = 18446744069414584321;  // the Goldilocks prime
150        type C = PoseidonGoldilocksConfig;
151        type F = <C as GenericConfig<D>>::F;
152
153        let config = CircuitConfig::standard_recursion_config();
154        let mut builder = CircuitBuilder::<F, D>::new(config);
155
156        let a64: Vec<u64> = vec![0, 1, 2, 65535, 65536, 65537, p - 4, p - 3, p - 2, p - 1];
157
158        let a: Vec<Target> = a64
159            .iter()
160            .map(|x| builder.constant(GoldilocksField::from_canonical_u64(*x)))
161            .collect();
162
163        let r: Vec<Target> = a.iter()
164            .map(|targ| Mod65537Builder::mod_65537(&mut builder, *targ))
165            .collect();
166
167        // check that the outputs are correct,
168        // obviously you don't need this in your own code
169        let r_expected64: Vec<u64> = a64.iter().map(|x| x % 65537).collect();
170        println!("Expected residues mod 64: {:?}", r_expected64);
171        let r_expected: Vec<Target> = r_expected64.iter()
172            .map(|x| builder.constant(GoldilocksField::from_canonical_u64(*x)))
173            .collect();
174        r.iter().zip(r_expected.iter())
175            .for_each(|(x, y)| builder.connect(*x, *y));
176        
177        let mut pw: PartialWitness<F> = PartialWitness::new();
178        let data = builder.build::<C>();
179        let proof = data.prove(pw).unwrap();
180
181        ()
182    }
183}