diff --git a/rln/utils.go b/rln/utils.go index a43e2a3..d87313d 100644 --- a/rln/utils.go +++ b/rln/utils.go @@ -4,6 +4,7 @@ import ( "encoding/hex" "hash" "math/big" + "sync" "github.com/consensys/gnark-crypto/ecc/bn254/fr" "golang.org/x/crypto/sha3" @@ -76,9 +77,9 @@ func Bytes128(b []byte) [128]byte { } func Flatten(b [][32]byte) []byte { - var result []byte - for _, v := range b { - result = append(result, v[:]...) + result := make([]byte, len(b)*32) + for i, v := range b { + copy(result[i*32:(i+1)*32], v[:]) } return result } @@ -137,6 +138,11 @@ type KeccakState interface { Read([]byte) (int, error) } +// Avoids multiple allocations if used frequently +var keccak256Pool = sync.Pool{New: func() interface{} { + return NewKeccakState() +}} + // NewKeccakState creates a new KeccakState func NewKeccakState() KeccakState { return sha3.NewLegacyKeccak256().(KeccakState) @@ -145,11 +151,17 @@ func NewKeccakState() KeccakState { // Keccak256 calculates and returns the Keccak256 hash of the input data. func Keccak256(data ...[]byte) []byte { b := make([]byte, 32) - d := NewKeccakState() - for _, b := range data { - d.Write(b) + h, ok := keccak256Pool.Get().(KeccakState) + if !ok { + h = NewKeccakState() } - d.Read(b) + defer keccak256Pool.Put(h) + h.Reset() + for _, b := range data { + h.Write(b) + } + + h.Read(b) return b } diff --git a/rln/utils_test.go b/rln/utils_test.go index e8ba168..c7610c2 100644 --- a/rln/utils_test.go +++ b/rln/utils_test.go @@ -2,7 +2,6 @@ package rln import ( "bytes" - "fmt" "math/big" "testing" @@ -48,7 +47,6 @@ func TestHashToBN255(t *testing.T) { msg := []byte{1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0} out := HashToBN255(msg) - fmt.Println("out: ", out) require.Equal(t, [32]byte{69, 7, 140, 46, 26, 131, 147, 30, 161, 68, 2, 5, 234, 195, 227, 223, 119, 187, 116, 97, 153, 70, 71, 254, 60, 149, 54, 109, 77, 79, 105, 20}, out)