From 72d3c39081cfea0a82cee4b611a168b6976d6362 Mon Sep 17 00:00:00 2001
From: lukechampine <luke.champine@gmail.com>
Date: Fri, 10 Jan 2020 17:17:39 -0500
Subject: [PATCH] make Sum functions zero-alloc

---
 blake3.go      | 40 +++++++++++++++++++++-------------------
 blake3_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/blake3.go b/blake3.go
index 44118b3..f141aa1 100644
--- a/blake3.go
+++ b/blake3.go
@@ -314,6 +314,18 @@ func (h *Hasher) addChunkChainingValue(cv [8]uint32, totalChunks uint64) {
 	h.stackSize++
 }
 
+func (h *Hasher) finalNode() node {
+	// Starting with the node from the current chunk, compute all the
+	// parent chaining values along the right edge of the tree, until we
+	// have the root node.
+	n := h.cs.node()
+	for i := h.stackSize - 1; i >= 0; i-- {
+		n = parentNode(h.chainStack[i], n.chainingValue(), h.key, h.flags)
+	}
+	n.flags |= flagRoot
+	return n
+}
+
 // Reset implements hash.Hash.
 func (h *Hasher) Reset() {
 	h.cs = newChunkState(h.key, 0, h.flags)
@@ -359,35 +371,25 @@ func (h *Hasher) Sum(b []byte) []byte {
 
 // XOF returns an OutputReader initialized with the current hash state.
 func (h *Hasher) XOF() *OutputReader {
-	// Starting with the node from the current chunk, compute all the
-	// parent chaining values along the right edge of the tree, until we
-	// have the root node.
-	n := h.cs.node()
-	for i := h.stackSize - 1; i >= 0; i-- {
-		n = parentNode(h.chainStack[i], n.chainingValue(), h.key, h.flags)
-	}
-	n.flags |= flagRoot
 	return &OutputReader{
-		n: n,
+		n: h.finalNode(),
 	}
 }
 
 // Sum256 returns the unkeyed BLAKE3 hash of b, truncated to 256 bits.
-func Sum256(b []byte) [32]byte {
-	var out [32]byte
-	h := New(32, nil)
+func Sum256(b []byte) (out [32]byte) {
+	h := newHasher(iv, 0, 0)
 	h.Write(b)
-	h.Sum(out[:0])
-	return out
+	h.XOF().Read(out[:])
+	return
 }
 
 // Sum512 returns the unkeyed BLAKE3 hash of b, truncated to 512 bits.
-func Sum512(b []byte) [64]byte {
-	var out [64]byte
-	h := New(64, nil)
+func Sum512(b []byte) (out [64]byte) {
+	h := newHasher(iv, 0, 0)
 	h.Write(b)
-	h.Sum(out[:0])
-	return out
+	h.XOF().Read(out[:])
+	return
 }
 
 // DeriveKey derives a subkey from ctx and srcKey.
diff --git a/blake3_test.go b/blake3_test.go
index f31d737..1f885cc 100644
--- a/blake3_test.go
+++ b/blake3_test.go
@@ -120,16 +120,40 @@ func TestXOF(t *testing.T) {
 	}
 }
 
+func TestSum(t *testing.T) {
+	for _, vec := range testVectors.Cases {
+		in := testInput[:vec.InputLen]
+
+		var exp256 [32]byte
+		h := blake3.New(32, nil)
+		h.Write(in)
+		h.Sum(exp256[:0])
+		if got256 := blake3.Sum256(in); exp256 != got256 {
+			t.Errorf("Sum256 output did not match Sum output:\n\texpected: %v...\n\t     got: %v...", exp256[:10], got256[:10])
+		}
+
+		var exp512 [64]byte
+		h = blake3.New(64, nil)
+		h.Write(in)
+		h.Sum(exp512[:0])
+		if got512 := blake3.Sum512(in); exp512 != got512 {
+			t.Errorf("Sum512 output did not match Sum output:\n\texpected: %v...\n\t     got: %v...", exp512[:10], got512[:10])
+		}
+	}
+}
+
 type nopReader struct{}
 
 func (nopReader) Read(p []byte) (int, error) { return len(p), nil }
 
 func BenchmarkWrite(b *testing.B) {
+	b.ReportAllocs()
 	b.SetBytes(1)
 	io.CopyN(blake3.New(0, nil), nopReader{}, int64(b.N))
 }
 
 func BenchmarkSum256(b *testing.B) {
+	b.ReportAllocs()
 	buf := make([]byte, 1024)
 	for i := 0; i < b.N; i++ {
 		blake3.Sum256(buf)
@@ -137,6 +161,7 @@ func BenchmarkSum256(b *testing.B) {
 }
 
 func BenchmarkXOF(b *testing.B) {
+	b.ReportAllocs()
 	b.SetBytes(1)
 	io.CopyN(ioutil.Discard, blake3.New(0, nil).XOF(), int64(b.N))
 }