1use alloc::vec::Vec;
  2use core::marker::PhantomData;
  3
  4use itertools::Itertools;
  5use plonky2::field::extension::Extendable;
  6use plonky2::field::types::Field;
  7use plonky2::hash::hash_types::RichField;
  8use plonky2::iop::target::Target;
  9use plonky2::plonk::circuit_builder::CircuitBuilder;
 10use plonky2_u32::gadgets::arithmetic_u32::{CircuitBuilderU32, U32Target};
 11
 12use crate::gadgets::biguint::BigUintTarget;
 13use crate::gadgets::nonnative::NonNativeTarget;
 14
 15pub trait CircuitBuilderSplit<F: RichField + Extendable<D>, const D: usize> {
 16    fn split_u32_to_4_bit_limbs(&mut self, val: U32Target) -> Vec<Target>;
 17
 18    fn split_nonnative_to_4_bit_limbs<FF: Field>(
 19        &mut self,
 20        val: &NonNativeTarget<FF>,
 21    ) -> Vec<Target>;
 22
 23    fn split_nonnative_to_2_bit_limbs<FF: Field>(
 24        &mut self,
 25        val: &NonNativeTarget<FF>,
 26    ) -> Vec<Target>;
 27
 28    // Note: assumes its inputs are 4-bit limbs, and does not range-check.
 29    fn recombine_nonnative_4_bit_limbs<FF: Field>(
 30        &mut self,
 31        limbs: Vec<Target>,
 32    ) -> NonNativeTarget<FF>;
 33}
 34
 35impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilderSplit<F, D>
 36    for CircuitBuilder<F, D>
 37{
 38    fn split_u32_to_4_bit_limbs(&mut self, val: U32Target) -> Vec<Target> {
 39        let two_bit_limbs = self.split_le_base::<4>(val.0, 16);
 40        let four = self.constant(F::from_canonical_usize(4));
 41        let combined_limbs = two_bit_limbs
 42            .iter()
 43            .tuples()
 44            .map(|(&a, &b)| self.mul_add(b, four, a))
 45            .collect();
 46
 47        combined_limbs
 48    }
 49
 50    fn split_nonnative_to_4_bit_limbs<FF: Field>(
 51        &mut self,
 52        val: &NonNativeTarget<FF>,
 53    ) -> Vec<Target> {
 54        val.value
 55            .limbs
 56            .iter()
 57            .flat_map(|&l| self.split_u32_to_4_bit_limbs(l))
 58            .collect()
 59    }
 60
 61    fn split_nonnative_to_2_bit_limbs<FF: Field>(
 62        &mut self,
 63        val: &NonNativeTarget<FF>,
 64    ) -> Vec<Target> {
 65        val.value
 66            .limbs
 67            .iter()
 68            .flat_map(|&l| self.split_le_base::<4>(l.0, 16))
 69            .collect()
 70    }
 71
 72    // Note: assumes its inputs are 4-bit limbs, and does not range-check.
 73    fn recombine_nonnative_4_bit_limbs<FF: Field>(
 74        &mut self,
 75        limbs: Vec<Target>,
 76    ) -> NonNativeTarget<FF> {
 77        let base = self.constant_u32(1 << 4);
 78        let u32_limbs = limbs
 79            .chunks(8)
 80            .map(|chunk| {
 81                let mut combined_chunk = self.zero_u32();
 82                for i in (0..8).rev() {
 83                    let (low, _high) = self.mul_add_u32(combined_chunk, base, U32Target(chunk[i]));
 84                    combined_chunk = low;
 85                }
 86                combined_chunk
 87            })
 88            .collect();
 89
 90        NonNativeTarget {
 91            value: BigUintTarget { limbs: u32_limbs },
 92            _phantom: PhantomData,
 93        }
 94    }
 95}
 96
 97#[cfg(test)]
 98mod tests {
 99    use anyhow::Result;
100    use plonky2::field::secp256k1_scalar::Secp256K1Scalar;
101    use plonky2::field::types::Sample;
102    use plonky2::iop::witness::PartialWitness;
103    use plonky2::plonk::circuit_data::CircuitConfig;
104    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
105
106    use super::*;
107    use crate::gadgets::nonnative::{CircuitBuilderNonNative, NonNativeTarget};
108
109    #[test]
110    fn test_split_nonnative() -> Result<()> {
111        type FF = Secp256K1Scalar;
112        const D: usize = 2;
113        type C = PoseidonGoldilocksConfig;
114        type F = <C as GenericConfig<D>>::F;
115
116        let config = CircuitConfig::standard_ecc_config();
117        let pw = PartialWitness::new();
118        let mut builder = CircuitBuilder::<F, D>::new(config);
119
120        let x = FF::rand();
121        let x_target = builder.constant_nonnative(x);
122        let split = builder.split_nonnative_to_4_bit_limbs(&x_target);
123        let combined: NonNativeTarget<Secp256K1Scalar> =
124            builder.recombine_nonnative_4_bit_limbs(split);
125        builder.connect_nonnative(&x_target, &combined);
126
127        let data = builder.build::<C>();
128        let proof = data.prove(pw).unwrap();
129        data.verify(proof)
130    }
131}