Skip to main content

poly_commitment/
lagrange_basis.rs

1use crate::{commitment::CommitmentCurve, ipa::SRS, SRS as _};
2use ark_ec::AffineRepr;
3use ark_poly::{EvaluationDomain, Radix2EvaluationDomain as D};
4use cache::LagrangeCache;
5use mina_curves::pasta::{Pallas, Vesta};
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 crate::PolyComm;
53    use ark_ec::AffineRepr;
54    use ark_poly::{EvaluationDomain, Radix2EvaluationDomain as D};
55    use core::marker::PhantomData;
56    use mina_curves::pasta::{Pallas, Vesta};
57    use std::{
58        env, fs,
59        fs::File,
60        path::{Path, PathBuf},
61        sync::LazyLock,
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        const fn new(cache_dir: PathBuf) -> Self {
95            Self {
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
114                .clone()
115                .join(format!("lagrange_basis_{srs_length}-{}", domain.size()))
116        }
117
118        fn load_lagrange_basis_from_cache(
119            &self,
120            srs_length: usize,
121            domain: &D<G::ScalarField>,
122        ) -> Option<Vec<PolyComm<G>>> {
123            let cache_key = self.lagrange_basis_cache_key(srs_length, domain);
124            if Path::exists(&cache_key) {
125                let f = File::open(cache_key.clone()).unwrap_or_else(|_| {
126                    panic!("Missing lagrange basis cache file {}", cache_key.display())
127                });
128                let basis: Vec<PolyComm<G>> =
129                    rmp_serde::decode::from_read(f).unwrap_or_else(|_| {
130                        panic!("Error decoding lagrange cache file {}", cache_key.display())
131                    });
132                Some(basis)
133            } else {
134                None
135            }
136        }
137
138        fn cache_lagrange_basis(
139            &self,
140            srs_length: usize,
141            domain: &D<G::ScalarField>,
142            basis: &[PolyComm<G>],
143        ) {
144            let cache_key = self.lagrange_basis_cache_key(srs_length, domain);
145            if !Path::exists(&cache_key) {
146                let mut f = File::create(cache_key.clone()).unwrap_or_else(|_| {
147                    panic!(
148                        "Error creating lagrange basis cache file {}",
149                        cache_key.display()
150                    )
151                });
152                rmp_serde::encode::write(&mut f, basis).unwrap_or_else(|_| {
153                    panic!(
154                        "Error encoding lagrange basis to file {}",
155                        cache_key.display()
156                    )
157                });
158            }
159        }
160    }
161
162    // The following two caches are all that we need for mina tests. These will
163    // not be initialized unless they are explicitly called.
164    static VESTA_FILE_CACHE: LazyLock<FileCache<Vesta>> = LazyLock::new(|| {
165        let cache_base_dir: String =
166            env::var("LAGRANGE_CACHE_DIR").expect("LAGRANGE_CACHE_DIR missing in env");
167        let cache_dir = PathBuf::from(format!("{cache_base_dir}/vesta"));
168        if !cache_dir.exists() {
169            fs::create_dir_all(&cache_dir).unwrap();
170        }
171        FileCache::new(cache_dir)
172    });
173
174    pub fn get_vesta_file_cache() -> &'static FileCache<Vesta> {
175        &VESTA_FILE_CACHE
176    }
177
178    static PALLAS_FILE_CACHE: LazyLock<FileCache<Pallas>> = LazyLock::new(|| {
179        let cache_base_dir: String =
180            env::var("LAGRANGE_CACHE_DIR").expect("LAGRANGE_CACHE_DIR missing in env");
181        let cache_dir = PathBuf::from(format!("{cache_base_dir}/pallas"));
182        if !cache_dir.exists() {
183            fs::create_dir_all(&cache_dir).unwrap();
184        }
185        FileCache::new(cache_dir)
186    });
187
188    pub fn get_pallas_file_cache() -> &'static FileCache<Pallas> {
189        &PALLAS_FILE_CACHE
190    }
191}