implement right way of arc unwrap

scroll-tech · Dec 19, 2024 · c7263ab · c7263ab
1 parent a15a32d
commit c7263ab
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 22 deletions.
diff --git a/ceno_zkvm/src/scheme/utils.rs b/ceno_zkvm/src/scheme/utils.rs
@@ -5,15 +5,15 @@ use ff_ext::ExtensionField;
 use itertools::Itertools;
 use multilinear_extensions::{
     commutative_op_mle_pair_pool,
-    mle::{DenseMultilinearExtension, FieldType, IntoMLE, MultilinearExtension},
+    mle::{DenseMultilinearExtension, FieldType, IntoMLE},
     op_mle_xa_b_pool, op_mle3_range_pool,
     util::ceil_log2,
     virtual_poly_v2::ArcMultilinearExtension,
 };
 
 use ff::Field;
 
-const POOL_CAP: usize = 12;
+const POOL_CAP: usize = 3;
 
 use rayon::{
     iter::{
@@ -268,15 +268,11 @@ fn try_recycle_arcpoly<E: ExtensionField>(
             Cow::Borrowed(_) => (),
             Cow::Owned(_) => {
                 let poly = poly.into_owned();
-                let poly = poly.dyn_try_unwrap().unwrap();
-
-                // let poly = Box::downcast::<MultilinearExtension<E>>(poly).unwrap();
-
-                // match poly.evaluations {
-                //     FieldType::Base(vec) => pool_b.return_to_pool(vec),
-                //     FieldType::Ext(vec) => pool_e.return_to_pool(vec),
-                //     _ => unreachable!(),
-                // };
+                match poly.arc_try_unwrap().unwrap() {
+                    FieldType::Base(vec) => pool_b.return_to_pool(vec),
+                    FieldType::Ext(vec) => pool_e.return_to_pool(vec),
+                    _ => unreachable!(),
+                };
             }
         };
     }

diff --git a/ceno_zkvm/src/utils.rs b/ceno_zkvm/src/utils.rs
@@ -260,7 +260,6 @@ impl<T> SimpleVecPool<T> {
 
     // Return an item to the pool
     pub fn return_to_pool(&mut self, item: T) {
-        println!("got return!");
         self.pool.push_back(item);
     }
 }
diff --git a/multilinear_extensions/src/mle.rs b/multilinear_extensions/src/mle.rs
@@ -63,9 +63,7 @@ pub trait MultilinearExtension<E: ExtensionField>: Send + Sync {
         }
     }
 
-    fn dyn_try_unwrap(
-        self: Arc<Self>,
-    ) -> Option<Box<dyn MultilinearExtension<E, Output = Self::Output> + Send + Sync>>;
+    fn arc_try_unwrap(self: Arc<Self>) -> Option<FieldType<E>>;
 }
 
 impl<E: ExtensionField> Debug for dyn MultilinearExtension<E, Output = DenseMultilinearExtension<E>> {
@@ -826,10 +824,10 @@ impl<E: ExtensionField> MultilinearExtension<E> for DenseMultilinearExtension<E>
         }
     }
 
-    fn dyn_try_unwrap(
-        self: Arc<Self>,
-    ) -> Option<Box<dyn MultilinearExtension<E, Output = Self::Output> + Send + Sync>> {
-        Arc::try_unwrap(self).ok().map(|it| Box::new(it) as _)
+    fn arc_try_unwrap(self: Arc<Self>) -> Option<FieldType<E>> {
+        Arc::try_unwrap(self)
+            .ok()
+            .map(|it| it.evaluations_to_owned())
     }
 }
 
@@ -1003,9 +1001,7 @@ impl<'a, E: ExtensionField> MultilinearExtension<E> for RangedMultilinearExtensi
         unimplemented!()
     }
 
-    fn dyn_try_unwrap(
-        self: Arc<Self>,
-    ) -> Option<Box<dyn MultilinearExtension<E, Output = Self::Output> + Send + Sync>> {
+    fn arc_try_unwrap(self: Arc<Self>) -> Option<FieldType<E>> {
         unimplemented!()
     }
 }