1use alloc::string::{String, ToString};
  2use alloc::vec::Vec;
  3use alloc::{format, vec};
  4use anyhow::Result;
  5use core::marker::PhantomData;
  6use num::integer::div_ceil;
  7use plonky2::plonk::circuit_data::CommonCircuitData;
  8use plonky2::util::serialization::{Buffer, IoResult, Read, Write};
  9
 10use plonky2::field::extension::Extendable;
 11use plonky2::field::types::Field;
 12use plonky2::gates::gate::Gate;
 13use plonky2::gates::util::StridedConstraintConsumer;
 14use plonky2::hash::hash_types::RichField;
 15use plonky2::iop::ext_target::ExtensionTarget;
 16use plonky2::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGeneratorRef};
 17use plonky2::iop::target::Target;
 18use plonky2::iop::witness::{PartitionWitness, Witness, WitnessWrite};
 19use plonky2::plonk::circuit_builder::CircuitBuilder;
 20use plonky2::plonk::plonk_common::{reduce_with_powers, reduce_with_powers_ext_circuit};
 21use plonky2::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
 22
 23/// A gate which can decompose a number into base B little-endian limbs.
 24#[derive(Copy, Clone, Debug)]
 25pub struct U32RangeCheckGate<F: RichField + Extendable<D>, const D: usize> {
 26    pub num_input_limbs: usize,
 27    _phantom: PhantomData<F>,
 28}
 29
 30impl<F: RichField + Extendable<D>, const D: usize> U32RangeCheckGate<F, D> {
 31    pub fn new(num_input_limbs: usize) -> Self {
 32        Self {
 33            num_input_limbs,
 34            _phantom: PhantomData,
 35        }
 36    }
 37
 38    pub const AUX_LIMB_BITS: usize = 2;
 39    pub const BASE: usize = 1 << Self::AUX_LIMB_BITS;
 40
 41    fn aux_limbs_per_input_limb(&self) -> usize {
 42        div_ceil(32, Self::AUX_LIMB_BITS)
 43    }
 44    pub fn wire_ith_input_limb(&self, i: usize) -> usize {
 45        debug_assert!(i < self.num_input_limbs);
 46        i
 47    }
 48    pub fn wire_ith_input_limb_jth_aux_limb(&self, i: usize, j: usize) -> usize {
 49        debug_assert!(i < self.num_input_limbs);
 50        debug_assert!(j < self.aux_limbs_per_input_limb());
 51        self.num_input_limbs + self.aux_limbs_per_input_limb() * i + j
 52    }
 53}
 54
 55impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for U32RangeCheckGate<F, D> {
 56    fn id(&self) -> String {
 57        format!("{self:?}")
 58    }
 59
 60    fn serialize(&self, dst: &mut Vec<u8>, _common_data: &CommonCircuitData<F, D>) -> IoResult<()> {
 61        dst.write_usize(self.num_input_limbs)
 62    }
 63
 64    fn deserialize(src: &mut Buffer, _common_data: &CommonCircuitData<F, D>) -> IoResult<Self> {
 65        let num_input_limbs = src.read_usize()?;
 66        Ok(Self {
 67            num_input_limbs,
 68            _phantom: PhantomData,
 69        })
 70    }
 71
 72    fn eval_unfiltered(&self, vars: EvaluationVars<F, D>) -> Vec<F::Extension> {
 73        let mut constraints = Vec::with_capacity(self.num_constraints());
 74
 75        let base = F::Extension::from_canonical_usize(Self::BASE);
 76        for i in 0..self.num_input_limbs {
 77            let input_limb = vars.local_wires[self.wire_ith_input_limb(i)];
 78            let aux_limbs: Vec<_> = (0..self.aux_limbs_per_input_limb())
 79                .map(|j| vars.local_wires[self.wire_ith_input_limb_jth_aux_limb(i, j)])
 80                .collect();
 81            let computed_sum = reduce_with_powers(&aux_limbs, base);
 82
 83            constraints.push(computed_sum - input_limb);
 84            for aux_limb in aux_limbs {
 85                constraints.push(
 86                    (0..Self::BASE)
 87                        .map(|i| aux_limb - F::Extension::from_canonical_usize(i))
 88                        .product(),
 89                );
 90            }
 91        }
 92
 93        constraints
 94    }
 95
 96    fn eval_unfiltered_base_one(
 97        &self,
 98        vars: EvaluationVarsBase<F>,
 99        mut yield_constr: StridedConstraintConsumer<F>,
100    ) {
101        let base = F::from_canonical_usize(Self::BASE);
102        for i in 0..self.num_input_limbs {
103            let input_limb = vars.local_wires[self.wire_ith_input_limb(i)];
104            let aux_limbs: Vec<_> = (0..self.aux_limbs_per_input_limb())
105                .map(|j| vars.local_wires[self.wire_ith_input_limb_jth_aux_limb(i, j)])
106                .collect();
107            let computed_sum = reduce_with_powers(&aux_limbs, base);
108
109            yield_constr.one(computed_sum - input_limb);
110            for aux_limb in aux_limbs {
111                yield_constr.one(
112                    (0..Self::BASE)
113                        .map(|i| aux_limb - F::from_canonical_usize(i))
114                        .product(),
115                );
116            }
117        }
118    }
119
120    fn eval_unfiltered_circuit(
121        &self,
122        builder: &mut CircuitBuilder<F, D>,
123        vars: EvaluationTargets<D>,
124    ) -> Vec<ExtensionTarget<D>> {
125        let mut constraints = Vec::with_capacity(self.num_constraints());
126
127        let base = builder.constant(F::from_canonical_usize(Self::BASE));
128        for i in 0..self.num_input_limbs {
129            let input_limb = vars.local_wires[self.wire_ith_input_limb(i)];
130            let aux_limbs: Vec<_> = (0..self.aux_limbs_per_input_limb())
131                .map(|j| vars.local_wires[self.wire_ith_input_limb_jth_aux_limb(i, j)])
132                .collect();
133            let computed_sum = reduce_with_powers_ext_circuit(builder, &aux_limbs, base);
134
135            constraints.push(builder.sub_extension(computed_sum, input_limb));
136            for aux_limb in aux_limbs {
137                constraints.push({
138                    let mut acc = builder.one_extension();
139                    (0..Self::BASE).for_each(|i| {
140                        // We update our accumulator as:
141                        // acc' = acc (x - i)
142                        //      = acc x + (-i) acc
143                        // Since -i is constant, we can do this in one arithmetic_extension call.
144                        let neg_i = -F::from_canonical_usize(i);
145                        acc = builder.arithmetic_extension(F::ONE, neg_i, acc, aux_limb, acc)
146                    });
147                    acc
148                });
149            }
150        }
151
152        constraints
153    }
154
155    fn generators(&self, row: usize, _local_constants: &[F]) -> Vec<WitnessGeneratorRef<F, D>> {
156        let gen = U32RangeCheckGenerator { gate: *self, row };
157        vec![WitnessGeneratorRef::new(gen.adapter())]
158    }
159
160    fn num_wires(&self) -> usize {
161        self.num_input_limbs * (1 + self.aux_limbs_per_input_limb())
162    }
163
164    fn num_constants(&self) -> usize {
165        0
166    }
167
168    // Bounded by the range-check (x-0)*(x-1)*...*(x-BASE+1).
169    fn degree(&self) -> usize {
170        Self::BASE
171    }
172
173    // 1 for checking the each sum of aux limbs, plus a range check for each aux limb.
174    fn num_constraints(&self) -> usize {
175        self.num_input_limbs * (1 + self.aux_limbs_per_input_limb())
176    }
177}
178
179#[derive(Debug)]
180pub struct U32RangeCheckGenerator<F: RichField + Extendable<D>, const D: usize> {
181    gate: U32RangeCheckGate<F, D>,
182    row: usize,
183}
184
185impl<F: RichField + Extendable<D>, const D: usize> SimpleGenerator<F, D>
186    for U32RangeCheckGenerator<F, D>
187{
188    fn id(&self) -> String {
189        "U32RangeCheckGenerator".to_string()
190    }
191
192    fn dependencies(&self) -> Vec<Target> {
193        let num_input_limbs = self.gate.num_input_limbs;
194        (0..num_input_limbs)
195            .map(|i| Target::wire(self.row, self.gate.wire_ith_input_limb(i)))
196            .collect()
197    }
198
199    fn run_once(
200        &self,
201        witness: &PartitionWitness<F>,
202        out_buffer: &mut GeneratedValues<F>,
203    ) -> Result<()> {
204        let num_input_limbs = self.gate.num_input_limbs;
205        for i in 0..num_input_limbs {
206            let sum_value = witness
207                .get_target(Target::wire(self.row, self.gate.wire_ith_input_limb(i)))
208                .to_canonical_u64() as u32;
209
210            let base = U32RangeCheckGate::<F, D>::BASE as u32;
211            let limbs = (0..self.gate.aux_limbs_per_input_limb())
212                .map(|j| Target::wire(self.row, self.gate.wire_ith_input_limb_jth_aux_limb(i, j)));
213            let limbs_value = (0..self.gate.aux_limbs_per_input_limb())
214                .scan(sum_value, |acc, _| {
215                    let tmp = *acc % base;
216                    *acc /= base;
217                    Some(F::from_canonical_u32(tmp))
218                })
219                .collect::<Vec<_>>();
220
221            for (b, b_value) in limbs.zip(limbs_value) {
222                out_buffer.set_target(b, b_value)?;
223            }
224        }
225
226        Ok(())
227    }
228
229    fn serialize(&self, dst: &mut Vec<u8>, common_data: &CommonCircuitData<F, D>) -> IoResult<()> {
230        self.gate.serialize(dst, common_data)?;
231        dst.write_usize(self.row)
232    }
233
234    fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData<F, D>) -> IoResult<Self> {
235        let gate = U32RangeCheckGate::deserialize(src, common_data)?;
236        let row = src.read_usize()?;
237        Ok(Self { row, gate })
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use anyhow::Result;
244    use itertools::unfold;
245    use plonky2::field::extension::quartic::QuarticExtension;
246    use plonky2::field::goldilocks_field::GoldilocksField;
247    use plonky2::field::types::{Field, Sample};
248    use plonky2::gates::gate_testing::{test_eval_fns, test_low_degree};
249    use plonky2::hash::hash_types::HashOut;
250    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
251    use rand::rngs::OsRng;
252    use rand::Rng;
253
254    use super::*;
255
256    #[test]
257    fn low_degree() {
258        test_low_degree::<GoldilocksField, _, 4>(U32RangeCheckGate::new(8))
259    }
260
261    #[test]
262    fn eval_fns() -> Result<()> {
263        const D: usize = 2;
264        type C = PoseidonGoldilocksConfig;
265        type F = <C as GenericConfig<D>>::F;
266        test_eval_fns::<F, C, _, D>(U32RangeCheckGate::new(8))
267    }
268
269    fn test_gate_constraint(input_limbs: Vec<u64>) {
270        type F = GoldilocksField;
271        type FF = QuarticExtension<GoldilocksField>;
272        const D: usize = 4;
273        const AUX_LIMB_BITS: usize = 2;
274        const BASE: usize = 1 << AUX_LIMB_BITS;
275        const AUX_LIMBS_PER_INPUT_LIMB: usize = (32 + AUX_LIMB_BITS - 1) / AUX_LIMB_BITS;
276
277        fn get_wires(input_limbs: Vec<u64>) -> Vec<FF> {
278            let num_input_limbs = input_limbs.len();
279            let mut v = Vec::new();
280
281            for i in 0..num_input_limbs {
282                let input_limb = input_limbs[i];
283
284                let split_to_limbs = |mut val, num| {
285                    unfold((), move |_| {
286                        let ret = val % (BASE as u64);
287                        val /= BASE as u64;
288                        Some(ret)
289                    })
290                    .take(num)
291                    .map(F::from_canonical_u64)
292                };
293
294                let mut aux_limbs: Vec<_> =
295                    split_to_limbs(input_limb, AUX_LIMBS_PER_INPUT_LIMB).collect();
296
297                v.append(&mut aux_limbs);
298            }
299
300            input_limbs
301                .iter()
302                .cloned()
303                .map(F::from_canonical_u64)
304                .chain(v.iter().cloned())
305                .map(|x| x.into())
306                .collect()
307        }
308
309        let gate = U32RangeCheckGate::<F, D> {
310            num_input_limbs: 8,
311            _phantom: PhantomData,
312        };
313
314        let vars = EvaluationVars {
315            local_constants: &[],
316            local_wires: &get_wires(input_limbs),
317            public_inputs_hash: &HashOut::rand(),
318        };
319
320        assert!(
321            gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()),
322            "Gate constraints are not satisfied."
323        );
324    }
325
326    #[test]
327    fn test_gate_constraint_good() {
328        let mut rng = OsRng;
329        let input_limbs: Vec<_> = (0..8).map(|_| rng.gen::<u32>() as u64).collect();
330
331        test_gate_constraint(input_limbs);
332    }
333
334    #[test]
335    #[should_panic]
336    fn test_gate_constraint_bad() {
337        let mut rng = OsRng;
338        let input_limbs: Vec<_> = (0..8).map(|_| rng.gen()).collect();
339
340        test_gate_constraint(input_limbs);
341    }
342}