1use ark_bn254::{Bn254, Fr, G1Projective as G1};
  2use ark_crypto_primitives::sponge::poseidon::PoseidonConfig;
  3use ark_grumpkin::Projective as G2;
  4use experimental_frontends::{circom::CircomFCircuit, utils::VecF};
  5use sonobe::{
  6    commitment::{kzg::KZG, pedersen::Pedersen},
  7    folding::{hypernova::HyperNova, nova::Nova},
  8    transcript::poseidon::poseidon_canonical_config,
  9    Error, FoldingScheme, MultiFolding,
 10};
 11use tracing::info_span;
 12
 13use crate::circuit::STEP_INPUT_WIDTH;
 14
 15pub type NovaFolding =
 16    Nova<G1, G2, CircomFCircuit<Fr, STEP_INPUT_WIDTH>, KZG<'static, Bn254>, Pedersen<G2>, false>;
 17pub type HyperNovaFolding<const M: usize, const N: usize> = HyperNova<
 18    G1,
 19    G2,
 20    CircomFCircuit<Fr, STEP_INPUT_WIDTH>,
 21    KZG<'static, Bn254>,
 22    Pedersen<G2>,
 23    M,
 24    N,
 25    false,
 26>;
 27
 28pub struct StepInput<OtherInstances> {
 29    pub external_inputs: VecF<Fr, STEP_INPUT_WIDTH>,
 30    pub other_instances: Option<OtherInstances>,
 31}
 32
 33pub trait FoldingSchemeExt: FoldingScheme<G1, G2, CircomFCircuit<Fr, STEP_INPUT_WIDTH>> {
 34    const MULTISTEP_SIZE: usize;
 35
 36    fn num_steps(num_inputs: usize) -> usize {
 37        assert_eq!(num_inputs % Self::MULTISTEP_SIZE, 0);
 38        num_inputs / Self::MULTISTEP_SIZE
 39    }
 40
 41    fn prepreprocess(
 42        poseidon_config: PoseidonConfig<Fr>,
 43        circuit: CircomFCircuit<Fr, STEP_INPUT_WIDTH>,
 44    ) -> Self::PreprocessorParam;
 45
 46    fn transform_multi_input(
 47        &self,
 48        multi_input: Vec<Vec<Fr>>,
 49        initial_state: Vec<Fr>,
 50        rng: &mut impl rand::RngCore,
 51    ) -> StepInput<Self::MultiCommittedInstanceWithWitness>;
 52
 53    fn prove_multistep(
 54        &mut self,
 55        multi_input: Vec<Vec<Fr>>,
 56        initial_state: Vec<Fr>,
 57        rng: &mut impl rand::RngCore,
 58    ) -> Result<(), Error> {
 59        let step_input = info_span!("Input prep")
 60            .in_scope(|| self.transform_multi_input(multi_input, initial_state, rng));
 61
 62        info_span!("Proving").in_scope(|| {
 63            self.prove_step(rng, step_input.external_inputs, step_input.other_instances)
 64        })
 65    }
 66}
 67
 68impl FoldingSchemeExt for NovaFolding {
 69    const MULTISTEP_SIZE: usize = 1;
 70
 71    fn prepreprocess(
 72        poseidon_config: PoseidonConfig<Fr>,
 73        circuit: CircomFCircuit<Fr, STEP_INPUT_WIDTH>,
 74    ) -> Self::PreprocessorParam {
 75        Self::PreprocessorParam::new(poseidon_config, circuit)
 76    }
 77
 78    fn transform_multi_input(
 79        &self,
 80        input: Vec<Vec<Fr>>,
 81        _initial_state: Vec<Fr>,
 82        _rng: &mut impl rand::RngCore,
 83    ) -> StepInput<Self::MultiCommittedInstanceWithWitness> {
 84        assert_eq!(input.len(), 1);
 85        StepInput {
 86            external_inputs: VecF(input[0].clone()),
 87            other_instances: None,
 88        }
 89    }
 90}
 91
 92impl<const M: usize, const N: usize> FoldingSchemeExt for HyperNovaFolding<M, N> {
 93    const MULTISTEP_SIZE: usize = M + N - 1;
 94
 95    fn prepreprocess(
 96        poseidon_config: PoseidonConfig<Fr>,
 97        circuit: CircomFCircuit<Fr, STEP_INPUT_WIDTH>,
 98    ) -> Self::PreprocessorParam {
 99        Self::PreprocessorParam::new(poseidon_config, circuit)
100    }
101
102    fn transform_multi_input(
103        &self,
104        multi_input: Vec<Vec<Fr>>,
105        initial_state: Vec<Fr>,
106        rng: &mut impl rand::RngCore,
107    ) -> StepInput<Self::MultiCommittedInstanceWithWitness> {
108        let (running, rest) = multi_input.split_at(M - 1);
109        let (incoming, [single]) = rest.split_at(N - 1) else {
110            panic!("Invalid input chunk size");
111        };
112
113        let new_running = |instance| {
114            self.new_running_instance(&mut *rng, initial_state.clone(), VecF(instance))
115                .expect("Failed to create running instance")
116        };
117
118        let new_instances =
119            |instances: Vec<Vec<Fr>>, maker| instances.into_iter().map(maker).collect();
120
121        let lcccs = new_instances(running.to_vec(), new_running);
122
123        let cccs = incoming
124            .iter()
125            .map(|instance| {
126                self.new_incoming_instance(&mut *rng, initial_state.clone(), VecF(instance.clone()))
127                    .expect("Failed to create incoming instance")
128            })
129            .collect();
130
131        StepInput {
132            external_inputs: VecF(single.clone()),
133            other_instances: Some((lcccs, cccs)),
134        }
135    }
136}
137
138pub fn prepare_folding<FS: FoldingSchemeExt>(
139    circuit: &CircomFCircuit<Fr, STEP_INPUT_WIDTH>,
140    start_ivc_state: Vec<Fr>,
141    rng: &mut impl rand::RngCore,
142) -> (FS, FS::VerifierParam) {
143    let preprocess_params = FS::prepreprocess(poseidon_canonical_config::<Fr>(), circuit.clone());
144    let params =
145        FS::preprocess(&mut *rng, &preprocess_params).expect("Failed to preprocess folding scheme");
146    let folding =
147        FS::init(&params, circuit.clone(), start_ivc_state).expect("Failed to init folding scheme");
148
149    (folding, params.1)
150}
151
152pub fn verify_folding<FS: FoldingSchemeExt>(folding: &FS, folding_vp: FS::VerifierParam) {
153    let ivc_proof = folding.ivc_proof();
154    FS::verify(folding_vp, ivc_proof).expect("Failed to verify folded proof");
155}