From 7b4b9585f56e4848917e8e13996193ed99b31d38 Mon Sep 17 00:00:00 2001
From: saitima <si@matterlabs.dev>
Date: Fri, 10 Jan 2025 19:34:25 +0300
Subject: [PATCH] feat(bellman): Crs with custom allocator

---
 crates/bellman/src/kate_commitment/mod.rs | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/crates/bellman/src/kate_commitment/mod.rs b/crates/bellman/src/kate_commitment/mod.rs
index b304b86..4d0ed71 100644
--- a/crates/bellman/src/kate_commitment/mod.rs
+++ b/crates/bellman/src/kate_commitment/mod.rs
@@ -16,13 +16,30 @@ impl CrsType for CrsForMonomialForm {}
 impl CrsType for CrsForLagrangeForm {}
 impl CrsType for CrsForLagrangeFormOnCoset {}
 
-pub struct Crs<E: Engine, T: CrsType> {
+pub struct Crs<E: Engine, T: CrsType, #[cfg(feature = "allocator")] A: std::alloc::Allocator + Default = std::alloc::Global> {
+    #[cfg(feature = "allocator")]
+    pub g1_bases: Arc<Vec<E::G1Affine, A>>,
+    #[cfg(not(feature = "allocator"))]
     pub g1_bases: Arc<Vec<E::G1Affine>>,
+    #[cfg(feature = "allocator")]
+    pub g2_monomial_bases: Arc<Vec<E::G2Affine, A>>,
+    #[cfg(not(feature = "allocator"))]
     pub g2_monomial_bases: Arc<Vec<E::G2Affine>>,
 
     _marker: std::marker::PhantomData<T>,
 }
 
+#[cfg(feature = "allocator")]
+impl<E: Engine, T: CrsType, A: std::alloc::Allocator + Default> Crs<E, T, A> {
+    pub fn new_in(g1_bases: Vec<E::G1Affine, A>, g2_monomial_bases: Vec<E::G2Affine, A>) -> Self {
+        Self {
+            g1_bases: Arc::new(g1_bases),
+            g2_monomial_bases: Arc::new(g2_monomial_bases),
+            _marker: std::marker::PhantomData,
+        }
+    }
+}
+
 use crate::byteorder::BigEndian;
 use crate::byteorder::ReadBytesExt;
 use crate::byteorder::WriteBytesExt;