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