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