Skip to content

Commit 4fec818

Browse files
diegokingstongabrielbosio
authored andcommitted
perf(prover): eliminate per-row Vec alloc in commit_columns_bit_reversed
Replace the per-row `Vec` allocation inside the iterator with a thread-local buffer via Rayon's `map_init` (parallel path) or a single buffer allocated outside the loop (sequential path). This avoids one heap allocation per row during Merkle tree construction.
1 parent a5d83f0 commit 4fec818

1 file changed

Lines changed: 28 additions & 12 deletions

File tree

crypto/stark/src/prover.rs

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -335,19 +335,35 @@ pub trait IsStarkProver<
335335
let num_cols = columns.len();
336336

337337
#[cfg(feature = "parallel")]
338-
let iter = (0..num_rows).into_par_iter();
339-
#[cfg(not(feature = "parallel"))]
340-
let iter = 0..num_rows;
338+
let hashed_leaves: Vec<Commitment> = {
339+
(0..num_rows)
340+
.into_par_iter()
341+
.map_init(
342+
|| vec![FieldElement::<E>::zero(); num_cols],
343+
|row_buf, row_idx| {
344+
let br_idx = reverse_index(row_idx, num_rows as u64);
345+
for col_idx in 0..num_cols {
346+
row_buf[col_idx] = columns[col_idx][br_idx].clone();
347+
}
348+
BatchedMerkleTreeBackend::<E>::hash_data(&*row_buf)
349+
},
350+
)
351+
.collect()
352+
};
341353

342-
let hashed_leaves: Vec<Commitment> = iter
343-
.map(|row_idx| {
344-
let br_idx = reverse_index(row_idx, num_rows as u64);
345-
let row: Vec<FieldElement<E>> = (0..num_cols)
346-
.map(|col_idx| columns[col_idx][br_idx].clone())
347-
.collect();
348-
BatchedMerkleTreeBackend::<E>::hash_data(&row)
349-
})
350-
.collect();
354+
#[cfg(not(feature = "parallel"))]
355+
let hashed_leaves: Vec<Commitment> = {
356+
let mut row_buf = vec![FieldElement::<E>::zero(); num_cols];
357+
(0..num_rows)
358+
.map(|row_idx| {
359+
let br_idx = reverse_index(row_idx, num_rows as u64);
360+
for col_idx in 0..num_cols {
361+
row_buf[col_idx] = columns[col_idx][br_idx].clone();
362+
}
363+
BatchedMerkleTreeBackend::<E>::hash_data(&row_buf)
364+
})
365+
.collect()
366+
};
351367

352368
let tree = BatchedMerkleTree::<E>::build_from_hashed_leaves(hashed_leaves)?;
353369
let root = tree.root;

0 commit comments

Comments
 (0)