1use alloc::string::{String, ToString};
  2use alloc::vec::Vec;
  3use alloc::{format, vec};
  4use anyhow::Result;
  5use num::integer::div_ceil;
  6use core::marker::PhantomData;
  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(&self, witness: &PartitionWitness<F>, out_buffer: &mut GeneratedValues<F>) -> Result<()> {
200        let num_input_limbs = self.gate.num_input_limbs;
201        for i in 0..num_input_limbs {
202            let sum_value = witness
203                .get_target(Target::wire(self.row, self.gate.wire_ith_input_limb(i)))
204                .to_canonical_u64() as u32;
205
206            let base = U32RangeCheckGate::<F, D>::BASE as u32;
207            let limbs = (0..self.gate.aux_limbs_per_input_limb())
208                .map(|j| Target::wire(self.row, self.gate.wire_ith_input_limb_jth_aux_limb(i, j)));
209            let limbs_value = (0..self.gate.aux_limbs_per_input_limb())
210                .scan(sum_value, |acc, _| {
211                    let tmp = *acc % base;
212                    *acc /= base;
213                    Some(F::from_canonical_u32(tmp))
214                })
215                .collect::<Vec<_>>();
216
217            for (b, b_value) in limbs.zip(limbs_value) {
218                out_buffer.set_target(b, b_value)?;
219            }
220        }
221
222        Ok(())
223    }
224
225    fn serialize(&self, dst: &mut Vec<u8>, common_data: &CommonCircuitData<F, D>) -> IoResult<()> {
226        self.gate.serialize(dst, common_data)?;
227        dst.write_usize(self.row)
228    }
229
230    fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData<F, D>) -> IoResult<Self> {
231        let gate = U32RangeCheckGate::deserialize(src, common_data)?;
232        let row = src.read_usize()?;
233        Ok(Self { row, gate })
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use anyhow::Result;
240    use itertools::unfold;
241    use plonky2::field::extension::quartic::QuarticExtension;
242    use plonky2::field::goldilocks_field::GoldilocksField;
243    use plonky2::field::types::{Field, Sample};
244    use plonky2::gates::gate_testing::{test_eval_fns, test_low_degree};
245    use plonky2::hash::hash_types::HashOut;
246    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
247    use rand::rngs::OsRng;
248    use rand::Rng;
249
250    use super::*;
251
252    #[test]
253    fn low_degree() {
254        test_low_degree::<GoldilocksField, _, 4>(U32RangeCheckGate::new(8))
255    }
256
257    #[test]
258    fn eval_fns() -> Result<()> {
259        const D: usize = 2;
260        type C = PoseidonGoldilocksConfig;
261        type F = <C as GenericConfig<D>>::F;
262        test_eval_fns::<F, C, _, D>(U32RangeCheckGate::new(8))
263    }
264
265    fn test_gate_constraint(input_limbs: Vec<u64>) {
266        type F = GoldilocksField;
267        type FF = QuarticExtension<GoldilocksField>;
268        const D: usize = 4;
269        const AUX_LIMB_BITS: usize = 2;
270        const BASE: usize = 1 << AUX_LIMB_BITS;
271        const AUX_LIMBS_PER_INPUT_LIMB: usize = (32 + AUX_LIMB_BITS - 1)/AUX_LIMB_BITS;
272
273        fn get_wires(input_limbs: Vec<u64>) -> Vec<FF> {
274            let num_input_limbs = input_limbs.len();
275            let mut v = Vec::new();
276
277            for i in 0..num_input_limbs {
278                let input_limb = input_limbs[i];
279
280                let split_to_limbs = |mut val, num| {
281                    unfold((), move |_| {
282                        let ret = val % (BASE as u64);
283                        val /= BASE as u64;
284                        Some(ret)
285                    })
286                    .take(num)
287                    .map(F::from_canonical_u64)
288                };
289
290                let mut aux_limbs: Vec<_> =
291                    split_to_limbs(input_limb, AUX_LIMBS_PER_INPUT_LIMB).collect();
292
293                v.append(&mut aux_limbs);
294            }
295
296            input_limbs
297                .iter()
298                .cloned()
299                .map(F::from_canonical_u64)
300                .chain(v.iter().cloned())
301                .map(|x| x.into())
302                .collect()
303        }
304
305        let gate = U32RangeCheckGate::<F, D> {
306            num_input_limbs: 8,
307            _phantom: PhantomData,
308        };
309
310        let vars = EvaluationVars {
311            local_constants: &[],
312            local_wires: &get_wires(input_limbs),
313            public_inputs_hash: &HashOut::rand(),
314        };
315
316        assert!(
317            gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()),
318            "Gate constraints are not satisfied."
319        );
320    }
321
322    #[test]
323    fn test_gate_constraint_good() {
324        let mut rng = OsRng;
325        let input_limbs: Vec<_> = (0..8).map(|_| rng.gen::<u32>() as u64).collect();
326
327        test_gate_constraint(input_limbs);
328    }
329
330    #[test]
331    #[should_panic]
332    fn test_gate_constraint_bad() {
333        let mut rng = OsRng;
334        let input_limbs: Vec<_> = (0..8).map(|_| rng.gen()).collect();
335
336        test_gate_constraint(input_limbs);
337    }
338}