1use plonky2::{
 2    field::extension::Extendable, hash::hash_types::RichField, iop::witness::PartialWitness,
 3    plonk::circuit_builder::CircuitBuilder,
 4};
 5use pod2::backends::plonky2::{
 6    basetypes::{D, F},
 7    primitives::ec::{
 8        bits::{BigUInt320Target, CircuitBuilderBits},
 9        curve::{CircuitBuilderElliptic, Point, PointTarget, WitnessWriteCurve},
10    },
11};
12
13use crate::ECGFP5SecretKey;
14
15/// EcGFp5 ecret key target type
16pub struct ECGFP5SecretKeyTarget(pub BigUInt320Target);
17
18/// EcGFp5 public key (point) type alias
19pub type PublicKeyTarget = PointTarget;
20
21pub trait CircuitBuilderECGFP5PublicKey<F: RichField + Extendable<D>, const D: usize> {
22    /// Adds secret key target
23    fn add_secret_key(&mut self) -> ECGFP5SecretKeyTarget;
24
25    /// Derives public key from secret key
26    fn public_key(&mut self, sk: &ECGFP5SecretKeyTarget) -> PublicKeyTarget;
27}
28
29impl CircuitBuilderECGFP5PublicKey<F, D> for CircuitBuilder<F, D> {
30    fn add_secret_key(&mut self) -> ECGFP5SecretKeyTarget {
31        ECGFP5SecretKeyTarget(self.add_virtual_biguint320_target())
32    }
33
34    fn public_key(&mut self, sk: &ECGFP5SecretKeyTarget) -> PublicKeyTarget {
35        let generator = self.constant_point(Point::generator());
36        self.multiply_point(&sk.0.bits, &generator)
37    }
38}
39
40pub trait PartialWitnessECGFP5PublicKey {
41    /// Sets EcGFp5 secret key witnesses
42    fn set_secret_key_target(
43        &mut self,
44        target: &ECGFP5SecretKeyTarget,
45        value: &ECGFP5SecretKey,
46    ) -> anyhow::Result<()>;
47}
48
49impl PartialWitnessECGFP5PublicKey for PartialWitness<F> {
50    fn set_secret_key_target(
51        &mut self,
52        target: &ECGFP5SecretKeyTarget,
53        value: &ECGFP5SecretKey,
54    ) -> anyhow::Result<()> {
55        self.set_biguint320_target(&target.0, &value.0)
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use plonky2::{
62        iop::witness::PartialWitness,
63        plonk::{
64            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
65            config::PoseidonGoldilocksConfig,
66        },
67    };
68    use pod2::backends::plonky2::{
69        basetypes::{D, F},
70        primitives::ec::curve::WitnessWriteCurve,
71    };
72
73    use crate::{
74        ECGFP5SecretKey,
75        circuit::{CircuitBuilderECGFP5PublicKey, PartialWitnessECGFP5PublicKey},
76    };
77
78    #[test]
79    fn public_key_calculation() -> anyhow::Result<()> {
80        let config = CircuitConfig::standard_recursion_config();
81        let mut builder = CircuitBuilder::<F, D>::new(config);
82
83        let sk_target = builder.add_secret_key();
84        let pk_target = builder.public_key(&sk_target);
85
86        let data = builder.build::<PoseidonGoldilocksConfig>();
87        (0..10).try_for_each(|_| {
88            let sk = ECGFP5SecretKey::rand();
89            let pk = sk.public_key();
90
91            let mut pw = PartialWitness::new();
92            pw.set_secret_key_target(&sk_target, &sk)?;
93            pw.set_point_target(&pk_target, &pk)?;
94
95            let proof = data.prove(pw)?;
96            data.verify(proof)
97        })
98    }
99}