perf: eliminate per-element Vec<u8> allocation in Merkle tree hashing by diegokingston · Pull Request #503 · yetanotherco/lambda_vm

diegokingston · 2026-04-16T19:23:18Z

Add WriteBytes trait for zero-allocation byte serialization and use it in the prover's Merkle commitment hot paths.

Problem: as_bytes() returns Vec (heap allocation) for every field element. For CPU table commitment (2^20 rows × 74 cols), that's ~77M allocations of 8-byte Vecs. Benchmarks show this as a 1.3-1.7x slowdown.

Solution:

WriteBytes trait: write_bytes_be(&self, buf: &mut [u8]) writes to caller-provided buffer, zero heap allocations per element.
Implemented for Goldilocks (8 bytes) and cubic extension (24 bytes).
commit_columns_bit_reversed: single buffer per row, write all columns directly, hash once. Eliminates both the row Vec and per-element Vec.
commit_composition_polynomial: same approach for composition pairs.

The AsBytes trait and FieldElementVectorBackend are unchanged — only the prover hot paths are optimized. Verifier and generic code still use as_bytes() for backward compatibility.

Benchmarked in lambdaworks (~/dev/lambdaworks bench/merkle-comparison):

16 cols: 1.3-1.5x speedup, beats Plonky3 by 1.3-1.4x
64 cols: 1.6-1.8x speedup, beats Plonky3 by 1.5-1.6x

Add WriteBytes trait for zero-allocation byte serialization and use it in the prover's Merkle commitment hot paths. Problem: as_bytes() returns Vec<u8> (heap allocation) for every field element. For CPU table commitment (2^20 rows × 74 cols), that's ~77M allocations of 8-byte Vecs. Benchmarks show this as a 1.3-1.7x slowdown. Solution: - WriteBytes trait: write_bytes_be(&self, buf: &mut [u8]) writes to caller-provided buffer, zero heap allocations per element. - Implemented for Goldilocks (8 bytes) and cubic extension (24 bytes). - commit_columns_bit_reversed: single buffer per row, write all columns directly, hash once. Eliminates both the row Vec and per-element Vec. - commit_composition_polynomial: same approach for composition pairs. The AsBytes trait and FieldElementVectorBackend are unchanged — only the prover hot paths are optimized. Verifier and generic code still use as_bytes() for backward compatibility. Benchmarked in lambdaworks (~/dev/lambdaworks bench/merkle-comparison): - 16 cols: 1.3-1.5x speedup, beats Plonky3 by 1.3-1.4x - 64 cols: 1.6-1.8x speedup, beats Plonky3 by 1.5-1.6x

diegokingston · 2026-04-16T19:23:27Z

/bench 3

github-actions · 2026-04-16T19:25:43Z

Benchmark — fib_iterative_8M (median of 10)

_{Table parallelism: 32 (auto = cores / 3)}

Metric	main	PR	Δ
Peak heap	68463 MB	67127 MB	-1336 MB (-2.0%) ⚪
Prove time	35.331s	33.832s	-1.499s (-4.2%) ⚪

✅ No significant change.

⚠️ Heap spread: 5.3% (65916 MB / 67533 MB / 66527 MB / 67743 MB / 66892 MB / 67271 MB / 67727 MB / 66074 MB / 67127 MB / 69454 MB)
⚠️ Baseline heap spread: 6.4% (67325 MB / 68951 MB / 67123 MB / 69453 MB / 68463 MB / 70561 MB / 67278 MB / 68502 MB / 71295 MB / 66898 MB) — comparison may be less reliable
Consider re-running /bench

_{Commit: 65c73af · Baseline: built from main · Runner: self-hosted bench}

github-actions · 2026-04-16T19:25:49Z

Codex Code Review

Medium (Potential bug / protocol compatibility): hardcoded Keccak in prover path bypasses configured Merkle backend

In prover.rs:416 and prover.rs:752, leaf hashing is now hardcoded with sha3::Keccak256.
Previously this used BatchedMerkleTreeBackend::<...>::hash_data(...), which kept prover hashing aligned with backend configuration.
Impact: if backend type/config changes (or is swapped in another build), prover commitments can silently diverge from verifier expectations.

Low (Robustness): new WriteBytes impls can panic on short buffers

goldilocks.rs:489 and extensions_goldilocks.rs:532 slice fixed ranges without validating buf.len().
This is fine for current call sites, but as a public trait it is easy to misuse and trigger panics.

Assumptions:

I assumed current BatchedMerkleTreeBackend is still Keccak256 (so no immediate runtime mismatch today), and reviewed only the provided PR diff.

claude · 2026-04-16T19:27:42Z

Review: perf/merkle-zero-alloc-hashing

Hash correctness confirmed. The old path called FieldElementVectorBackend::hash_data which iterates and calls hasher.update(element.as_bytes()) per element. The new path writes all bytes into a contiguous buffer then calls hasher.update(&buf) once. Since Keccak256 update is streaming, these produce identical digests — and as_bytes() and write_bytes_be both use canonical_u64().to_be_bytes(), so byte-for-byte output is identical. No soundness regression.

Two issues found:

Low — write_bytes_be trait has no enforced precondition on buffer length. Each impl will panic on an undersized buffer with no helpful message. Add a debug_assert!(buf.len() >= Self::BYTE_LEN) in each implementation. See inline comment on traits.rs:44.
Medium — WriteBytes bounds added to the IsStarkProver trait signature and impl Prover are broader than needed and constitute a breaking API change. The two hot-path methods already carry their own where FieldElement<E>: WriteBytes clauses, so the trait-level bounds aren't required for correctness — they just propagate the constraint up. Removing them from the trait/impl level would keep the change non-breaking for downstream implementors. See inline comment on prover.rs:359.

github-actions · 2026-04-16T20:41:42Z

Codex Code Review

Findings

Medium - Potential correctness/config mismatch
In crypto/stark/src/prover.rs:416 and crypto/stark/src/prover.rs:752, leaf hashing is now hardcoded to sha3::Keccak256.
Previously this path used BatchedMerkleTreeBackend::<...>::hash_data(...), which followed the backend alias in config.
If BatchedMerkleTreeBackend is ever changed (or made generic), prover commitments can silently diverge from verifier expectations.
Action: hash through the configured backend (or tie the hash type to backend type) instead of hardcoding Keccak in prover logic.
Low - Public API/genericity regression
IsStarkProver now has top-level bounds requiring WriteBytes for both base and extension field elements in crypto/stark/src/prover.rs:357.
This makes the entire prover trait unavailable for fields that still support AsBytes but do not implement WriteBytes, even for methods that do not need WriteBytes.
Action: move WriteBytes bounds down to only the methods that require them (commit_columns_bit_reversed, commit_composition_polynomial) to avoid unnecessary global restriction.

Security review note

No direct memory-safety/unsafe/crypto-implementation vulnerability was found in this diff.
I could not run cargo check in this environment due rustup temp-file sandbox restrictions (/home/runner/.rustup/tmp read-only).

claude · 2026-04-16T20:45:52Z

PR Review: Zero-Allocation Merkle Leaf Hashing

The optimization direction (eliminating per-element heap allocations in the Merkle leaf hash path) is correct and meaningful. One high-severity issue and a few medium/low issues below.

High: Prover–Verifier Hash Function Coupling

The most critical issue: the prover now hardcodes sha3::Keccak256 directly instead of routing through BatchedMerkleTreeBackend::hash_data. This silently breaks if BatchedMerkleTreeBackend is ever changed to a different hash, or if as_bytes() serialization order changes — the prover and verifier would diverge without any compile error or assertion failure. The old code routed through BatchedMerkleTreeBackend::<E>::hash_data, ensuring hash function consistency.

Inline comments on specific lines below.

MauroToscano

Changing a bit the design to keep the optimization but be able to swap hashes smoothly, and use algebraics in the future

MauroToscano · 2026-04-17T15:27:58Z

/bench 10

…ytes (#506) Add hash_bytes() inherent method to FieldElementVectorBackend that surfaces the generic digest parameter D. The prover now calls BatchedMerkleTreeBackend::<E>::hash_bytes() instead of hardcoding sha3::Keccak256, so changing the config.rs type alias automatically propagates the hash function to the commit hot paths. Also adds missing debug_assert for power-of-two in commit_columns_bit_reversed for consistency with commit_composition_polynomial.

Consolidate serialization traits by adding BYTE_LEN and write_bytes_be to ByteConversion instead of having a separate WriteBytes trait. Goldilocks and cubic extension override write_bytes_be with zero-alloc implementations; other types use the default which delegates to to_bytes_be.

MauroToscano · 2026-04-17T17:35:41Z

/bench 10