1use alloc::vec;
  2
  3use num::BigUint;
  4use plonky2::field::extension::Extendable;
  5use plonky2::field::types::Field;
  6use plonky2::hash::hash_types::RichField;
  7use plonky2::hash::keccak::KeccakHash;
  8use plonky2::plonk::circuit_builder::CircuitBuilder;
  9use plonky2::plonk::config::{GenericHashOut, Hasher};
 10
 11use crate::curve::curve_types::{Curve, CurveScalar};
 12use crate::gadgets::curve::{AffinePointTarget, CircuitBuilderCurve};
 13use crate::gadgets::curve_windowed_mul::CircuitBuilderWindowedMul;
 14use crate::gadgets::nonnative::NonNativeTarget;
 15use crate::gadgets::split_nonnative::CircuitBuilderSplit;
 16
 17/// Computes `n*p + m*q` using windowed MSM, with a 2-bit window.
 18/// See Algorithm 9.23 in Handbook of Elliptic and Hyperelliptic Curve Cryptography for a
 19/// description.
 20/// Note: Doesn't work if `p == q`.
 21pub fn curve_msm_circuit<C: Curve, F: RichField + Extendable<D>, const D: usize>(
 22    builder: &mut CircuitBuilder<F, D>,
 23    p: &AffinePointTarget<C>,
 24    q: &AffinePointTarget<C>,
 25    n: &NonNativeTarget<C::ScalarField>,
 26    m: &NonNativeTarget<C::ScalarField>,
 27) -> AffinePointTarget<C> {
 28    let limbs_n = builder.split_nonnative_to_2_bit_limbs(n);
 29    let limbs_m = builder.split_nonnative_to_2_bit_limbs(m);
 30    assert_eq!(limbs_n.len(), limbs_m.len());
 31    let num_limbs = limbs_n.len();
 32
 33    let hash_0 = KeccakHash::<32>::hash_no_pad(&[F::ZERO]);
 34    let hash_0_scalar = C::ScalarField::from_noncanonical_biguint(BigUint::from_bytes_le(
 35        &GenericHashOut::<F>::to_bytes(&hash_0),
 36    ));
 37    let rando = (CurveScalar(hash_0_scalar) * C::GENERATOR_PROJECTIVE).to_affine();
 38    let rando_t = builder.constant_affine_point(rando);
 39    let neg_rando = builder.constant_affine_point(-rando);
 40
 41    // Precomputes `precomputation[i + 4*j] = i*p + j*q` for `i,j=0..4`.
 42    let mut precomputation = vec![p.clone(); 16];
 43    let mut cur_p = rando_t.clone();
 44    let mut cur_q = rando_t.clone();
 45    for i in 0..4 {
 46        precomputation[i] = cur_p.clone();
 47        precomputation[4 * i] = cur_q.clone();
 48        cur_p = builder.curve_add(&cur_p, p);
 49        cur_q = builder.curve_add(&cur_q, q);
 50    }
 51    for i in 1..4 {
 52        precomputation[i] = builder.curve_add(&precomputation[i], &neg_rando);
 53        precomputation[4 * i] = builder.curve_add(&precomputation[4 * i], &neg_rando);
 54    }
 55    for i in 1..4 {
 56        for j in 1..4 {
 57            precomputation[i + 4 * j] =
 58                builder.curve_add(&precomputation[i], &precomputation[4 * j]);
 59        }
 60    }
 61
 62    let four = builder.constant(F::from_canonical_usize(4));
 63
 64    let zero = builder.zero();
 65    let mut result = rando_t;
 66    for (limb_n, limb_m) in limbs_n.into_iter().zip(limbs_m).rev() {
 67        result = builder.curve_repeated_double(&result, 2);
 68        let index = builder.mul_add(four, limb_m, limb_n);
 69        let r = builder.random_access_curve_points(index, precomputation.clone());
 70        let is_zero = builder.is_equal(index, zero);
 71        let should_add = builder.not(is_zero);
 72        result = builder.curve_conditional_add(&result, &r, should_add);
 73    }
 74    let starting_point_multiplied = (0..2 * num_limbs).fold(rando, |acc, _| acc.double());
 75    let to_add = builder.constant_affine_point(-starting_point_multiplied);
 76    result = builder.curve_add(&result, &to_add);
 77
 78    result
 79}
 80
 81#[cfg(test)]
 82mod tests {
 83    use anyhow::Result;
 84    use plonky2::field::secp256k1_scalar::Secp256K1Scalar;
 85    use plonky2::field::types::Sample;
 86    use plonky2::iop::witness::PartialWitness;
 87    use plonky2::plonk::circuit_builder::CircuitBuilder;
 88    use plonky2::plonk::circuit_data::CircuitConfig;
 89    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
 90
 91    use crate::curve::curve_types::{Curve, CurveScalar};
 92    use crate::curve::secp256k1::Secp256K1;
 93    use crate::gadgets::curve::CircuitBuilderCurve;
 94    use crate::gadgets::curve_msm::curve_msm_circuit;
 95    use crate::gadgets::nonnative::CircuitBuilderNonNative;
 96
 97    #[test]
 98    #[ignore]
 99    fn test_curve_msm() -> Result<()> {
100        const D: usize = 2;
101        type C = PoseidonGoldilocksConfig;
102        type F = <C as GenericConfig<D>>::F;
103
104        let config = CircuitConfig::standard_ecc_config();
105
106        let pw = PartialWitness::new();
107        let mut builder = CircuitBuilder::<F, D>::new(config);
108
109        let p =
110            (CurveScalar(Secp256K1Scalar::rand()) * Secp256K1::GENERATOR_PROJECTIVE).to_affine();
111        let q =
112            (CurveScalar(Secp256K1Scalar::rand()) * Secp256K1::GENERATOR_PROJECTIVE).to_affine();
113        let n = Secp256K1Scalar::rand();
114        let m = Secp256K1Scalar::rand();
115
116        let res =
117            (CurveScalar(n) * p.to_projective() + CurveScalar(m) * q.to_projective()).to_affine();
118        let res_expected = builder.constant_affine_point(res);
119        builder.curve_assert_valid(&res_expected);
120
121        let p_target = builder.constant_affine_point(p);
122        let q_target = builder.constant_affine_point(q);
123        let n_target = builder.constant_nonnative(n);
124        let m_target = builder.constant_nonnative(m);
125
126        let res_target =
127            curve_msm_circuit(&mut builder, &p_target, &q_target, &n_target, &m_target);
128        builder.curve_assert_valid(&res_target);
129
130        builder.connect_affine_point(&res_target, &res_expected);
131
132        dbg!(builder.num_gates());
133        let data = builder.build::<C>();
134        let proof = data.prove(pw).unwrap();
135
136        data.verify(proof)
137    }
138}