Skip to main content

kimchi_stubs/
pasta_fq_plonk_index.rs

1use crate::{arkworks::CamlFq, gate_vector::fq::CamlPastaFqPlonkGateVectorPtr, srs::fq::CamlFqSrs};
2use ark_poly::EvaluationDomain;
3use kimchi::{
4    cached_prover_index::MmapProverIndex,
5    circuits::{
6        constraints::ConstraintSystem,
7        gate::CircuitGate,
8        lookup::{
9            runtime_tables::{caml::CamlRuntimeTableCfg, RuntimeTableCfg},
10            tables::{caml::CamlLookupTable, LookupTable},
11        },
12    },
13    linearization::expr_linearization,
14    prover_index::ProverIndex,
15};
16use mina_curves::pasta::{Fq, Pallas, PallasParameters, Vesta};
17use mina_poseidon::{
18    constants::PlonkSpongeConstantsKimchi, pasta::FULL_ROUNDS, sponge::DefaultFqSponge,
19};
20use poly_commitment::{ipa::OpeningProof, lagrange_basis::WithLagrangeBasis, SRS as _};
21use serde::{Deserialize, Serialize};
22use std::{
23    fs::{File, OpenOptions},
24    io::{BufReader, BufWriter, Seek, SeekFrom::Start},
25};
26
27type Srs =
28    <OpeningProof<Pallas, FULL_ROUNDS> as poly_commitment::OpenProof<Pallas, FULL_ROUNDS>>::SRS;
29
30/// Fq companion of the Fp-side [`crate::pasta_fp_plonk_index::IndexHandle`] —
31/// selects between an owned `ProverIndex` and an `MmapProverIndex` whose
32/// bulk `Vec<F>` fields point into a mmap'd cache file. Both variants
33/// [`Deref`] to the same `&ProverIndex` so existing field-access patterns
34/// in this file continue to work via auto-deref.
35pub enum IndexHandle {
36    Owned(Box<ProverIndex<FULL_ROUNDS, Pallas, Srs>>),
37    Mmap(Box<MmapProverIndex<FULL_ROUNDS, Pallas, Srs>>),
38}
39
40impl core::ops::Deref for IndexHandle {
41    type Target = ProverIndex<FULL_ROUNDS, Pallas, Srs>;
42    fn deref(&self) -> &Self::Target {
43        match self {
44            IndexHandle::Owned(b) => b,
45            IndexHandle::Mmap(b) => b,
46        }
47    }
48}
49
50/// Boxed so that we don't store large proving indexes in the OCaml heap.
51#[derive(ocaml_gen::CustomType)]
52pub struct CamlPastaFqPlonkIndex(pub IndexHandle);
53pub type CamlPastaFqPlonkIndexPtr<'a> = ocaml::Pointer<'a, CamlPastaFqPlonkIndex>;
54
55extern "C" fn caml_pasta_fq_plonk_index_finalize(v: ocaml::Raw) {
56    unsafe {
57        let mut v: CamlPastaFqPlonkIndexPtr = v.as_pointer();
58        v.as_mut_ptr().drop_in_place();
59    }
60}
61
62impl ocaml::custom::Custom for CamlPastaFqPlonkIndex {
63    const NAME: &'static str = "CamlPastaFqPlonkIndex\0";
64    const USED: usize = 1;
65    /// Encourage the GC to free when there are > 12 in memory
66    const MAX: usize = 12;
67    const OPS: ocaml::custom::CustomOps = ocaml::custom::CustomOps {
68        identifier: Self::NAME.as_ptr() as *const ocaml::sys::Char,
69        finalize: Some(caml_pasta_fq_plonk_index_finalize),
70        ..ocaml::custom::DEFAULT_CUSTOM_OPS
71    };
72}
73
74#[ocaml_gen::func]
75#[ocaml::func]
76pub fn caml_pasta_fq_plonk_index_create(
77    gates: CamlPastaFqPlonkGateVectorPtr,
78    public: ocaml::Int,
79    lookup_tables: Vec<CamlLookupTable<CamlFq>>,
80    runtime_tables: Vec<CamlRuntimeTableCfg<CamlFq>>,
81    prev_challenges: ocaml::Int,
82    srs: CamlFqSrs,
83    lazy_mode: bool,
84) -> Result<CamlPastaFqPlonkIndex, ocaml::Error> {
85    let gates: Vec<_> = gates
86        .as_ref()
87        .0
88        .iter()
89        .map(|gate| CircuitGate::<Fq> {
90            typ: gate.typ,
91            wires: gate.wires,
92            coeffs: gate.coeffs.clone(),
93        })
94        .collect();
95
96    let runtime_tables: Vec<RuntimeTableCfg<Fq>> =
97        runtime_tables.into_iter().map(Into::into).collect();
98
99    let lookup_tables: Vec<LookupTable<Fq>> = lookup_tables.into_iter().map(Into::into).collect();
100
101    // create constraint system
102    let cs = ConstraintSystem::<Fq>::create(gates)
103        .public(public as usize)
104        .prev_challenges(prev_challenges as usize)
105        .lookup(lookup_tables)
106        .runtime(if runtime_tables.is_empty() {
107            None
108        } else {
109            Some(runtime_tables)
110        })
111        .lazy_mode(lazy_mode)
112        .build()?;
113
114    // endo
115    let (endo_q, _endo_r) = poly_commitment::ipa::endos::<Vesta>();
116
117    srs.0.with_lagrange_basis(cs.domain.d1);
118
119    // create index
120    let mut index =
121        ProverIndex::<FULL_ROUNDS, Pallas, Srs>::create(cs, endo_q, srs.clone(), lazy_mode);
122    // Compute and cache the verifier index digest
123    index.compute_verifier_index_digest::<DefaultFqSponge<
124        PallasParameters,
125        PlonkSpongeConstantsKimchi,
126        FULL_ROUNDS,
127    >>();
128
129    Ok(CamlPastaFqPlonkIndex(IndexHandle::Owned(Box::new(index))))
130}
131
132#[ocaml_gen::func]
133#[ocaml::func]
134pub fn caml_pasta_fq_plonk_index_max_degree(index: CamlPastaFqPlonkIndexPtr) -> ocaml::Int {
135    index.as_ref().0.srs.max_poly_size() as isize
136}
137
138#[ocaml_gen::func]
139#[ocaml::func]
140pub fn caml_pasta_fq_plonk_index_public_inputs(index: CamlPastaFqPlonkIndexPtr) -> ocaml::Int {
141    index.as_ref().0.cs.public as isize
142}
143
144#[ocaml_gen::func]
145#[ocaml::func]
146pub fn caml_pasta_fq_plonk_index_domain_d1_size(index: CamlPastaFqPlonkIndexPtr) -> ocaml::Int {
147    index.as_ref().0.cs.domain.d1.size() as isize
148}
149
150#[ocaml_gen::func]
151#[ocaml::func]
152pub fn caml_pasta_fq_plonk_index_domain_d4_size(index: CamlPastaFqPlonkIndexPtr) -> ocaml::Int {
153    index.as_ref().0.cs.domain.d4.size() as isize
154}
155
156#[ocaml_gen::func]
157#[ocaml::func]
158pub fn caml_pasta_fq_plonk_index_domain_d8_size(index: CamlPastaFqPlonkIndexPtr) -> ocaml::Int {
159    index.as_ref().0.cs.domain.d8.size() as isize
160}
161
162#[ocaml_gen::func]
163#[ocaml::func]
164pub fn caml_pasta_fq_plonk_index_read(
165    offset: Option<ocaml::Int>,
166    srs: CamlFqSrs,
167    path: String,
168) -> Result<CamlPastaFqPlonkIndex, ocaml::Error> {
169    // open the file for reading
170    let file = match File::open(path) {
171        Err(_) => {
172            return Err(
173                ocaml::Error::invalid_argument("caml_pasta_fp_plonk_index_read")
174                    .err()
175                    .unwrap(),
176            )
177        }
178        Ok(file) => file,
179    };
180    let mut r = BufReader::new(file);
181
182    // optional offset in file
183    if let Some(offset) = offset {
184        r.seek(Start(offset as u64))?;
185    }
186
187    // deserialize the index
188    let mut t =
189        ProverIndex::<FULL_ROUNDS, Pallas, Srs>::deserialize(&mut rmp_serde::Deserializer::new(r))?;
190    t.srs = srs.clone();
191
192    let (linearization, powers_of_alpha) = expr_linearization(Some(&t.cs.feature_flags), true);
193    t.linearization = linearization;
194    t.powers_of_alpha = powers_of_alpha;
195
196    Ok(CamlPastaFqPlonkIndex(IndexHandle::Owned(Box::new(t))))
197}
198
199#[ocaml_gen::func]
200#[ocaml::func]
201pub fn caml_pasta_fq_plonk_index_write(
202    append: Option<bool>,
203    index: CamlPastaFqPlonkIndexPtr<'static>,
204    path: String,
205) -> Result<(), ocaml::Error> {
206    let file = OpenOptions::new()
207        .append(append.unwrap_or(true))
208        .open(path)
209        .map_err(|_| {
210            ocaml::Error::invalid_argument("caml_pasta_fq_plonk_index_write")
211                .err()
212                .unwrap()
213        })?;
214    let w = BufWriter::new(file);
215    match &index.as_ref().0 {
216        IndexHandle::Owned(b) => b
217            .serialize(&mut rmp_serde::Serializer::new(w))
218            .map_err(|e| e.into()),
219        IndexHandle::Mmap(_) => Err(ocaml::Error::Message(
220            "caml_pasta_fq_plonk_index_write: legacy serde write is not \
221             supported for mmap-backed indexes; use \
222             caml_pasta_fq_plonk_index_write_cached instead",
223        )),
224    }
225}
226
227/// Writes the proving index to `path` in the mmap-backed cache format.
228/// See the Fp variant for full documentation; this is the Pallas/Fq
229/// companion used by the Wrap side of pickles.
230#[ocaml_gen::func]
231#[ocaml::func]
232pub fn caml_pasta_fq_plonk_index_write_cached(
233    identifier: String,
234    index: CamlPastaFqPlonkIndexPtr<'static>,
235    path: String,
236) -> Result<(), ocaml::Error> {
237    kimchi::cached_prover_index::write_cache(
238        &identifier,
239        &*index.as_ref().0,
240        std::path::Path::new(&path),
241    )
242    .map_err(|e| {
243        crate::cache_error::CacheFfiError::wrap(format!(
244            "caml_pasta_fq_plonk_index_write_cached: {e}"
245        ))
246    })
247}
248
249/// Reads a proving index from `path` in the mmap-backed cache format.
250/// See the Fp variant for full documentation.
251#[ocaml_gen::func]
252#[ocaml::func]
253pub fn caml_pasta_fq_plonk_index_read_cached(
254    identifier: String,
255    srs: CamlFqSrs,
256    path: String,
257) -> Result<CamlPastaFqPlonkIndex, ocaml::Error> {
258    let index = kimchi::cached_prover_index::read_cache::<FULL_ROUNDS, Pallas, Srs>(
259        &identifier,
260        std::path::Path::new(&path),
261        srs.clone(),
262    )
263    .map_err(|e| {
264        crate::cache_error::CacheFfiError::wrap(format!(
265            "caml_pasta_fq_plonk_index_read_cached: {e}"
266        ))
267    })?;
268    Ok(CamlPastaFqPlonkIndex(IndexHandle::Mmap(Box::new(index))))
269}