Skip to content

Commit d3a2f4b

Browse files
hero78119hero78119
committed
code cosmetics
1 parent dd98c57 commit d3a2f4b

File tree

3 files changed

+31
-61
lines changed

3 files changed

+31
-61
lines changed

‎ceno_zkvm/src/scheme/utils.rs‎

Lines changed: 15 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@ use ark_std::iterable::Iterable;
44
use ff_ext::ExtensionField;
55
use itertools::Itertools;
66
use multilinear_extensions::{
7-
commutative_op_mle_pair_pool,
7+
commutative_op_mle_pair,
88
mle::{DenseMultilinearExtension, FieldType, IntoMLE},
9-
op_mle_xa_b_pool, op_mle3_range_pool,
9+
op_mle_xa_b, op_mle3_range,
1010
util::{ceil_log2, max_usable_threads},
1111
virtual_poly_v2::ArcMultilinearExtension,
1212
};
@@ -238,28 +238,12 @@ pub(crate) fn infer_tower_product_witness<E: ExtensionField>(
238238
wit_layers
239239
}
240240

241-
fn try_recycle_arcpoly<E: ExtensionField, PF1: Fn() -> Vec<E>, PF2: Fn() -> Vec<E::BaseField>>(
241+
fn optional_arcpoly_unwrap_pushback<E: ExtensionField>(
242242
poly: Cow<ArcMultilinearExtension<'_, E>>,
243-
pool_e: &mut SimpleVecPool<Vec<E>, PF1>,
244-
pool_b: &mut SimpleVecPool<Vec<E::BaseField>, PF2>,
243+
pool_e: &mut SimpleVecPool<Vec<E>, impl Fn() -> Vec<E>>,
244+
pool_b: &mut SimpleVecPool<Vec<E::BaseField>, impl Fn() -> Vec<E::BaseField>>,
245245
pool_expected_size_vec: usize,
246246
) {
247-
// fn downcast_arc<E: ExtensionField>(
248-
// arc: ArcMultilinearExtension<'_, E>,
249-
// ) -> DenseMultilinearExtension<E> {
250-
// unsafe {
251-
// // get the raw pointer from the Arc
252-
// assert_eq!(Arc::strong_count(&arc), 1);
253-
// let raw = Arc::into_raw(arc);
254-
// // cast the raw pointer to the desired concrete type
255-
// let typed_ptr = raw as *const DenseMultilinearExtension<E>;
256-
// // manually drop the Arc without dropping the value
257-
// Arc::decrement_strong_count(raw);
258-
// // reconstruct the Arc with the concrete type
259-
// // Move the value out
260-
// ptr::read(typed_ptr)
261-
// }
262-
// }
263247
let len = poly.evaluations().len();
264248
if len == pool_expected_size_vec {
265249
match poly {
@@ -348,7 +332,7 @@ pub(crate) fn wit_infer_by_expr_pool<'a, E: ExtensionField, const N: usize>(
348332
&|cow_a, cow_b, pool_e, pool_b| {
349333
let (a, b) = (cow_a.as_ref(), cow_b.as_ref());
350334
let poly =
351-
commutative_op_mle_pair_pool!(
335+
commutative_op_mle_pair!(
352336
|a, b, res| {
353337
match (a.len(), b.len()) {
354338
(1, 1) => {
@@ -401,14 +385,14 @@ pub(crate) fn wit_infer_by_expr_pool<'a, E: ExtensionField, const N: usize>(
401385
pool_e,
402386
pool_b
403387
);
404-
try_recycle_arcpoly(cow_a, pool_e, pool_b, len);
405-
try_recycle_arcpoly(cow_b, pool_e, pool_b, len);
388+
optional_arcpoly_unwrap_pushback(cow_a, pool_e, pool_b, len);
389+
optional_arcpoly_unwrap_pushback(cow_b, pool_e, pool_b, len);
406390
poly
407391
},
408392
&|cow_a, cow_b, pool_e, pool_b| {
409393
let (a, b) = (cow_a.as_ref(), cow_b.as_ref());
410394
let poly =
411-
commutative_op_mle_pair_pool!(
395+
commutative_op_mle_pair!(
412396
|a, b, res| {
413397
match (a.len(), b.len()) {
414398
(1, 1) => {
@@ -464,13 +448,13 @@ pub(crate) fn wit_infer_by_expr_pool<'a, E: ExtensionField, const N: usize>(
464448
pool_e,
465449
pool_b
466450
);
467-
try_recycle_arcpoly(cow_a, pool_e, pool_b, len);
468-
try_recycle_arcpoly(cow_b, pool_e, pool_b, len);
451+
optional_arcpoly_unwrap_pushback(cow_a, pool_e, pool_b, len);
452+
optional_arcpoly_unwrap_pushback(cow_b, pool_e, pool_b, len);
469453
poly
470454
},
471455
&|cow_x, cow_a, cow_b, pool_e, pool_b| {
472456
let (x, a, b) = (cow_x.as_ref(), cow_a.as_ref(), cow_b.as_ref());
473-
let poly = op_mle_xa_b_pool!(
457+
let poly = op_mle_xa_b!(
474458
|x, a, b, res| {
475459
let res = SyncUnsafeCell::new(res);
476460
assert_eq!(a.len(), 1);
@@ -490,9 +474,9 @@ pub(crate) fn wit_infer_by_expr_pool<'a, E: ExtensionField, const N: usize>(
490474
pool_e,
491475
pool_b
492476
);
493-
try_recycle_arcpoly(cow_a, pool_e, pool_b, len);
494-
try_recycle_arcpoly(cow_b, pool_e, pool_b, len);
495-
try_recycle_arcpoly(cow_x, pool_e, pool_b, len);
477+
optional_arcpoly_unwrap_pushback(cow_a, pool_e, pool_b, len);
478+
optional_arcpoly_unwrap_pushback(cow_b, pool_e, pool_b, len);
479+
optional_arcpoly_unwrap_pushback(cow_x, pool_e, pool_b, len);
496480
poly
497481
},
498482
pool_e,

‎ceno_zkvm/src/utils.rs‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -230,33 +230,28 @@ where
230230
result
231231
}
232232

233+
/// a simple vector pool
234+
/// not support multi-thread access
233235
pub struct SimpleVecPool<T, F: Fn() -> T> {
234236
pool: VecDeque<T>,
235237
factory_fn: F,
236238
}
237239

238240
impl<T, F: Fn() -> T> SimpleVecPool<T, F> {
239-
// Create a new pool with a factory closure
241+
// new pool with a factory closure
240242
pub fn new(init: F) -> Self {
241-
let pool = SimpleVecPool {
243+
SimpleVecPool {
242244
pool: VecDeque::new(),
243245
factory_fn: init,
244-
};
245-
246-
pool
247-
}
248-
249-
// Add a new item to the pool
250-
pub fn add(&mut self, item: T) {
251-
self.pool.push_back(item);
246+
}
252247
}
253248

254-
// Borrow an item from the pool, or create a new one if empty
249+
// borrow an item from the pool, or create a new one if empty
255250
pub fn borrow(&mut self) -> T {
256251
self.pool.pop_front().unwrap_or_else(|| (self.factory_fn)())
257252
}
258253

259-
// Return an item to the pool
254+
// push an item to the pool
260255
pub fn return_to_pool(&mut self, item: T) {
261256
self.pool.push_back(item);
262257
}

‎multilinear_extensions/src/mle.rs‎

Lines changed: 9 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1062,10 +1062,7 @@ macro_rules! op_mle3_range {
10621062
let $bb_out = $op;
10631063
$op_bb_out
10641064
}};
1065-
}
10661065

1067-
#[macro_export]
1068-
macro_rules! op_mle3_range_pool {
10691066
($x:ident, $a:ident, $b:ident, $res:ident, $x_vec:ident, $a_vec:ident, $b_vec:ident, $res_vec:ident, $op:expr, |$bb_out:ident| $op_bb_out:expr) => {{
10701067
let $x = if let Some((start, offset)) = $x.evaluations_range() {
10711068
&$x_vec[start..][..offset]
@@ -1091,7 +1088,7 @@ macro_rules! op_mle3_range_pool {
10911088

10921089
/// deal with x * a + b
10931090
#[macro_export]
1094-
macro_rules! op_mle_xa_b_pool {
1091+
macro_rules! op_mle_xa_b {
10951092
(|$x:ident, $a:ident, $b:ident, $res:ident| $op:expr, $pool_e:ident, $pool_b:ident, |$bb_out:ident| $op_bb_out:expr) => {
10961093
match (&$x.evaluations(), &$a.evaluations(), &$b.evaluations()) {
10971094
(
@@ -1100,7 +1097,7 @@ macro_rules! op_mle_xa_b_pool {
11001097
$crate::mle::FieldType::Base(b_vec),
11011098
) => {
11021099
let res_vec = $pool_b.borrow();
1103-
op_mle3_range_pool!(
1100+
op_mle3_range!(
11041101
$x,
11051102
$a,
11061103
$b,
@@ -1119,7 +1116,7 @@ macro_rules! op_mle_xa_b_pool {
11191116
$crate::mle::FieldType::Base(b_vec),
11201117
) => {
11211118
let res_vec = $pool_e.borrow();
1122-
op_mle3_range_pool!(
1119+
op_mle3_range!(
11231120
$x,
11241121
$a,
11251122
$b,
@@ -1138,7 +1135,7 @@ macro_rules! op_mle_xa_b_pool {
11381135
$crate::mle::FieldType::Ext(b_vec),
11391136
) => {
11401137
let res_vec = $pool_e.borrow();
1141-
op_mle3_range_pool!(
1138+
op_mle3_range!(
11421139
$x,
11431140
$a,
11441141
$b,
@@ -1160,7 +1157,7 @@ macro_rules! op_mle_xa_b_pool {
11601157
}
11611158
};
11621159
(|$x:ident, $a:ident, $b:ident, $res:ident| $op:expr, $pool_e:ident, $pool_b:ident) => {
1163-
op_mle_xa_b_pool!(|$x, $a, $b, $res| $op, $pool_e, $pool_b, |out| out)
1160+
op_mle_xa_b!(|$x, $a, $b, $res| $op, $pool_e, $pool_b, |out| out)
11641161
};
11651162
}
11661163

@@ -1297,15 +1294,6 @@ macro_rules! commutative_op_mle_pair {
12971294
_ => unreachable!(),
12981295
}
12991296
};
1300-
(|$a:ident, $b:ident| $op:expr) => {
1301-
commutative_op_mle_pair!(|$a, $b| $op, |out| out)
1302-
};
1303-
}
1304-
1305-
/// macro support op(a, b) and tackles type matching internally.
1306-
/// Please noted that op must satisfy commutative rule w.r.t op(b, a) operand swap.
1307-
#[macro_export]
1308-
macro_rules! commutative_op_mle_pair_pool {
13091297
(|$first:ident, $second:ident, $res:ident| $op:expr, $pool_e:ident, $pool_b:ident, |$bb_out:ident| $op_bb_out:expr) => {
13101298
match (&$first.evaluations(), &$second.evaluations()) {
13111299
($crate::mle::FieldType::Base(base1), $crate::mle::FieldType::Base(base2)) => {
@@ -1374,6 +1362,9 @@ macro_rules! commutative_op_mle_pair_pool {
13741362
}
13751363
};
13761364
(|$a:ident, $b:ident, $res:ident| $op:expr, $pool_e:ident, $pool_b:ident) => {
1377-
commutative_op_mle_pair_pool!(|$a, $b, $res| $op, $pool_e, $pool_b, |out| out)
1365+
commutative_op_mle_pair!(|$a, $b, $res| $op, $pool_e, $pool_b, |out| out)
1366+
};
1367+
(|$a:ident, $b:ident| $op:expr) => {
1368+
commutative_op_mle_pair!(|$a, $b| $op, |out| out)
13781369
};
13791370
}

0 commit comments

Comments
 (0)