perf(poseidon): eliminate heap allocations in poseidon (#712)

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) {
+func exp5state(state []ff.Element, t int) {
-	for i := 0; i < len(state); i++ {
+	for i := 0; i < t; i++ {
-		exp5(state[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) {
+func ark(state []ff.Element, c []*ff.Element, it int, t int) {
-	for i := 0; i < len(state); i++ {
+	for i := 0; i < t; i++ {
-		state[i].Add(state[i], c[it+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++ {
+		for j := 0; j < t; j++ {
-			mul.Mul(m[j][i], state[j])
+			mul.Mul(m[j][i], &state[j])
-			newState[i].Add(newState[i], mul)
+			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)
+		exp5state(state, t)
-		ark(state, C, (i+1)*t)
+		ark(state, C, (i+1)*t, t)
 		state = mix(state, t, M)
 	}
-	exp5state(state)
+	exp5state(state, t)
-	ark(state, C, (nRoundsF/2)*t)
+	ark(state, C, (nRoundsF/2)*t, t)
 	state = mix(state, t, P)
 
 	for i := 0; i < nRoundsP; i++ {
-		exp5(state[0])
+		exp5(&state[0])
-		state[0].Add(state[0], C[(nRoundsF/2+1)*t+i])
+		state[0].Add(&state[0], C[(nRoundsF/2+1)*t+i])
 
 		mul := zero()
 		newState0 := zero()
-		for j := 0; j < len(state); j++ {
+		for j := 0; j < t; j++ {
-			mul.Mul(S[(t*2-1)*i+j], state[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].Add(&state[k], mul.Mul(&state[0], S[(t*2-1)*i+t+k-1]))
-			state[k] = 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)
+		exp5state(state, t)
-		ark(state, C, (nRoundsF/2+1)*t+nRoundsP+i*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) {
+	if width > MAX_WIDTH {
-		return nil, fmt.Errorf("invalid inputs width %d, max %d", width, len(NROUNDSP)+1) //nolint:gomnd,lll
+		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 := make([]ff.Element, MAX_WIDTH)
-	state[0] = capflag
+	state[0] = *capflag
-	for i := 1; i < width; i++ {
-		state[i] = zero()
-	}
 
 	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 := make([]ff.Element, MAX_WIDTH)
-	state[0] = ff.NewElement().SetBigInt(domain)
+	state[0] = *ff.NewElement().SetBigInt(domain)
 	copy(state[1:], inp[:])
 
 	state = permute(state, t)

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
 	}
 }

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
 )