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