argon2.patch

diff --git a/argon2/argon2.go b/argon2/argon2.go
index b423fea..4ce5cac 100644
--- a/argon2/argon2.go
+++ b/argon2/argon2.go
@@ -99,6 +99,37 @@ func IDKey(password, salt []byte, time, memory uint32, threads uint8, keyLen uin
 	return deriveKey(argon2id, password, salt, nil, nil, time, memory, threads, keyLen)
 }
 
+func New(time, memory uint32, threads uint8) func([]byte, []byte, uint32) []byte {
+	if time < 1 {
+		panic("argon2: number of rounds too small")
+	}
+	if threads < 1 {
+		panic("argon2: parallelism degree too low")
+	}
+
+	mem := memory / (syncPoints * uint32(threads)) * (syncPoints * uint32(threads))
+	if mem < 2*syncPoints*uint32(threads) {
+		mem = 2 * syncPoints * uint32(threads)
+	}
+	pool := sync.Pool{
+		New: func() interface{} {
+			b := make([]block, mem)
+			return &b
+		},
+	}
+
+	return func(password, salt []byte, keyLen uint32) []byte {
+		blocks := pool.Get().(*[]block)
+		defer pool.Put(blocks)
+
+		h0 := initHash(password, salt, nil, nil, time, memory, uint32(threads), keyLen, argon2id)
+		B := initBlocks(&h0, *blocks, uint32(threads))
+		processBlocks(B, time, memory, uint32(threads), argon2id)
+		key := extractKey(B, memory, uint32(threads), keyLen)
+		return key
+	}
+}
+
 func deriveKey(mode int, password, salt, secret, data []byte, time, memory uint32, threads uint8, keyLen uint32) []byte {
 	if time < 1 {
 		panic("argon2: number of rounds too small")
@@ -112,7 +143,8 @@ func deriveKey(mode int, password, salt, secret, data []byte, time, memory uint3
 	if memory < 2*syncPoints*uint32(threads) {
 		memory = 2 * syncPoints * uint32(threads)
 	}
-	B := initBlocks(&h0, memory, uint32(threads))
+	B := make([]block, memory)
+	B = initBlocks(&h0, B, uint32(threads))
 	processBlocks(B, time, memory, uint32(threads), mode)
 	return extractKey(B, memory, uint32(threads), keyLen)
 }
@@ -155,11 +187,11 @@ func initHash(password, salt, key, data []byte, time, memory, threads, keyLen ui
 	return h0
 }
 
-func initBlocks(h0 *[blake2b.Size + 8]byte, memory, threads uint32) []block {
+func initBlocks(h0 *[blake2b.Size + 8]byte, blocks []block, threads uint32) []block {
 	var block0 [1024]byte
-	B := make([]block, memory)
+	B := blocks
 	for lane := uint32(0); lane < threads; lane++ {
-		j := lane * (memory / threads)
+		j := lane * (uint32(len(B)) / threads)
 		binary.LittleEndian.PutUint32(h0[blake2b.Size+4:], lane)
 
 		binary.LittleEndian.PutUint32(h0[blake2b.Size:], 0)
diff --git a/argon2/argon2_test.go b/argon2/argon2_test.go
index 775b97a..d08fdfc 100644
--- a/argon2/argon2_test.go
+++ b/argon2/argon2_test.go
@@ -127,6 +127,25 @@ func BenchmarkArgon2id(b *testing.B) {
 	b.Run(" Time: 5, Memory: 64 MB, Threads: 4", func(b *testing.B) { benchmarkArgon2(argon2id, 5, 64*1024, 4, 32, b) })
 }
 
+func benchmarkArgon2idNew(time, memory uint32, threads uint8, keyLen uint32, b *testing.B) {
+	f := New(time, memory, threads)
+	password := []byte("password")
+	salt := make([]byte, 32)
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_ = f(password, salt, 32)
+	}
+}
+
+func BenchmarkArgon2idNew(b *testing.B) {
+	b.Run(" Time: 3, Memory: 32 MB, Threads: 1", func(b *testing.B) { benchmarkArgon2idNew(3, 32*1024, 1, 32, b) })
+	b.Run(" Time: 4, Memory: 32 MB, Threads: 1", func(b *testing.B) { benchmarkArgon2idNew(4, 32*1024, 1, 32, b) })
+	b.Run(" Time: 5, Memory: 32 MB, Threads: 1", func(b *testing.B) { benchmarkArgon2idNew(5, 32*1024, 1, 32, b) })
+	b.Run(" Time: 3, Memory: 64 MB, Threads: 4", func(b *testing.B) { benchmarkArgon2idNew(3, 64*1024, 4, 32, b) })
+	b.Run(" Time: 4, Memory: 64 MB, Threads: 4", func(b *testing.B) { benchmarkArgon2idNew(4, 64*1024, 4, 32, b) })
+	b.Run(" Time: 5, Memory: 64 MB, Threads: 4", func(b *testing.B) { benchmarkArgon2idNew(5, 64*1024, 4, 32, b) })
+}
+