From 629f14ea88b925f7ec0f09413c7923963e6ad07c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=96mer=20Faruk=20Irmak?= Date: Wed, 24 Apr 2024 11:12:27 +0300 Subject: [PATCH] perf(poseidon): eliminate heap allocations in poseidon (#712) --- crypto/poseidon/poseidon.go | 85 ++++++++++++++++---------------- crypto/poseidon/poseidon_test.go | 2 +- params/version.go | 2 +- 3 files changed, 45 insertions(+), 44 deletions(-) diff --git a/crypto/poseidon/poseidon.go b/crypto/poseidon/poseidon.go index 0902f6631d..8a696db108 100644 --- a/crypto/poseidon/poseidon.go +++ b/crypto/poseidon/poseidon.go @@ -16,6 +16,7 @@ import ( const NROUNDSF = 8 //nolint:golint var NROUNDSP = []int{56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68} //nolint:golint +const MAX_WIDTH = 18 // len(NROUNDSP)+2 func zero() *ff.Element { return ff.NewElement() @@ -28,37 +29,38 @@ func exp5(a *ff.Element) { } // exp5state perform exp5 for whole state -func exp5state(state []*ff.Element) { - for i := 0; i < len(state); i++ { - exp5(state[i]) +func exp5state(state []ff.Element, t int) { + for i := 0; i < t; i++ { + exp5(&state[i]) } } // ark computes Add-Round Key, from the paper https://eprint.iacr.org/2019/458.pdf -func ark(state []*ff.Element, c []*ff.Element, it int) { - for i := 0; i < len(state); i++ { - state[i].Add(state[i], c[it+i]) +func ark(state []ff.Element, c []*ff.Element, it int, t int) { + for i := 0; i < t; i++ { + state[i].Add(&state[i], c[it+i]) } } // mix returns [[matrix]] * [vector] -func mix(state []*ff.Element, t int, m [][]*ff.Element) []*ff.Element { +func mix(state []ff.Element, t int, m [][]*ff.Element) []ff.Element { mul := zero() - newState := make([]*ff.Element, t) + newState := make([]ff.Element, MAX_WIDTH) for i := 0; i < t; i++ { - newState[i] = zero() - } - for i := 0; i < len(state); i++ { newState[i].SetUint64(0) - for j := 0; j < len(state); j++ { - mul.Mul(m[j][i], state[j]) - newState[i].Add(newState[i], mul) + for j := 0; j < t; j++ { + mul.Mul(m[j][i], &state[j]) + newState[i].Add(&newState[i], mul) } } - return newState + + for i := 0; i < t; i++ { + state[i].Set(&newState[i]) + } + return state } -func permute(state []*ff.Element, t int) []*ff.Element { +func permute(state []ff.Element, t int) []ff.Element { nRoundsF := NROUNDSF nRoundsP := NROUNDSP[t-2] @@ -67,41 +69,40 @@ func permute(state []*ff.Element, t int) []*ff.Element { M := c.m[t-2] P := c.p[t-2] - ark(state, C, 0) + ark(state, C, 0, t) for i := 0; i < nRoundsF/2-1; i++ { - exp5state(state) - ark(state, C, (i+1)*t) + exp5state(state, t) + ark(state, C, (i+1)*t, t) state = mix(state, t, M) } - exp5state(state) - ark(state, C, (nRoundsF/2)*t) + exp5state(state, t) + ark(state, C, (nRoundsF/2)*t, t) state = mix(state, t, P) for i := 0; i < nRoundsP; i++ { - exp5(state[0]) - state[0].Add(state[0], C[(nRoundsF/2+1)*t+i]) + exp5(&state[0]) + state[0].Add(&state[0], C[(nRoundsF/2+1)*t+i]) mul := zero() newState0 := zero() - for j := 0; j < len(state); j++ { - mul.Mul(S[(t*2-1)*i+j], state[j]) + for j := 0; j < t; j++ { + mul.Mul(S[(t*2-1)*i+j], &state[j]) newState0.Add(newState0, mul) } for k := 1; k < t; k++ { - mul = zero() - state[k] = state[k].Add(state[k], mul.Mul(state[0], S[(t*2-1)*i+t+k-1])) + state[k].Add(&state[k], mul.Mul(&state[0], S[(t*2-1)*i+t+k-1])) } - state[0] = newState0 + state[0].Set(newState0) } for i := 0; i < nRoundsF/2-1; i++ { - exp5state(state) - ark(state, C, (nRoundsF/2+1)*t+nRoundsP+i*t) + exp5state(state, t) + ark(state, C, (nRoundsF/2+1)*t+nRoundsP+i*t, t) state = mix(state, t, M) } - exp5state(state) + exp5state(state, t) return mix(state, t, M) } @@ -116,8 +117,8 @@ func HashWithCap(inpBI []*big.Int, width int, nBytes int64) (*big.Int, error) { if width < 2 { return nil, fmt.Errorf("width must be ranged from 2 to 16") } - if width-2 > len(NROUNDSP) { - return nil, fmt.Errorf("invalid inputs width %d, max %d", width, len(NROUNDSP)+1) //nolint:gomnd,lll + if width > MAX_WIDTH { + return nil, fmt.Errorf("invalid inputs width %d, max %d", width, MAX_WIDTH) //nolint:gomnd,lll } // capflag = nBytes * 2^64 @@ -127,11 +128,8 @@ func HashWithCap(inpBI []*big.Int, width int, nBytes int64) (*big.Int, error) { capflag.Mul(capflag, ff.NewElement().SetBigInt(pow64)) // initialize the state - state := make([]*ff.Element, width) - state[0] = capflag - for i := 1; i < width; i++ { - state[i] = zero() - } + state := make([]ff.Element, MAX_WIDTH) + state[0] = *capflag rate := width - 1 i := 0 @@ -139,7 +137,7 @@ func HashWithCap(inpBI []*big.Int, width int, nBytes int64) (*big.Int, error) { for { // each round absorb at most `rate` elements from `inpBI` for j := 0; j < rate && i < len(inpBI); i, j = i+1, j+1 { - state[j+1].Add(state[j+1], ff.NewElement().SetBigInt(inpBI[i])) + state[j+1].Add(&state[j+1], ff.NewElement().SetBigInt(inpBI[i])) } state = permute(state, width) if i == len(inpBI) { @@ -164,10 +162,13 @@ func HashFixedWithDomain(inpBI []*big.Int, domain *big.Int) (*big.Int, error) { if !utils.CheckBigIntArrayInField(inpBI[:]) { return nil, errors.New("inputs values not inside Finite Field") } - inp := utils.BigIntArrayToElementArray(inpBI[:]) + inp := make([]ff.Element, MAX_WIDTH) + for idx, bi := range inpBI { + inp[idx].SetBigInt(bi) + } - state := make([]*ff.Element, t) - state[0] = ff.NewElement().SetBigInt(domain) + state := make([]ff.Element, MAX_WIDTH) + state[0] = *ff.NewElement().SetBigInt(domain) copy(state[1:], inp[:]) state = permute(state, t) diff --git a/crypto/poseidon/poseidon_test.go b/crypto/poseidon/poseidon_test.go index 5037c3da2d..c18207e154 100644 --- a/crypto/poseidon/poseidon_test.go +++ b/crypto/poseidon/poseidon_test.go @@ -173,6 +173,6 @@ func BenchmarkPoseidonHash(b *testing.B) { bigArray4 := []*big.Int{b1, b2, b0, b0, b0, b0} for i := 0; i < b.N; i++ { - HashFixed(bigArray4) //nolint:errcheck,gosec + HashFixedWithDomain(bigArray4, b0) //nolint:errcheck,gosec } } diff --git a/params/version.go b/params/version.go index 4465390647..dd76b01c77 100644 --- a/params/version.go +++ b/params/version.go @@ -24,7 +24,7 @@ import ( const ( VersionMajor = 5 // Major version component of the current release VersionMinor = 3 // Minor version component of the current release - VersionPatch = 2 // Patch version component of the current release + VersionPatch = 3 // Patch version component of the current release VersionMeta = "mainnet" // Version metadata to append to the version string )