kimchi_stubs/
lagrange_basis.rs

1use ark_ec::AffineRepr;
2use ark_poly::{EvaluationDomain, Radix2EvaluationDomain as D};
3use cache::LagrangeCache;
4use mina_curves::pasta::{Pallas, Vesta};
5use poly_commitment::{commitment::CommitmentCurve, ipa::SRS, SRS as _};
6use std::env;
7
8pub trait WithLagrangeBasis<G: AffineRepr> {
9    fn with_lagrange_basis(&self, domain: D<G::ScalarField>);
10}
11
12impl WithLagrangeBasis<Vesta> for SRS<Vesta> {
13    fn with_lagrange_basis(&self, domain: D<<Vesta as AffineRepr>::ScalarField>) {
14        match env::var("LAGRANGE_CACHE_DIR") {
15            Ok(_) => add_lagrange_basis_with_cache(self, domain, cache::get_vesta_file_cache()),
16            Err(_) => {
17                self.get_lagrange_basis(domain);
18            }
19        }
20    }
21}
22
23impl WithLagrangeBasis<Pallas> for SRS<Pallas> {
24    fn with_lagrange_basis(&self, domain: D<<Pallas as AffineRepr>::ScalarField>) {
25        match env::var("LAGRANGE_CACHE_DIR") {
26            Ok(_) => add_lagrange_basis_with_cache(self, domain, cache::get_pallas_file_cache()),
27            Err(_) => {
28                self.get_lagrange_basis(domain);
29            }
30        }
31    }
32}
33
34fn add_lagrange_basis_with_cache<G: CommitmentCurve, C: LagrangeCache<G>>(
35    srs: &SRS<G>,
36    domain: D<G::ScalarField>,
37    cache: &C,
38) {
39    let n = domain.size();
40    if srs.lagrange_bases.contains_key(&n) {
41        return;
42    }
43    if let Some(basis) = cache.load_lagrange_basis_from_cache(srs.g.len(), &domain) {
44        srs.lagrange_bases.get_or_generate(n, || basis);
45    } else {
46        let basis = srs.get_lagrange_basis(domain);
47        cache.cache_lagrange_basis(srs.g.len(), &domain, basis);
48    }
49}
50
51mod cache {
52    use ark_ec::AffineRepr;
53    use ark_poly::{EvaluationDomain, Radix2EvaluationDomain as D};
54    use core::marker::PhantomData;
55    use mina_curves::pasta::{Pallas, Vesta};
56    use once_cell::sync::Lazy;
57    use poly_commitment::PolyComm;
58    use std::{
59        env, fs,
60        fs::File,
61        path::{Path, PathBuf},
62    };
63
64    pub trait LagrangeCache<G: AffineRepr> {
65        type CacheKey;
66
67        fn lagrange_basis_cache_key(
68            &self,
69            srs_length: usize,
70            domain: &D<G::ScalarField>,
71        ) -> Self::CacheKey;
72
73        fn load_lagrange_basis_from_cache(
74            &self,
75            srs_length: usize,
76            domain: &D<G::ScalarField>,
77        ) -> Option<Vec<PolyComm<G>>>;
78
79        fn cache_lagrange_basis(
80            &self,
81            srs_length: usize,
82            domain: &D<G::ScalarField>,
83            basis: &[PolyComm<G>],
84        );
85    }
86
87    #[derive(Debug, Clone, PartialEq, Eq)]
88    pub struct FileCache<G> {
89        cache_dir: PathBuf,
90        point_type: PhantomData<G>,
91    }
92
93    impl<G> FileCache<G> {
94        fn new(cache_dir: PathBuf) -> Self {
95            FileCache {
96                cache_dir,
97                point_type: PhantomData,
98            }
99        }
100    }
101
102    // The FileCache implementation uses a directory as a cache for the Lagrange
103    // basis hash map -- i.e every file corresponds to a Lagrange basis for a
104    // given G-basis and domain size.
105    impl<G: AffineRepr> LagrangeCache<G> for FileCache<G> {
106        type CacheKey = PathBuf;
107
108        fn lagrange_basis_cache_key(
109            &self,
110            srs_length: usize,
111            domain: &D<G::ScalarField>,
112        ) -> Self::CacheKey {
113            self.cache_dir.clone().join(format!(
114                "lagrange_basis_{:}-{:}",
115                srs_length,
116                domain.size()
117            ))
118        }
119
120        fn load_lagrange_basis_from_cache(
121            &self,
122            srs_length: usize,
123            domain: &D<G::ScalarField>,
124        ) -> Option<Vec<PolyComm<G>>> {
125            let cache_key = self.lagrange_basis_cache_key(srs_length, domain);
126            if Path::exists(&cache_key) {
127                let f = File::open(cache_key.clone()).unwrap_or_else(|_| {
128                    panic!("Missing lagrange basis cache file {:?}", cache_key)
129                });
130                let basis: Vec<PolyComm<G>> =
131                    rmp_serde::decode::from_read(f).unwrap_or_else(|_| {
132                        panic!("Error decoding lagrange cache file {:?}", cache_key)
133                    });
134                Some(basis)
135            } else {
136                None
137            }
138        }
139
140        fn cache_lagrange_basis(
141            &self,
142            srs_length: usize,
143            domain: &D<G::ScalarField>,
144            basis: &[PolyComm<G>],
145        ) {
146            let cache_key = self.lagrange_basis_cache_key(srs_length, domain);
147            if !Path::exists(&cache_key) {
148                let mut f = File::create(cache_key.clone()).unwrap_or_else(|_| {
149                    panic!("Error creating lagrange basis cache file {:?}", cache_key)
150                });
151                rmp_serde::encode::write(&mut f, basis).unwrap_or_else(|_| {
152                    panic!("Error encoding lagrange basis to file {:?}", cache_key)
153                });
154            }
155        }
156    }
157
158    // The following two caches are all that we need for mina tests. These will
159    // not be initialized unless they are explicitly called.
160    static VESTA_FILE_CACHE: Lazy<FileCache<Vesta>> = Lazy::new(|| {
161        let cache_base_dir: String =
162            env::var("LAGRANGE_CACHE_DIR").expect("LAGRANGE_CACHE_DIR missing in env");
163        let cache_dir = PathBuf::from(format!("{}/vesta", cache_base_dir));
164        if !cache_dir.exists() {
165            fs::create_dir_all(&cache_dir).unwrap();
166        }
167        FileCache::new(cache_dir)
168    });
169
170    pub fn get_vesta_file_cache() -> &'static FileCache<Vesta> {
171        &VESTA_FILE_CACHE
172    }
173
174    static PALLAS_FILE_CACHE: Lazy<FileCache<Pallas>> = Lazy::new(|| {
175        let cache_base_dir: String =
176            env::var("LAGRANGE_CACHE_DIR").expect("LAGRANGE_CACHE_DIR missing in env");
177        let cache_dir = PathBuf::from(format!("{}/pallas", cache_base_dir));
178        if !cache_dir.exists() {
179            fs::create_dir_all(&cache_dir).unwrap();
180        }
181        FileCache::new(cache_dir)
182    });
183
184    pub fn get_pallas_file_cache() -> &'static FileCache<Pallas> {
185        &PALLAS_FILE_CACHE
186    }
187}