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