kimchi_stubs/
lagrange_basis.rs1use 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 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 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}