Skip to main content

kimchi_stubs/
pasta_fp_plonk_index.rs

1use crate::{arkworks::CamlFp, gate_vector::fp::CamlPastaFpPlonkGateVectorPtr, srs::fp::CamlFpSrs};
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::{Fp, Pallas, Vesta, VestaParameters};
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<Vesta, FULL_ROUNDS> as poly_commitment::OpenProof<Vesta, FULL_ROUNDS>>::SRS;
29
30/// Holds a prover index behind one of two backing stores. The enum
31/// variants both [`Deref`] to the same `&ProverIndex`, so all existing
32/// `.0.cs…` / `.0.srs…` field-access patterns in this file continue to
33/// work via auto-deref. Concretely:
34///
35/// - [`IndexHandle::Owned`] is the classic path: a heap-allocated
36///   `ProverIndex` whose `Vec<F>` fields are Rust-owned. Populated by
37///   `caml_pasta_fp_plonk_index_create` / `_read`.
38/// - [`IndexHandle::Mmap`] wraps an [`MmapProverIndex`] whose bulk
39///   `Vec<F>` fields point into an mmap'd cache file. Populated by
40///   `caml_pasta_fp_plonk_index_read_cached`. Dropping an `Mmap` variant
41///   skips the inner `Vec<F>`s' destructors (they would otherwise call
42///   `dealloc` on mmap memory — UB) and unmaps the file.
43pub enum IndexHandle {
44    Owned(Box<ProverIndex<FULL_ROUNDS, Vesta, Srs>>),
45    Mmap(Box<MmapProverIndex<FULL_ROUNDS, Vesta, Srs>>),
46}
47
48impl core::ops::Deref for IndexHandle {
49    type Target = ProverIndex<FULL_ROUNDS, Vesta, Srs>;
50    fn deref(&self) -> &Self::Target {
51        match self {
52            IndexHandle::Owned(b) => b,
53            IndexHandle::Mmap(b) => b,
54        }
55    }
56}
57
58/// Boxed so that we don't store large proving indexes in the OCaml heap.
59#[derive(ocaml_gen::CustomType)]
60pub struct CamlPastaFpPlonkIndex(pub IndexHandle);
61pub type CamlPastaFpPlonkIndexPtr<'a> = ocaml::Pointer<'a, CamlPastaFpPlonkIndex>;
62
63extern "C" fn caml_pasta_fp_plonk_index_finalize(v: ocaml::Raw) {
64    unsafe {
65        let mut v: CamlPastaFpPlonkIndexPtr = v.as_pointer();
66        v.as_mut_ptr().drop_in_place();
67    }
68}
69
70impl ocaml::custom::Custom for CamlPastaFpPlonkIndex {
71    const NAME: &'static str = "CamlPastaFpPlonkIndex\0";
72    const USED: usize = 1;
73    /// Encourage the GC to free when there are > 12 in memory
74    const MAX: usize = 12;
75    const OPS: ocaml::custom::CustomOps = ocaml::custom::CustomOps {
76        identifier: Self::NAME.as_ptr() as *const ocaml::sys::Char,
77        finalize: Some(caml_pasta_fp_plonk_index_finalize),
78        ..ocaml::custom::DEFAULT_CUSTOM_OPS
79    };
80}
81
82#[ocaml_gen::func]
83#[ocaml::func]
84pub fn caml_pasta_fp_plonk_index_create(
85    gates: CamlPastaFpPlonkGateVectorPtr,
86    public: ocaml::Int,
87    lookup_tables: Vec<CamlLookupTable<CamlFp>>,
88    runtime_tables: Vec<CamlRuntimeTableCfg<CamlFp>>,
89    prev_challenges: ocaml::Int,
90    srs: CamlFpSrs,
91    lazy_mode: bool,
92) -> Result<CamlPastaFpPlonkIndex, ocaml::Error> {
93    let gates: Vec<_> = gates
94        .as_ref()
95        .0
96        .iter()
97        .map(|gate| CircuitGate::<Fp> {
98            typ: gate.typ,
99            wires: gate.wires,
100            coeffs: gate.coeffs.clone(),
101        })
102        .collect();
103
104    let runtime_tables: Vec<RuntimeTableCfg<Fp>> =
105        runtime_tables.into_iter().map(Into::into).collect();
106
107    let lookup_tables: Vec<LookupTable<Fp>> = lookup_tables.into_iter().map(Into::into).collect();
108
109    // create constraint system
110    let cs = ConstraintSystem::<Fp>::create(gates)
111        .public(public as usize)
112        .prev_challenges(prev_challenges as usize)
113        .max_poly_size(Some(srs.0.max_poly_size()))
114        .lookup(lookup_tables)
115        .runtime(if runtime_tables.is_empty() {
116            None
117        } else {
118            Some(runtime_tables)
119        })
120        .lazy_mode(lazy_mode)
121        .build()?;
122
123    // endo
124    let (endo_q, _endo_r) = poly_commitment::ipa::endos::<Pallas>();
125
126    srs.0.with_lagrange_basis(cs.domain.d1);
127
128    // create index
129    let mut index = ProverIndex::create(cs, endo_q, srs.clone(), lazy_mode);
130    // Compute and cache the verifier index digest
131    index.compute_verifier_index_digest::<DefaultFqSponge<
132        VestaParameters,
133        PlonkSpongeConstantsKimchi,
134        FULL_ROUNDS,
135    >>();
136
137    Ok(CamlPastaFpPlonkIndex(IndexHandle::Owned(Box::new(index))))
138}
139
140#[ocaml_gen::func]
141#[ocaml::func]
142pub fn caml_pasta_fp_plonk_index_max_degree(index: CamlPastaFpPlonkIndexPtr) -> ocaml::Int {
143    index.as_ref().0.srs.max_poly_size() as isize
144}
145
146#[ocaml_gen::func]
147#[ocaml::func]
148pub fn caml_pasta_fp_plonk_index_public_inputs(index: CamlPastaFpPlonkIndexPtr) -> ocaml::Int {
149    index.as_ref().0.cs.public as isize
150}
151
152#[ocaml_gen::func]
153#[ocaml::func]
154pub fn caml_pasta_fp_plonk_index_domain_d1_size(index: CamlPastaFpPlonkIndexPtr) -> ocaml::Int {
155    index.as_ref().0.cs.domain.d1.size() as isize
156}
157
158#[ocaml_gen::func]
159#[ocaml::func]
160pub fn caml_pasta_fp_plonk_index_domain_d4_size(index: CamlPastaFpPlonkIndexPtr) -> ocaml::Int {
161    index.as_ref().0.cs.domain.d4.size() as isize
162}
163
164#[ocaml_gen::func]
165#[ocaml::func]
166pub fn caml_pasta_fp_plonk_index_domain_d8_size(index: CamlPastaFpPlonkIndexPtr) -> ocaml::Int {
167    index.as_ref().0.cs.domain.d8.size() as isize
168}
169
170#[ocaml_gen::func]
171#[ocaml::func]
172pub fn caml_pasta_fp_plonk_index_read(
173    offset: Option<ocaml::Int>,
174    srs: CamlFpSrs,
175    path: String,
176) -> Result<CamlPastaFpPlonkIndex, ocaml::Error> {
177    // open the file for reading
178    let file = match File::open(path) {
179        Err(_) => {
180            return Err(
181                ocaml::Error::invalid_argument("caml_pasta_fp_plonk_index_read")
182                    .err()
183                    .unwrap(),
184            )
185        }
186        Ok(file) => file,
187    };
188    let mut r = BufReader::new(file);
189
190    // optional offset in file
191    if let Some(offset) = offset {
192        r.seek(Start(offset as u64))?;
193    }
194
195    // deserialize the index
196    let mut t =
197        ProverIndex::<FULL_ROUNDS, Vesta, Srs>::deserialize(&mut rmp_serde::Deserializer::new(r))?;
198    t.srs = srs.clone();
199
200    let (linearization, powers_of_alpha) = expr_linearization(Some(&t.cs.feature_flags), true);
201    t.linearization = linearization;
202    t.powers_of_alpha = powers_of_alpha;
203
204    Ok(CamlPastaFpPlonkIndex(IndexHandle::Owned(Box::new(t))))
205}
206
207#[ocaml_gen::func]
208#[ocaml::func]
209pub fn caml_pasta_fp_plonk_index_write(
210    append: Option<bool>,
211    index: CamlPastaFpPlonkIndexPtr<'static>,
212    path: String,
213) -> Result<(), ocaml::Error> {
214    let file = OpenOptions::new()
215        .append(append.unwrap_or(true))
216        .open(path)
217        .map_err(|_| {
218            ocaml::Error::invalid_argument("caml_pasta_fp_plonk_index_write")
219                .err()
220                .unwrap()
221        })?;
222    let w = BufWriter::new(file);
223    // Legacy rmp_serde write only supports the fully-owned `ProverIndex`
224    // backing; serialising an mmap-backed index would require copying
225    // every Vec<F> back through serde, which defeats its purpose.
226    match &index.as_ref().0 {
227        IndexHandle::Owned(b) => b
228            .serialize(&mut rmp_serde::Serializer::new(w))
229            .map_err(|e| e.into()),
230        IndexHandle::Mmap(_) => Err(ocaml::Error::Message(
231            "caml_pasta_fp_plonk_index_write: legacy serde write is not \
232             supported for mmap-backed indexes; use \
233             caml_pasta_fp_plonk_index_write_cached instead",
234        )),
235    }
236}
237
238/// Writes the proving index to `path` in the mmap-backed cache format.
239///
240/// `identifier` is a caller-supplied string (bounded at 512 bytes) that is
241/// round-tripped through the file header. Reads must pass the same
242/// identifier; mismatches return a descriptive error instead of loading
243/// the wrong key.
244///
245/// Writes are atomic: the file is staged at `path.tmp` then renamed into
246/// place so concurrent readers never observe a half-written file.
247#[ocaml_gen::func]
248#[ocaml::func]
249pub fn caml_pasta_fp_plonk_index_write_cached(
250    identifier: String,
251    index: CamlPastaFpPlonkIndexPtr<'static>,
252    path: String,
253) -> Result<(), ocaml::Error> {
254    // Both `IndexHandle` variants Deref to `&ProverIndex`, which is what
255    // `write_cache` consumes; a double-borrow (`&*`) forces the coercion.
256    kimchi::cached_prover_index::write_cache(
257        &identifier,
258        &*index.as_ref().0,
259        std::path::Path::new(&path),
260    )
261    .map_err(|e| {
262        crate::cache_error::CacheFfiError::wrap(format!(
263            "caml_pasta_fp_plonk_index_write_cached: {e}"
264        ))
265    })
266}
267
268/// Reads a proving index from `path` in the mmap-backed cache format,
269/// binding the supplied `srs` onto the reconstructed index.
270///
271/// `identifier` must match the value used at write time or the call fails
272/// without loading the key.
273#[ocaml_gen::func]
274#[ocaml::func]
275pub fn caml_pasta_fp_plonk_index_read_cached(
276    identifier: String,
277    srs: CamlFpSrs,
278    path: String,
279) -> Result<CamlPastaFpPlonkIndex, ocaml::Error> {
280    let index = kimchi::cached_prover_index::read_cache::<FULL_ROUNDS, Vesta, Srs>(
281        &identifier,
282        std::path::Path::new(&path),
283        srs.clone(),
284    )
285    .map_err(|e| {
286        crate::cache_error::CacheFfiError::wrap(format!(
287            "caml_pasta_fp_plonk_index_read_cached: {e}"
288        ))
289    })?;
290    Ok(CamlPastaFpPlonkIndex(IndexHandle::Mmap(Box::new(index))))
291}
292
293#[cfg(test)]
294mod tests {
295    /// ocaml-rs raises `ocaml::Error::Error` by calling `caml_failwith` on
296    /// `format!("{:?}", e)` — the boxed error's Debug rendering. Whatever the
297    /// read_cached/write_cached wrappers box must therefore Debug-render as
298    /// the plain message, or OCaml sees the text wrapped in literal quotes
299    /// with inner escapes (`Failure "\"caml_...: ...\""`), garbling logs and
300    /// breaking any caller that matches on the failure string.
301    #[test]
302    fn cached_ffi_failure_text_is_the_plain_message() {
303        // Same construction as the map_err sites on the cached-index FFI
304        // paths.
305        let err = crate::cache_error::CacheFfiError::wrap(format!(
306            "caml_pasta_fp_plonk_index_write_cached: {}",
307            "boom"
308        ));
309        let ocaml::Error::Error(boxed) = err else {
310            unreachable!()
311        };
312        assert_eq!(
313            format!("{boxed:?}"),
314            "caml_pasta_fp_plonk_index_write_cached: boom"
315        );
316    }
317}