Skip to content

feat(hash): Field interfaces, range extenders, and a base-to-extension adapter - #920

Open
Tabaie wants to merge 24 commits into
masterfrom
refactor/hashes
Open

Tabaie wants to merge 24 commits into
masterfrom
refactor/hashes

Conversation

@Tabaie

@Tabaie Tabaie commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor
  • FieldHasher interface
  • Rewriting the small-field Compress function to use CompressField
  • Adapters converting a FieldHasher over the base into one for any given extension
  • A FieldMerkleDamgardHasher implementing all the interfaces MerkleDamgardHasher did, as well as FieldHasher

Note

Medium Risk
Large refactor of cryptographic hash implementations across many curves; behavior is regression-tested with vectors, but removal of MiMC state APIs and the internal rewrite could affect callers relying on old semantics.

Overview
Refactors MiMC and Poseidon2 across all supported scalar fields to sit on the shared hash stack: typed FieldHasher / FieldCompressor with LargeFieldMerkleDamgardHasher, and byte APIs via NewByteHasher plus a per-curve codec.

Each field’s fr package now aliases FieldHasher, and MiMC/Poseidon2 register constructors with RegisterFieldHasher alongside existing byte hash registration. MiMC exposes NewHasher for field elements and keeps NewMiMC / deprecated NewFieldHasher (legacyHasher) for backward-compatible WriteElement / SumElement usage. SetState / State are removed from MiMC.

Poseidon2 adds NewMerkleDamgardFieldHasher, a Compressor with CompressField, and routes the old byte Permutation.Compress through it. Docs and tests now describe left-zero padding for a final short block; JSON test vectors and hashtesting checks (Merkle–Damgard equivalence, vectors, no-allocation) replace TestSetState.

Reviewed by Cursor Bugbot for commit 36dbd98. Bugbot is set up for automated code reviews on this repo. Configure here.

Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
@Tabaie
Tabaie requested a review from a team as a code owner October 14, 2026 23:30

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit 7abb2ce. Configure here.

Comment thread hash/merkle-damgard.go Resolved
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
@Tabaie Tabaie changed the title Refactor/hashes feat(hash): Field interfaces, range extenders, and a base-to-extension adapter Oct 15, 2026
@Tabaie
Tabaie requested a review from ivokub October 15, 2026 13:44
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>

@ivokub ivokub left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Generic FieldHasher[F, D] is the right direction: one interface over large fields (D = F), small fields (D = several elements) and extension adapters. Digests are unchanged versus master on every input master accepted (checked differentially over all 8 curves × {MiMC, Poseidon2} and koalabear/babybear/goldilocks/mamabear), generated code is in sync with templates, and gnark builds/tests unchanged against this branch.

The interface is however new public API that downstream (gnark, linea prover-ray) will write generic code against, so its shape is hard to change after merge. I'd like to change it before merging. Concrete proposal first, then the reasons and measurements.

Proposed interface

// FieldHasher hashes a sequence of elements of type F into a digest of type D.
type FieldHasher[F, D any] interface {
	// Reset restores the initial state.
	Reset()
	// WriteElement writes one element.
	WriteElement(F)
	// WriteElements writes the elements; equivalent to WriteElement on each.
	WriteElements(...F)
	// PadBlock ends the current block, left-padding it with zeros.
	PadBlock()
	// SumElement returns the digest of the elements written so far, without
	// changing the state. An incomplete block is padded as by PadBlock.
	SumElement() D
	// BlockLen returns the number of elements of type F in a block (1 for
	// large fields). Not BlockSize: that is hash.Hash's size in bytes.
	BlockLen() int
}

// FieldCompressor compresses two digests into one. D is a value type
// (fr.Element, [8]koalabear.Element, ...). dst may alias left or right.
// Implementations must be safe for concurrent use.
type FieldCompressor[D any] interface {
	CompressField(dst, left, right *D)
}

Small-field digests become fixed-size arrays, e.g. [8]koalabear.Element. A generic Merkle–Damgård implementation over arrays is possible without unsafe or reflection:

type Block[F any] interface{ ~[4]F | ~[8]F | ~[12]F | ~[16]F }

type SmallFieldMerkleDamgardHasher[F any, D Block[F]] struct {
	f              FieldCompressor[D]
	iv, state, buf D
	n              int
}

Go allows indexing and len on such a type parameter (not slicing), which is enough for MD. Large fields use FieldCompressor[F]; the bn254 Compressor.CompressField(dst, left, right *fr.Element) already has that signature. Whether these hashers stay generic in hash/ or get generated per package is discussed in recommendation 10. A working prototype of both, cross-checked against this PR's hashers on 300 random write/PadBlock sequences, is at the end of this review.

Recommendations

1. Small-field digest: [N]F instead of []F (blocker: changing it later breaks every generic caller's type arguments)

Why:

  • Allocations. FieldMerkleDamgardHasher.SumField (hash/merkle-damgard.go:185) allocates on every call, and paddedPending allocates again when a block is incomplete. An array is returned by value, without allocating.
  • Type safety. The length of []F isn't part of the type. The adapter has to check it at runtime, calling SumField() in its constructor (hashadapter/template/adapters.go.tmpl:21). With an array, [8]F vs [6]F is checked by the compiler.
  • It's what callers already use. field/koalabear/vortex.Hash = [8]koalabear.Element; Permutation.Compressx16* return [][8]fr.Element; linea prover-ray crypto/koalabear/hash.Digest = [8]koalabear.Element. The getter we want later, GetPoseidon2Hasher[koalabear.Element, [8]koalabear.Element](), needs D to be an array type.
  • Digests can be compared with == and used as map keys (linea's tests already rely on got != want).

Measured on koalabear Poseidon2 (AMD 7940HS, -count=6, benchstat medians), each called from a generic function func f[F, D any, H FieldHasher[F, D]](h H, ...) as downstream code would be:

koalabear, generic caller PR ([]F, variadic) Alt generic ([8]F) Alt code-generated ([8]F)
8 elements, one call 296 ns, 1 alloc 302 ns, 0 289 ns, 0
8 elements, one per call 384 ns ±55%, 9 allocs 302 ns, 0 299 ns, 0
64 elements, one call 2.23 µs, 1 alloc 2.27 µs, 0 2.22 µs, 0
64 elements, one per call 2.84 µs, 65 allocs 2.28 µs, 0 2.26 µs, 0
1-block hashes, 16 goroutines 35.9 ns ±30%, 1 alloc — 29.2 ns, 0
floor: 1 raw permutation 266 ns

Taking one call with all elements, the array version is about as fast; the gain is zero allocations, so no GC pressure when building a Merkle tree over millions of leaves. Writing one element per call, the array version is 20–27% faster. A generic [8]F hasher is within 2% of a hand-written one, so generic vs generated is a design choice rather than a performance one (see recommendation 10).

2. Replace variadic-only WriteField(...F) with WriteElement(F) + WriteElements(...F) (blocker)

Why: when the callee isn't known at compile time (generic code with a pointer type parameter, or an interface value), Go must assume the implicit variadic slice escapes, so h.WriteField(x) allocates once per call. Measured:

Generic caller PR WriteField Alt WriteElement
bn254 Poseidon2, 16 elements one per call 16 allocs, 512 B 0
bn254 MiMC, 16 elements one per call 16 allocs, 512 B 0
koalabear, 64 elements one per call 65 allocs 0
two elements: WriteField(a, b) / WriteElements(a, b) / WriteElement(a); WriteElement(b) 3 allocs 1 / 0

A single element passed by value never escapes. The existing test TestMerkleDamgardWriteFieldAllocs doesn't catch this because it calls the concrete type, which the compiler can see through.

Keep WriteElements variadic rather than []F. The cost is the same when the caller already has a slice (es...), and linea's interface (WriteElements(...koalabear.Element)) and vortex's test hasher already match it. Document on WriteElements that literal argument lists allocate in generic code; prefer WriteElement there.

3. Names: WriteElement / WriteElements / SumElement instead of WriteField / SumField

This doesn't break compatibility. WriteField/SumField are new in this PR, so renaming them costs nothing today. The existing names make adoption easier:

  • mimc.Hasher already has WriteElement(fr.Element) and SumElement() fr.Element from the deprecated mimc.FieldHasher. With these names it satisfies the new interface with three added methods. The PR instead adds a second, duplicate set (WriteField/SumField next to WriteElement/SumElement).
  • field/koalabear/vortex/hash_test.go already has WriteElements(...) + SumElement() Hash.
  • linea prover-ray FieldHasher has Reset, WriteElements(...), Sum() Digest. Only Sum would need an alias. It can't be Sum in gnark-crypto because hash.Hash.Sum([]byte) []byte is taken.
  • mimc.SumElements([]fr.Element) fr.Element (write, then sum) keeps its meaning; nothing collides.

Then deprecate mimc.FieldHasher in favour of hash.FieldHasher[fr.Element, fr.Element] rather than keeping two parallel method sets.

4. Add Reset() and BlockLen() to the interface

  • Reset: generic consumers (Merkle tree builders, Fiat–Shamir transcripts) reuse hashers. Without it every generic signature needs an extra constraint, and the interface as written can't be used on its own.

  • BlockLen: PadBlock is defined in terms of blocks, so callers need to know the block length to use it:

    • to write full blocks (the fast path that skips copying in FieldMerkleDamgardHasher.WriteField);
    • to place domain separation at block boundaries;
    • so the extension adapter can document or align E6 (6 coordinates) against base blocks of 8 elements.

    It must not be called BlockSize. Every concrete hasher in this PR also implements hash.Hash, whose BlockSize() int is in bytes, and one type can't have both. For large fields it returns 1. With array digests the digest length is already known from the type (len(D)), so BlockLen only reports the rate. That stays meaningful if we later add sponge hashers, where rate ≠ digest size (linea already has Poseidon2SpongeHasher).

5. Small-field Compressor must be stateless (blocker: silent wrong hashes)

internal/generator/field/template/poseidon2/poseidon2.go.tmpl:600 stores a shared x []fr.Element scratch buffer. The bn254 Compressor is safe for concurrent use and the old Permutation.Compress was too, but FieldCompressor says nothing about concurrency. Sharing one compressor across goroutines, as a parallel Merkle tree builder would: 1127 of 16000 results were wrong with 8 goroutines (TestSharedCompressor in the prototype at the end). An array-based compressor with a stack buffer is safe, allocates nothing, and is as fast:

func (c Compressor) CompressField(dst, left, right *[8]fr.Element) {
	var x [16]fr.Element
	copy(x[:8], left[:]); copy(x[8:], right[:])
	_ = c.perm.Permutation(x[:])
	for i := range 8 { dst[i].Add(&right[i], &x[8+i]) }
}
koalabear 2-to-1 compression ns/op allocs safe to share
PR CompressField([]F) with buffer 276 0 no: 1127/16000 wrong
array + stack buffer 273 0 yes: 0/16000 wrong

Also state the concurrency requirement in the FieldCompressor doc.

6. MiMC Write regression: 0 → 16 allocations per 512 bytes

crypto/hash/mimc/template/mimc.go.tmpl:168-171: the local var block [BlockSize]byte escapes to the heap because d.byteOrder.Element(&block) is an interface call. Master didn't allocate on this path. The time cost is negligible (75 vs 76 µs), but this is the hot path for byte callers such as gnark's Fiat–Shamir transcript. Fix: keep the buffer in the struct (block [BlockSize]byte on Hasher), as LargeFieldMerkleDamgardHasher.block already does. Checked: local buffer 1 alloc, buffer in the struct 0.

7. Extension adapter: don't silently truncate the digest

BaseFieldHashAdapterE6.SumField returns the first 6 of the 8 base elements: 248 → 186 bits, about 93-bit collision resistance. mamabear E3 keeps 3 of 8 elements: 147 bits, about 73-bit. That's fine for sampling a Fiat–Shamir challenge but too weak for binding commitments, and the type doesn't make it visible. Alternative: the adapter adapts the input only, FieldHasher[E6, [8]koalabear.Element], keeping the base digest. Callers convert to E6 explicitly when they need a challenge (linea already does this: OutputToExt(Digest) E6). With array digests the constructor-time SumField() length check goes away too.

My preference would be to remove the extension adapter and have the downstream convert explicitly. Imo right now we would force what extension to use, but this can change (different extension, different towers). If needed then its easy to write the extension adapter in the downstream usage how it is necessary.

8. Batch hashing for AVX-512 (follow-up, but fix D now so it fits)

Through FieldHasher the single-permutation AVX-512 kernel is used automatically (Permutation → permutation16_avx512), but the 16-lane kernel can't be reached. Measured on 64-element koalabear hashes:

Path per hash
FieldHasher, purego build 9.27 µs
FieldHasher, single-permutation AVX-512 2.24 µs
Compressx16Columns (16 lanes) 0.48 µs

I checked that each lane of Compressx16Columns equals NewMerkleDamgardHasher().WriteField(row...), so a batch API produces the same digests. Suggest a separate optional interface with a looping fallback, e.g.

// BatchFieldHasher hashes inputs[i] into dst[i] independently.
type BatchFieldHasher[F, D any] interface {
	SumBatch(dst []D, inputs [][]F)
}

The koalabear implementation would map onto Compressx16Columns when lengths allow it. Its output [][8]fr.Element is exactly []D with array digests, with no copy. This doesn't need to be in this PR, but it only fits cleanly if the digest is an array.

9. Keep the field hasher and the byte hasher as separate objects

On master the small-field Poseidon2 hashers only handled bytes. This PR makes FieldMerkleDamgardHasher a FieldHasher and a hash.StateStorer on the same object. Nothing downstream needs that, and on small fields it causes real problems:

  • Two padding rules on one object. WriteField treats its input as a continuous stream. Write calls PadBlock first and left-pads the last block of every call. Measured on koalabear with 6 elements a–f:
    • WriteField(a..f) vs WriteField(a,b,c) then Write(bytes(d,e,f)): different digests.
    • WriteField(a..f) vs one Write per element: different digests.
  • State/SetState doesn't round-trip. With elements pending in an unfinished block, State() returns the padded digest. Running SetState(State()) and then the same writes gives a different digest (measured). That goes against the StateStorer contract ("allow continue the use of the current hasher").
  • Names and units clash. BlockSize()/Size() are in bytes and Sum([]byte) is taken, which pushes the field side towards awkward names (SumField, and BlockLen instead of BlockSize).
  • The byte side drives the generic design. The PF interface{ *F; SetBytesCanonical; Marshal } type parameter is only there for the byte methods. Without it the array-digest generic hasher from recommendation 1 needs only [F, D].
  • Embedding a dual object leads to silent bugs. linea prover-ray/crypto/koalabear/poseidon2/poseidon2.go:19-79 embeds gnarkposeidon2.NewMerkleDamgardHasher() as hash.StateStorer and keeps its own field state, overriding Write, Sum and Reset. SetState, Size and BlockSize fall through to the embedded gnark hasher, whose state is never used for the digest, so SetState does nothing to the result. vortex/hash_test.go:101 follows the same pattern. Their byte semantics also differ: linea's Write adds to a continuous stream and Sum(msg) writes msg first. There is no single obvious "bytes view" of an algebraic hash, so the encoding should be an explicit adapter.

Where the byte interface is still needed (keep it, as its own object):

  • The hash.Hash registry (hash.POSEIDON2_*.New(), hash.MIMC_*.New()), used by:
    • gnark: GKR Prove/Verify, std/compress checksums, the groth16 mpcsetup tests, examples/rollup;
    • gnark-crypto: vortex.HashConstructor = func() hash.Hash and fiat-shamir.
  • StateStorer, used in gnark std/internal/mimc tests.
  • Existing callers of poseidon2.NewMerkleDamgardHasher() as hash.StateStorer (prover-ray, vortex tests).

Nothing in gnark, gnark-crypto or prover-ray uses bytes and elements on the same instance, or type-asserts a hash.Hash to a field hasher. The only API that requires both on one object is the deprecated mimc.FieldHasher (it embeds hash.StateStorer and adds WriteElement). No call sites found.

Proposal:

  • Field core: poseidon2.NewHasher() *poseidon2.Hasher, implementing FieldHasher[F, D], with no byte methods (see recommendation 10).
  • Byte wrapper: poseidon2.NewMerkleDamgardHasher() hash.StateStorer, as on master, implemented on top of the field core by a small generic adapter, e.g. hash.NewByteHasher[F, D](core, codec). Write = PadBlock + parse + WriteElements + PadBlock reproduces today's digests exactly. Every byte Write ends on a block boundary, so State/SetState become an exact round trip again. It also undoes the return-type change, so the function-value break goes away.
  • MiMC (large field): a block is one element, so there's never a partial block and mixing is harmless. Keeping mimc.Hasher dual during the mimc.FieldHasher deprecation is fine; drop the byte methods from the field type when that interface is removed.

Field-side state save/restore, if needed later, should be typed and include the partial block (e.g. Clone()), not go through bytes.

10. One interface; the two generic hashers in hash/ become generated code

FieldHasher[F, D] already covers both cases: bn254 is FieldHasher[fr.Element, fr.Element], koalabear is FieldHasher[koalabear.Element, [8]koalabear.Element]. Consumers only need the interface, so FieldMerkleDamgardHasher vs LargeFieldMerkleDamgardHasher is an implementation detail that shouldn't appear in the public API of hash/.

The PR has two types because the bodies differ: on large fields every element is a block (no buffering, PadBlock does nothing, the digest is F), while on small fields elements are buffered and the last block is left-padded. Go can't merge these into one generic body. interface{ F | [8]F } doesn't compile (term cannot be a type parameter). The remaining options are worse:

  • D = [1]F for large fields: the digest would be [1]fr.Element instead of fr.Element, which hurts every bn254 caller.
  • Get the block layout through an interface: costs an indirect call per element.

Proposal: since every package is generated anyway, generate one concrete Hasher per package from a single template, branching on block length at generation time. Drop the two exported generic types from hash/, so hash/ holds only the interfaces and the byte adapter from recommendation 9. Constructors return the concrete type (poseidon2.NewHasher() *poseidon2.Hasher), checked at compile time with var _ hash.FieldHasher[fr.Element, [8]fr.Element] = (*Hasher)(nil). Calls on the concrete type are direct and don't allocate, while generic code sees only the interface. A generated koalabear hasher was within 2% of a generic one (and slightly faster) in the prototype, so this costs nothing in performance. The later GetPoseidon2Hasher[F, D]() hash.FieldHasher[F, D] getter is a generated type switch over these concrete types.

11. Minor

  • NewFieldMerkleDamgardHasher can't check that len(initialState) matches the compressor's block, because FieldCompressor doesn't expose one. With D an array, the IV is a D and the mismatch can't happen.
  • Return types changed from interfaces to concrete types (poseidon2.NewMerkleDamgardHasher → *MerkleDamgardHasher, mimc.NewFieldHasher → *mimc.Hasher). This only breaks callers that use the functions as values (var f func() hash.StateStorer = poseidon2.NewMerkleDamgardHasher); none in gnark or prover-ray. Moot for Poseidon2 if recommendation 9 is adopted; otherwise worth a changelog line.
  • Changes in what's accepted, all improvements but worth a changelog line:
    • MiMC Write with lengths that aren't a multiple of the element size (master read past the slice, or panicked on bw6-633) now left-pads the last chunk.
    • Shorter states are now accepted by MiMC SetState.
    • Poseidon2 SetState now rejects non-canonical values.
    • Write now reports the true byte count.

Compatibility summary

Status
Digests for inputs master accepted identical (all curves, MiMC + Poseidon2, all small fields, LE byte order, Permutation.Compress)
Generated code vs templates in sync (go generate ./... leaves tree clean)
gnark go build, go vet, 44 hash-related test packages: same as master
linea prover-ray compiles; poseidon2.NewMerkleDamgardHasher() assigned to hash.StateStorer still works
Proposed renames no compatibility cost: WriteField/SumField are new in this PR; mimc.Hasher already has WriteElement/SumElement

Prototype

Throwaway package (not for merging) used for the numbers above, placed as zz_alt/ in this branch. go test ./zz_alt -run Test -v runs the equivalence check against this PR's hashers and the shared-compressor test; go test ./zz_alt -run '^$' -bench . runs the benchmarks.

alt.go
// Package zzalt is a throwaway prototype of an alternative FieldHasher API, used
// to benchmark against PR #920. Not for merging.
package zzalt

import (
	kb "github.com/consensys/gnark-crypto/field/koalabear"
	p2kb "github.com/consensys/gnark-crypto/field/koalabear/poseidon2"
)

// FieldHasher is the proposed interface.
type FieldHasher[F, D any] interface {
	Reset()
	WriteElement(F)
	WriteElements(...F)
	PadBlock()
	SumElement() D
	BlockLen() int
}

// FieldCompressor compresses two digests into one; D is the digest (= block) type.
type FieldCompressor[D any] interface {
	CompressField(dst, left, right *D)
}

// Block is the set of fixed-size small-field digests.
type Block[F any] interface {
	~[4]F | ~[8]F | ~[12]F | ~[16]F
}

// SmallMD is a generic Merkle-Damgard hasher with a fixed-size array digest.
type SmallMD[F any, D Block[F]] struct {
	f              FieldCompressor[D]
	iv, state, buf D
	n              int
}

func NewSmallMD[F any, D Block[F]](f FieldCompressor[D], iv D) *SmallMD[F, D] {
	return &SmallMD[F, D]{f: f, iv: iv, state: iv}
}

func (h *SmallMD[F, D]) Reset()        { h.state, h.n = h.iv, 0 }
func (h *SmallMD[F, D]) BlockLen() int { return len(h.buf) }

func (h *SmallMD[F, D]) WriteElement(e F) {
	h.buf[h.n] = e
	h.n++
	if h.n == len(h.buf) {
		h.f.CompressField(&h.state, &h.state, &h.buf)
		h.n = 0
	}
}

func (h *SmallMD[F, D]) WriteElements(es ...F) {
	for i := range es {
		h.WriteElement(es[i])
	}
}

// padded returns the pending elements left-padded with zeros into a block.
func (h *SmallMD[F, D]) padded() D {
	var b D
	off := len(b) - h.n
	for i := range h.n {
		b[off+i] = h.buf[i]
	}
	return b
}

func (h *SmallMD[F, D]) PadBlock() {
	if h.n != 0 {
		b := h.padded()
		h.f.CompressField(&h.state, &h.state, &b)
		h.n = 0
	}
}

func (h *SmallMD[F, D]) SumElement() D {
	if h.n == 0 {
		return h.state
	}
	b := h.padded()
	var res D
	h.f.CompressField(&res, &h.state, &b)
	return res
}

// LargeMD is the large-field variant: block = digest = one element.
type LargeMD[F any] struct {
	f            FieldCompressor[F]
	iv, state, b F
}

func NewLargeMD[F any](f FieldCompressor[F], iv F) *LargeMD[F] {
	return &LargeMD[F]{f: f, iv: iv, state: iv}
}

func (h *LargeMD[F]) Reset()        { h.state = h.iv }
func (h *LargeMD[F]) BlockLen() int { return 1 }
func (h *LargeMD[F]) PadBlock()     {}
func (h *LargeMD[F]) SumElement() F { return h.state }
func (h *LargeMD[F]) WriteElement(e F) {
	h.b = e
	h.f.CompressField(&h.state, &h.state, &h.b)
}
func (h *LargeMD[F]) WriteElements(es ...F) {
	for i := range es {
		h.WriteElement(es[i])
	}
}

// KBCompressor is a stateless (concurrency-safe) width-16 koalabear compressor
// over [8]kb.Element blocks, with the same feed-forward as poseidon2.Compressor.
type KBCompressor struct{ perm *p2kb.Permutation }

func NewKBCompressor() KBCompressor { return KBCompressor{p2kb.NewDefaultPermutation()} }

func (c KBCompressor) CompressField(dst, left, right *[8]kb.Element) {
	var x [16]kb.Element
	copy(x[:8], left[:])
	copy(x[8:], right[:])
	_ = c.perm.Permutation(x[:])
	for i := range 8 {
		dst[i].Add(&right[i], &x[8+i])
	}
}

// KBMD is what a code-generated, non-generic koalabear hasher would look like.
type KBMD struct {
	perm       *p2kb.Permutation
	state, buf [8]kb.Element
	n          int
}

func NewKBMD() *KBMD { return &KBMD{perm: p2kb.NewDefaultPermutation()} }

func (h *KBMD) compress(dst, left, right *[8]kb.Element) {
	var x [16]kb.Element
	copy(x[:8], left[:])
	copy(x[8:], right[:])
	_ = h.perm.Permutation(x[:])
	for i := range 8 {
		dst[i].Add(&right[i], &x[8+i])
	}
}

func (h *KBMD) Reset()        { h.state, h.n = [8]kb.Element{}, 0 }
func (h *KBMD) BlockLen() int { return 8 }
func (h *KBMD) WriteElement(e kb.Element) {
	h.buf[h.n] = e
	h.n++
	if h.n == 8 {
		h.compress(&h.state, &h.state, &h.buf)
		h.n = 0
	}
}

func (h *KBMD) WriteElements(es ...kb.Element) {
	if h.n == 0 {
		for len(es) >= 8 {
			h.compress(&h.state, &h.state, (*[8]kb.Element)(es[:8]))
			es = es[8:]
		}
	}
	for i := range es {
		h.WriteElement(es[i])
	}
}

func (h *KBMD) padded() [8]kb.Element {
	var b [8]kb.Element
	copy(b[8-h.n:], h.buf[:h.n])
	return b
}

func (h *KBMD) PadBlock() {
	if h.n != 0 {
		b := h.padded()
		h.compress(&h.state, &h.state, &b)
		h.n = 0
	}
}

func (h *KBMD) SumElement() [8]kb.Element {
	if h.n == 0 {
		return h.state
	}
	b := h.padded()
	var res [8]kb.Element
	h.compress(&res, &h.state, &b)
	return res
}
alt_test.go
package zzalt

import (
	"math/rand/v2"
	"sync"
	"testing"

	fr "github.com/consensys/gnark-crypto/ecc/bn254/fr"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr/mimc"
	p2bn "github.com/consensys/gnark-crypto/ecc/bn254/fr/poseidon2"
	kb "github.com/consensys/gnark-crypto/field/koalabear"
	p2kb "github.com/consensys/gnark-crypto/field/koalabear/poseidon2"
	gh "github.com/consensys/gnark-crypto/hash"
)

type kbD = [8]kb.Element

func kbElems(n int) []kb.Element {
	es := make([]kb.Element, n)
	kb.Vector(es).MustSetRandom()
	return es
}

func newAltGeneric() *SmallMD[kb.Element, kbD] {
	return NewSmallMD[kb.Element, kbD](NewKBCompressor(), kbD{})
}

// ---- equivalence with the PR implementation ----

func TestAltMatchesPR(t *testing.T) {
	r := rand.New(rand.NewPCG(1, 1))
	for trial := range 300 {
		pr := p2kb.NewMerkleDamgardHasher()
		ag := newAltGeneric()
		ac := NewKBMD()
		for range r.IntN(6) {
			es := kbElems(r.IntN(20))
			switch r.IntN(3) {
			case 0:
				pr.WriteField(es...)
				ag.WriteElements(es...)
				ac.WriteElements(es...)
			case 1:
				for i := range es {
					pr.WriteField(es[i])
					ag.WriteElement(es[i])
					ac.WriteElement(es[i])
				}
			case 2:
				pr.PadBlock()
				ag.PadBlock()
				ac.PadBlock()
			}
		}
		want := pr.SumField()
		g, c := ag.SumElement(), ac.SumElement()
		for i := range 8 {
			if !want[i].Equal(&g[i]) || !want[i].Equal(&c[i]) {
				t.Fatalf("trial %d: mismatch at %d", trial, i)
			}
		}
	}
	// large field
	for trial := range 50 {
		pr := p2bn.NewMerkleDamgardHasher()
		alt := NewLargeMD[fr.Element](p2bn.NewCompressor(p2bn.NewDefaultPermutation()), fr.Element{})
		es := make([]fr.Element, trial)
		fr.Vector(es).MustSetRandom()
		pr.WriteField(es...)
		alt.WriteElements(es...)
		w, g := pr.SumField(), alt.SumElement()
		if !w.Equal(&g) {
			t.Fatalf("large trial %d mismatch", trial)
		}
	}
}

// ---- generic consumers, as gnark / linea code would be written ----

//go:noinline
func prPerElem[F, D any, H gh.FieldHasher[F, D]](h H, es []F) D {
	for i := range es {
		h.WriteField(es[i])
	}
	return h.SumField()
}

//go:noinline
func prBulk[F, D any, H gh.FieldHasher[F, D]](h H, es []F) D {
	h.WriteField(es...)
	return h.SumField()
}

//go:noinline
func prTwo[F, D any, H gh.FieldHasher[F, D]](h H, a, b F) D {
	h.WriteField(a, b)
	return h.SumField()
}

//go:noinline
func altPerElem[F, D any, H FieldHasher[F, D]](h H, es []F) D {
	for i := range es {
		h.WriteElement(es[i])
	}
	return h.SumElement()
}

//go:noinline
func altBulk[F, D any, H FieldHasher[F, D]](h H, es []F) D {
	h.WriteElements(es...)
	return h.SumElement()
}

//go:noinline
func altTwo[F, D any, H FieldHasher[F, D]](h H, a, b F) D {
	h.WriteElements(a, b)
	return h.SumElement()
}

//go:noinline
func altTwoSingle[F, D any, H FieldHasher[F, D]](h H, a, b F) D {
	h.WriteElement(a)
	h.WriteElement(b)
	return h.SumElement()
}

type resetter interface{ Reset() }

func benchKB[D any](b *testing.B, n int, h resetter, f func([]kb.Element) D) {
	es := kbElems(n)
	b.ReportAllocs()
	var s D
	for b.Loop() {
		h.Reset()
		s = f(es)
	}
	_ = s
}

func BenchmarkKB(b *testing.B) {
	for _, n := range []int{8, 64} {
		pre := "n=" + map[int]string{8: "8", 64: "64"}[n] + "/"
		b.Run(pre+"PR/perElem", func(b *testing.B) {
			h := p2kb.NewMerkleDamgardHasher()
			benchKB(b, n, h, func(es []kb.Element) []kb.Element { return prPerElem[kb.Element, []kb.Element](h, es) })
		})
		b.Run(pre+"PR/bulk", func(b *testing.B) {
			h := p2kb.NewMerkleDamgardHasher()
			benchKB(b, n, h, func(es []kb.Element) []kb.Element { return prBulk[kb.Element, []kb.Element](h, es) })
		})
		b.Run(pre+"AltGeneric/perElem", func(b *testing.B) {
			h := newAltGeneric()
			benchKB(b, n, h, func(es []kb.Element) kbD { return altPerElem[kb.Element, kbD](h, es) })
		})
		b.Run(pre+"AltGeneric/bulk", func(b *testing.B) {
			h := newAltGeneric()
			benchKB(b, n, h, func(es []kb.Element) kbD { return altBulk[kb.Element, kbD](h, es) })
		})
		b.Run(pre+"AltGenerated/perElem", func(b *testing.B) {
			h := NewKBMD()
			benchKB(b, n, h, func(es []kb.Element) kbD { return altPerElem[kb.Element, kbD](h, es) })
		})
		b.Run(pre+"AltGenerated/bulk", func(b *testing.B) {
			h := NewKBMD()
			benchKB(b, n, h, func(es []kb.Element) kbD { return altBulk[kb.Element, kbD](h, es) })
		})
	}
	b.Run("n=8/RawPermutation", func(b *testing.B) {
		p := p2kb.NewDefaultPermutation()
		var x [16]kb.Element
		b.ReportAllocs()
		for b.Loop() {
			_ = p.Permutation(x[:])
		}
	})
}

// Merkle-tree style 2-to-1 compression of digests.
func BenchmarkKBCompress(b *testing.B) {
	l, r := kbElems(8), kbElems(8)
	b.Run("PR/slice", func(b *testing.B) {
		c := p2kb.NewCompressor(p2kb.NewDefaultPermutation())
		dst := make([]kb.Element, 8)
		b.ReportAllocs()
		for b.Loop() {
			c.CompressField(dst, l, r)
		}
	})
	b.Run("Alt/array", func(b *testing.B) {
		c := NewKBCompressor()
		var dst kbD
		la, ra := kbD(l), kbD(r)
		b.ReportAllocs()
		for b.Loop() {
			c.CompressField(&dst, &la, &ra)
		}
	})
}

func BenchmarkKBTwoElems(b *testing.B) {
	es := kbElems(2)
	b.Run("PR/WriteField(a,b)", func(b *testing.B) {
		h := p2kb.NewMerkleDamgardHasher()
		b.ReportAllocs()
		for b.Loop() {
			h.Reset()
			_ = prTwo[kb.Element, []kb.Element](h, es[0], es[1])
		}
	})
	b.Run("Alt/WriteElements(a,b)", func(b *testing.B) {
		h := NewKBMD()
		b.ReportAllocs()
		for b.Loop() {
			h.Reset()
			_ = altTwo[kb.Element, kbD](h, es[0], es[1])
		}
	})
	b.Run("Alt/WriteElement(a);WriteElement(b)", func(b *testing.B) {
		h := NewKBMD()
		b.ReportAllocs()
		for b.Loop() {
			h.Reset()
			_ = altTwoSingle[kb.Element, kbD](h, es[0], es[1])
		}
	})
}

func BenchmarkBN254(b *testing.B) {
	es := make([]fr.Element, 16)
	fr.Vector(es).MustSetRandom()
	b.Run("P2/PR/perElem", func(b *testing.B) {
		h := p2bn.NewMerkleDamgardHasher()
		b.ReportAllocs()
		for b.Loop() {
			h.Reset()
			_ = prPerElem[fr.Element, fr.Element](h, es)
		}
	})
	b.Run("P2/Alt/perElem", func(b *testing.B) {
		h := NewLargeMD[fr.Element](p2bn.NewCompressor(p2bn.NewDefaultPermutation()), fr.Element{})
		b.ReportAllocs()
		for b.Loop() {
			h.Reset()
			_ = altPerElem[fr.Element, fr.Element](h, es)
		}
	})
	b.Run("MiMC/PR/perElem", func(b *testing.B) {
		h := mimc.NewFieldHasher()
		b.ReportAllocs()
		for b.Loop() {
			h.Reset()
			_ = prPerElem[fr.Element, fr.Element](h, es)
		}
	})
	b.Run("MiMC/Alt/perElem", func(b *testing.B) {
		h := NewLargeMD[fr.Element](mimc.Compressor{}, fr.Element{})
		b.ReportAllocs()
		for b.Loop() {
			h.Reset()
			_ = altPerElem[fr.Element, fr.Element](h, es)
		}
	})
}

// Concurrency: one shared compressor, many goroutines (Merkle tree pattern).
func TestSharedCompressor(t *testing.T) {
	l, r := kbElems(8), kbElems(8)
	want := make([]kb.Element, 8)
	p2kb.NewCompressor(p2kb.NewDefaultPermutation()).CompressField(want, l, r)

	check := func(name string, compress func(dst []kb.Element)) {
		var wg sync.WaitGroup
		var mu sync.Mutex
		bad := 0
		for range 8 {
			wg.Add(1)
			go func() {
				defer wg.Done()
				dst := make([]kb.Element, 8)
				for range 2000 {
					compress(dst)
					for i := range 8 {
						if !dst[i].Equal(&want[i]) {
							mu.Lock()
							bad++
							mu.Unlock()
							break
						}
					}
				}
			}()
		}
		wg.Wait()
		t.Logf("%s: %d/16000 wrong results", name, bad)
	}
	prC := p2kb.NewCompressor(p2kb.NewDefaultPermutation())
	check("PR Compressor shared", func(dst []kb.Element) { prC.CompressField(dst, l, r) })
	altC := NewKBCompressor()
	la, ra := kbD(l), kbD(r)
	check("Alt KBCompressor shared", func(dst []kb.Element) {
		var d kbD
		altC.CompressField(&d, &la, &ra)
		copy(dst, d[:])
	})
}

// Merkle-leaf pattern: many 1-block hashes in parallel, one hasher per goroutine.
func BenchmarkKBParallelLeaves(b *testing.B) {
	es := kbElems(8)
	b.Run("PR", func(b *testing.B) {
		b.ReportAllocs()
		b.RunParallel(func(pb *testing.PB) {
			h := p2kb.NewMerkleDamgardHasher()
			for pb.Next() {
				h.Reset()
				_ = prBulk[kb.Element, []kb.Element](h, es)
			}
		})
	})
	b.Run("AltGenerated", func(b *testing.B) {
		b.ReportAllocs()
		b.RunParallel(func(pb *testing.PB) {
			h := NewKBMD()
			for pb.Next() {
				h.Reset()
				_ = altBulk[kb.Element, kbD](h, es)
			}
		})
	})
}

Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
@Tabaie
Tabaie requested a review from ivokub October 23, 2026 19:57
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Signed-off-by: Arya Tabaie <arya.pourtabatabaie@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants