From 3bac89032bf93919d54a8648b8d1bd9a367a3c8b Mon Sep 17 00:00:00 2001
From: saitima <si@matterlabs.dev>
Date: Wed, 13 Dec 2023 13:30:21 +0100
Subject: [PATCH] WitnessVec is abstract over Allocator

---
 src/cs/implementations/convenience.rs |  4 +++-
 src/cs/implementations/witness.rs     | 29 ++++++++++++++++++---------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/cs/implementations/convenience.rs b/src/cs/implementations/convenience.rs
index c5b108f..e3f97b0 100644
--- a/src/cs/implementations/convenience.rs
+++ b/src/cs/implementations/convenience.rs
@@ -8,6 +8,7 @@ use super::verifier::VerificationKey;
 use super::witness::WitnessVec;
 use crate::cs::cs_builder_verifier::CsVerifierBuilder;
 use crate::cs::oracle::merkle_tree::MerkleTreeWithCap;
+use crate::cs::traits::GoodAllocator;
 
 use super::transcript::Transcript;
 use crate::cs::traits::circuit::Circuit;
@@ -160,9 +161,10 @@ impl<
         TR: Transcript<F>,
         H: TreeHasher<F, Output = TR::CompatibleCap>,
         POW: PoWRunner,
+        A: GoodAllocator,
     >(
         &self,
-        witness_vector: &WitnessVec<F>,
+        witness_vector: &WitnessVec<F, A>,
         proof_config: ProofConfig,
         setup_base: &SetupBaseStorage<F, P>,
         setup: &SetupStorage<F, P>,
diff --git a/src/cs/implementations/witness.rs b/src/cs/implementations/witness.rs
index e7617b6..b586342 100644
--- a/src/cs/implementations/witness.rs
+++ b/src/cs/implementations/witness.rs
@@ -1,4 +1,5 @@
 use crate::cs::implementations::fast_serialization::MemcopySerializable;
+use crate::cs::traits::GoodAllocator;
 use crate::dag::WitnessSource;
 use std::alloc::Global;
 use std::sync::atomic::AtomicU32;
@@ -29,13 +30,20 @@ pub struct WitnessSet<F: SmallField> {
 #[derive(Derivative, serde::Serialize, serde::Deserialize)]
 #[derivative(Clone, Debug)]
 #[serde(bound = "F: serde::Serialize + serde::de::DeserializeOwned")]
-pub struct WitnessVec<F: SmallField> {
+pub struct WitnessVec<F: SmallField, A: GoodAllocator = Global> {
     pub public_inputs_locations: Vec<(usize, usize)>,
-    pub all_values: Vec<F>,
-    pub multiplicities: Vec<u32>,
+    #[serde(serialize_with = "crate::utils::serialize_vec_with_allocator")]
+    #[serde(deserialize_with = "crate::utils::deserialize_vec_with_allocator")]
+    pub all_values: Vec<F, A>,
+    #[serde(serialize_with = "crate::utils::serialize_vec_with_allocator")]
+    #[serde(deserialize_with = "crate::utils::deserialize_vec_with_allocator")]
+    pub multiplicities: Vec<u32, A>,
 }
 
-impl<F: SmallField> MemcopySerializable for WitnessVec<F> {
+impl<F: SmallField, A: GoodAllocator> MemcopySerializable for WitnessVec<F, A>
+where
+    A: 'static,
+{
     fn read_from_buffer<R: std::io::Read>(mut src: R) -> Result<Self, Box<dyn std::error::Error>> {
         let public_inputs_locations = read_vec_from_buffer(&mut src)?;
         let all_values = MemcopySerializable::read_from_buffer(&mut src)?;
@@ -428,7 +436,7 @@ impl<
         result
     }
 
-    pub fn materialize_witness_vec(&mut self) -> WitnessVec<F> {
+    pub fn materialize_witness_vec<A: GoodAllocator>(&mut self) -> WitnessVec<F, A> {
         assert!(
             CFG::WitnessConfig::EVALUATE_WITNESS,
             "CS is not configured to have witness available"
@@ -454,7 +462,7 @@ impl<
         assert!(max_idx > 0);
 
         // we should do memcopy instead later on
-        let mut all_values = Vec::with_capacity(max_idx);
+        let mut all_values = Vec::with_capacity_in(max_idx, A::default());
         let storage_ref = &self.variables_storage.read().unwrap();
         for idx in 0..max_idx {
             let place = Place(idx as u64);
@@ -463,9 +471,10 @@ impl<
         }
 
         let multiplicities = if self.lookup_parameters.lookup_is_allowed() == false {
-            Vec::new()
+            Vec::new_in(A::default())
         } else {
-            let mut multiplicities = Vec::with_capacity(self.lookups_tables_total_len());
+            let mut multiplicities =
+                Vec::with_capacity_in(self.lookups_tables_total_len(), A::default());
             for subtable in self.lookup_multiplicities.iter() {
                 multiplicities.extend(
                     subtable
@@ -484,9 +493,9 @@ impl<
         }
     }
 
-    pub fn witness_set_from_witness_vec(
+    pub fn witness_set_from_witness_vec<A: GoodAllocator>(
         &self,
-        witness_set: &WitnessVec<F>,
+        witness_set: &WitnessVec<F, A>,
         vars_hint: &DenseVariablesCopyHint,
         wits_hint: &DenseWitnessCopyHint,
         worker: &Worker,