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