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