Skip to content

Commit d8e432f

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 398e18f commit d8e432f

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
@@ -341,19 +341,35 @@ pub trait IsStarkProver<
341341
let num_cols = columns.len();
342342

343343
#[cfg(feature = "parallel")]
344-
let iter = (0..num_rows).into_par_iter();
345-
#[cfg(not(feature = "parallel"))]
346-
let iter = 0..num_rows;
344+
let hashed_leaves: Vec<Commitment> = {
345+
(0..num_rows)
346+
.into_par_iter()
347+
.map_init(
348+
|| vec![FieldElement::<E>::zero(); num_cols],
349+
|row_buf, row_idx| {
350+
let br_idx = reverse_index(row_idx, num_rows as u64);
351+
for col_idx in 0..num_cols {
352+
row_buf[col_idx] = columns[col_idx][br_idx].clone();
353+
}
354+
BatchedMerkleTreeBackend::<E>::hash_data(row_buf)
355+
},
356+
)
357+
.collect()
358+
};
347359

348-
let hashed_leaves: Vec<Commitment> = iter
349-
.map(|row_idx| {
350-
let br_idx = reverse_index(row_idx, num_rows as u64);
351-
let row: Vec<FieldElement<E>> = (0..num_cols)
352-
.map(|col_idx| columns[col_idx][br_idx].clone())
353-
.collect();
354-
BatchedMerkleTreeBackend::<E>::hash_data(&row)
355-
})
356-
.collect();
360+
#[cfg(not(feature = "parallel"))]
361+
let hashed_leaves: Vec<Commitment> = {
362+
let mut row_buf = vec![FieldElement::<E>::zero(); num_cols];
363+
(0..num_rows)
364+
.map(|row_idx| {
365+
let br_idx = reverse_index(row_idx, num_rows as u64);
366+
for col_idx in 0..num_cols {
367+
row_buf[col_idx] = columns[col_idx][br_idx].clone();
368+
}
369+
BatchedMerkleTreeBackend::<E>::hash_data(&row_buf)
370+
})
371+
.collect()
372+
};
357373

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

0 commit comments

Comments
 (0)