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