From c7263ab18d807b04d1713a442b48c537e7ecc452 Mon Sep 17 00:00:00 2001
From: "sm.wu" <hero78119@gmail.com>
Date: Thu, 19 Dec 2024 16:47:07 +0800
Subject: [PATCH] implement right way of arc unwrap

---
 ceno_zkvm/src/scheme/utils.rs     | 18 +++++++-----------
 ceno_zkvm/src/utils.rs            |  1 -
 multilinear_extensions/src/mle.rs | 16 ++++++----------
 3 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/ceno_zkvm/src/scheme/utils.rs b/ceno_zkvm/src/scheme/utils.rs
index 624cf89bf..4d9f81d52 100644
--- a/ceno_zkvm/src/scheme/utils.rs
+++ b/ceno_zkvm/src/scheme/utils.rs
@@ -5,7 +5,7 @@ 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,
@@ -13,7 +13,7 @@ use multilinear_extensions::{
 
 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
index 53d745194..8d89431e3 100644
--- 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
index 576f4856d..0c8d51393 100644
--- 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!()
     }
 }