From 08875a596f5eaa682f4782b1fa6ed36120239952 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Fri, 12 Nov 2021 12:39:00 +0000
Subject: [PATCH 01/19] Buffer and pool usage fixed and improved - Fix buffers
 not put back in pool - Write to buffer passed as arguments - Decouple pools
 for encoding, digests and hashers - Improve sync.Pool usage generally -
 Improve parallel encoding of branches

---
 lib/trie/hash.go      | 289 ++++++++++++++++++++++++++----------------
 lib/trie/hash_test.go |  79 +++++-------
 lib/trie/node.go      |  84 +++++++-----
 lib/trie/node_test.go |  96 ++++++--------
 lib/trie/print.go     |  31 +++--
 lib/trie/trie.go      |  16 ++-
 6 files changed, 329 insertions(+), 266 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index 228fc7802a..439c3bd286 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -5,195 +5,262 @@ package trie
 
 import (
 	"bytes"
-	"context"
+	"errors"
+	"fmt"
 	"hash"
 	"sync"
 
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 	"golang.org/x/crypto/blake2b"
-	"golang.org/x/sync/errgroup"
 )
 
-// Hasher is a wrapper around a hash function
-type hasher struct {
-	hash     hash.Hash
-	tmp      bytes.Buffer
-	parallel bool // Whether to use parallel threads when hashing
+var encodingBufferPool = &sync.Pool{
+	New: func() interface{} {
+		const initialBufferCapacity = 1900000 // 1.9MB, from checking capacities at runtime
+		b := make([]byte, 0, initialBufferCapacity)
+		return bytes.NewBuffer(b)
+	},
 }
 
-// hasherPool creates a pool of Hasher.
-var hasherPool = sync.Pool{
+var digestBufferPool = &sync.Pool{
 	New: func() interface{} {
-		h, _ := blake2b.New256(nil)
-		var buf bytes.Buffer
-		// This allocation will be helpful for encoding keys. This is the min buffer size.
-		buf.Grow(700)
-
-		return &hasher{
-			tmp:  buf,
-			hash: h,
-		}
+		const bufferCapacity = 32
+		b := make([]byte, 0, bufferCapacity)
+		return bytes.NewBuffer(b)
 	},
 }
 
-// NewHasher create new Hasher instance
-func newHasher(parallel bool) *hasher {
-	h := hasherPool.Get().(*hasher)
-	h.parallel = parallel
-	return h
+var hasherPool = &sync.Pool{
+	New: func() interface{} {
+		hasher, err := blake2b.New256(nil)
+		if err != nil {
+			panic("cannot create Blake2b-256 hasher: " + err.Error())
+		}
+		return hasher
+	},
 }
 
-func (h *hasher) returnToPool() {
-	h.tmp.Reset()
-	h.hash.Reset()
-	hasherPool.Put(h)
-}
+func hashNode(n node, digestBuffer *bytes.Buffer) (err error) {
+	encodingBuffer := encodingBufferPool.Get().(*bytes.Buffer)
+	encodingBuffer.Reset()
+	defer encodingBufferPool.Put(encodingBuffer)
+
+	const parallel = false
 
-// Hash encodes the node and then hashes it if its encoded length is > 32 bytes
-func (h *hasher) Hash(n node) (res []byte, err error) {
-	encNode, err := h.encode(n)
+	err = encodeNode(n, encodingBuffer, parallel)
 	if err != nil {
-		return nil, err
+		return fmt.Errorf("cannot encode node: %w", err)
 	}
 
 	// if length of encoded leaf is less than 32 bytes, do not hash
-	if len(encNode) < 32 {
-		return encNode, nil
+	if encodingBuffer.Len() < 32 {
+		_, err = digestBuffer.Write(encodingBuffer.Bytes())
+		return err
 	}
 
-	h.hash.Reset()
 	// otherwise, hash encoded node
-	_, err = h.hash.Write(encNode)
-	if err == nil {
-		res = h.hash.Sum(nil)
+	hasher := hasherPool.Get().(hash.Hash)
+	hasher.Reset()
+	defer hasherPool.Put(hasher)
+
+	// Note: using the sync.Pool's buffer is useful here.
+	_, err = hasher.Write(encodingBuffer.Bytes())
+	if err != nil {
+		return fmt.Errorf("cannot hash encoded node: %w", err)
 	}
 
-	return res, err
+	_, err = digestBuffer.Write(hasher.Sum(nil))
+	return err
 }
 
-// encode is the high-level function wrapping the encoding for different node types
-// encoding has the following format:
+var ErrNodeTypeUnsupported = errors.New("node type is not supported")
+
+// encodeNode writes the encoding of the node to the buffer given.
+// It is the high-level function wrapping the encoding for different
+// node types. The encoding has the following format:
 // NodeHeader | Extra partial key length | Partial Key | Value
-func (h *hasher) encode(n node) ([]byte, error) {
+func encodeNode(n node, buffer *bytes.Buffer, parallel bool) (err error) {
 	switch n := n.(type) {
 	case *branch:
-		return h.encodeBranch(n)
+		err := encodeBranch(n, buffer, parallel)
+		if err != nil {
+			return fmt.Errorf("cannot encode branch: %w", err)
+		}
+		return nil
 	case *leaf:
-		return h.encodeLeaf(n)
+		err := encodeLeaf(n, buffer)
+		if err != nil {
+			return fmt.Errorf("cannot encode leaf: %w", err)
+		}
+		n.encoding = make([]byte, buffer.Len())
+		copy(n.encoding, buffer.Bytes())
+		return nil
 	case nil:
-		return []byte{0}, nil
+		buffer.Write([]byte{0})
+		return nil
+	default:
+		return fmt.Errorf("%w: %T", ErrNodeTypeUnsupported, n)
 	}
-
-	return nil, nil
 }
 
 func encodeAndHash(n node) ([]byte, error) {
-	h := newHasher(false)
-	defer h.returnToPool()
+	buffer := digestBufferPool.Get().(*bytes.Buffer)
+	buffer.Reset()
+	defer digestBufferPool.Put(buffer)
 
-	encChild, err := h.Hash(n)
+	err := hashNode(n, buffer)
 	if err != nil {
 		return nil, err
 	}
 
-	scEncChild, err := scale.Marshal(encChild)
+	scEncChild, err := scale.Marshal(buffer.Bytes())
 	if err != nil {
 		return nil, err
 	}
 	return scEncChild, nil
 }
 
-// encodeBranch encodes a branch with the encoding specified at the top of this package
-func (h *hasher) encodeBranch(b *branch) ([]byte, error) {
+// encodeBranch encodes a branch with the encoding specified in hashedOrEncodedNode
+// to the buffer given.
+func encodeBranch(b *branch, buffer *bytes.Buffer, parallel bool) (err error) {
 	if !b.dirty && b.encoding != nil {
-		return b.encoding, nil
+		_, err = buffer.Write(b.encoding)
+		return err
 	}
-	h.tmp.Reset()
 
 	encoding, err := b.header()
-	h.tmp.Write(encoding)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	h.tmp.Write(nibblesToKeyLE(b.key))
-	h.tmp.Write(common.Uint16ToBytes(b.childrenBitmap()))
+	buffer.Write(encoding)
+	buffer.Write(nibblesToKeyLE(b.key))
+	buffer.Write(common.Uint16ToBytes(b.childrenBitmap()))
 
 	if b.value != nil {
 		bytes, err := scale.Marshal(b.value)
 		if err != nil {
-			return nil, err
-		}
-		h.tmp.Write(bytes)
-	}
-
-	if h.parallel {
-		wg, _ := errgroup.WithContext(context.Background())
-		resBuff := make([][]byte, 16)
-		for i := 0; i < 16; i++ {
-			func(i int) {
-				wg.Go(func() error {
-					child := b.children[i]
-					if child == nil {
-						return nil
-					}
-
-					var err error
-					resBuff[i], err = encodeAndHash(child)
-					if err != nil {
-						return err
-					}
-					return nil
-				})
-			}(i)
-		}
-		if err := wg.Wait(); err != nil {
-			return nil, err
+			return err
 		}
+		buffer.Write(bytes)
+	}
+
+	if parallel {
+		return encodeChildsInParallel(b.children, buffer)
+	}
+	return encodeChildsSequentially(b.children, buffer)
+}
+
+func encodeChildsInParallel(children [16]node, buffer *bytes.Buffer) (err error) {
+	type result struct {
+		index  int
+		buffer *bytes.Buffer
+	}
+
+	resultsCh := make(chan result)
+	errorCh := make(chan error)
 
-		for _, v := range resBuff {
-			if v != nil {
-				h.tmp.Write(v)
+	for i, child := range children {
+		go func(index int, child node) {
+			buffer := encodingBufferPool.Get().(*bytes.Buffer)
+			buffer.Reset()
+			// buffer is put back in the pool after processing its
+			// data in the select block below.
+
+			err := encodeChild(child, buffer)
+			if err != nil {
+				errorCh <- err
+				return
 			}
-		}
-	} else {
-		for i := 0; i < 16; i++ {
-			if child := b.children[i]; child != nil {
-				scEncChild, err := encodeAndHash(child)
-				if err != nil {
-					return nil, err
+
+			resultsCh <- result{
+				index:  index,
+				buffer: buffer,
+			}
+		}(i, child)
+	}
+
+	currentIndex := 0
+	resultBuffers := make([]*bytes.Buffer, len(children))
+	for range children {
+		select {
+		case result := <-resultsCh:
+			resultBuffers[result.index] = result.buffer
+
+			// write as many completed buffers to the result buffer.
+			for currentIndex < len(children) &&
+				resultBuffers[currentIndex] != nil {
+				// note buffer.Write copies the byte slice given as argument
+				_, writeErr := buffer.Write(resultBuffers[currentIndex].Bytes())
+				if writeErr != nil && err == nil {
+					err = writeErr
 				}
-				h.tmp.Write(scEncChild)
+
+				encodingBufferPool.Put(resultBuffers[currentIndex])
+				resultBuffers[currentIndex] = nil
+
+				currentIndex++
+			}
+		case newErr := <-errorCh:
+			if err == nil { // only set the first error we get
+				err = newErr
 			}
 		}
 	}
 
-	return h.tmp.Bytes(), nil
+	return err
 }
 
-// encodeLeaf encodes a leaf with the encoding specified at the top of this package
-func (h *hasher) encodeLeaf(l *leaf) ([]byte, error) {
-	if !l.dirty && l.encoding != nil {
-		return l.encoding, nil
+func encodeChildsSequentially(children [16]node, buffer *bytes.Buffer) (err error) {
+	for _, child := range children {
+		err = encodeChild(child, buffer)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func encodeChild(child node, buffer *bytes.Buffer) (err error) {
+	if child == nil {
+		return nil
+	}
+
+	scaleEncodedChild, err := encodeAndHash(child)
+	if err != nil {
+		return fmt.Errorf("failed to hash and scale encode child: %w", err)
+	}
+
+	_, err = buffer.Write(scaleEncodedChild)
+	if err != nil {
+		return fmt.Errorf("failed to write child to buffer: %w", err)
 	}
 
-	h.tmp.Reset()
+	return nil
+}
+
+// encodeLeaf encodes a leaf to the buffer given, with the encoding
+// specified at the top of this package.
+func encodeLeaf(l *leaf, buffer *bytes.Buffer) (err error) {
+	if !l.dirty && l.encoding != nil {
+		_, err = buffer.Write(l.encoding)
+		return err
+	}
 
 	encoding, err := l.header()
-	h.tmp.Write(encoding)
 	if err != nil {
-		return nil, err
+		return err
 	}
+	_, _ = buffer.Write(encoding)
 
-	h.tmp.Write(nibblesToKeyLE(l.key))
+	_, _ = buffer.Write(nibblesToKeyLE(l.key))
 
-	bytes, err := scale.Marshal(l.value)
+	bytes, err := scale.Marshal(l.value) // TODO scale encoder to write to buffer
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	h.tmp.Write(bytes)
-	l.encoding = h.tmp.Bytes()
-	return h.tmp.Bytes(), nil
+	_, _ = buffer.Write(bytes)
+	return nil
 }
diff --git a/lib/trie/hash_test.go b/lib/trie/hash_test.go
index cccb9f5a84..ffe46403f9 100644
--- a/lib/trie/hash_test.go
+++ b/lib/trie/hash_test.go
@@ -7,86 +7,69 @@ import (
 	"bytes"
 	"math/rand"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func generateRandBytes(size int) []byte {
-	r := *rand.New(rand.NewSource(rand.Int63()))
-	buf := make([]byte, r.Intn(size)+1)
-	r.Read(buf)
+	buf := make([]byte, rand.Intn(size)+1)
+	rand.Read(buf)
 	return buf
 }
 
 func generateRand(size int) [][]byte {
 	rt := make([][]byte, size)
-	r := *rand.New(rand.NewSource(rand.Int63()))
 	for i := range rt {
-		buf := make([]byte, r.Intn(379)+1)
-		r.Read(buf)
+		buf := make([]byte, rand.Intn(379)+1)
+		rand.Read(buf)
 		rt[i] = buf
 	}
 	return rt
 }
 
-func TestNewHasher(t *testing.T) {
-	hasher := newHasher(false)
-	defer hasher.returnToPool()
-
-	_, err := hasher.hash.Write([]byte("noot"))
-	if err != nil {
-		t.Error(err)
-	}
-
-	sum := hasher.hash.Sum(nil)
-	if sum == nil {
-		t.Error("did not sum hash")
-	}
-
-	hasher.hash.Reset()
-}
-
 func TestHashLeaf(t *testing.T) {
-	hasher := newHasher(false)
-	defer hasher.returnToPool()
-
 	n := &leaf{key: generateRandBytes(380), value: generateRandBytes(64)}
-	h, err := hasher.Hash(n)
+
+	buffer := bytes.NewBuffer(nil)
+	const parallel = false
+	err := encodeNode(n, buffer, parallel)
+
 	if err != nil {
 		t.Errorf("did not hash leaf node: %s", err)
-	} else if h == nil {
+	} else if buffer.Len() == 0 {
 		t.Errorf("did not hash leaf node: nil")
 	}
 }
 
 func TestHashBranch(t *testing.T) {
-	hasher := newHasher(false)
-	defer hasher.returnToPool()
-
 	n := &branch{key: generateRandBytes(380), value: generateRandBytes(380)}
 	n.children[3] = &leaf{key: generateRandBytes(380), value: generateRandBytes(380)}
-	h, err := hasher.Hash(n)
+
+	buffer := bytes.NewBuffer(nil)
+	const parallel = false
+	err := encodeNode(n, buffer, parallel)
+
 	if err != nil {
 		t.Errorf("did not hash branch node: %s", err)
-	} else if h == nil {
+	} else if buffer.Len() == 0 {
 		t.Errorf("did not hash branch node: nil")
 	}
 }
 
 func TestHashShort(t *testing.T) {
-	hasher := newHasher(false)
-	defer hasher.returnToPool()
-
-	n := &leaf{key: generateRandBytes(2), value: generateRandBytes(3)}
-	expected, err := hasher.encode(n)
-	if err != nil {
-		t.Fatal(err)
+	n := &leaf{
+		key:   generateRandBytes(2),
+		value: generateRandBytes(3),
 	}
 
-	h, err := hasher.Hash(n)
-	if err != nil {
-		t.Errorf("did not hash leaf node: %s", err)
-	} else if h == nil {
-		t.Errorf("did not hash leaf node: nil")
-	} else if !bytes.Equal(h[:], expected) {
-		t.Errorf("did not return encoded node padded to 32 bytes: got %s", h)
-	}
+	encodingBuffer := bytes.NewBuffer(nil)
+	const parallel = false
+	err := encodeNode(n, encodingBuffer, parallel)
+	require.NoError(t, err)
+
+	digestBuffer := bytes.NewBuffer(nil)
+	err = hashNode(n, digestBuffer)
+	require.NoError(t, err)
+	assert.Equal(t, encodingBuffer.Bytes(), digestBuffer.Bytes())
 }
diff --git a/lib/trie/node.go b/lib/trie/node.go
index 080cc6063a..8b4728babd 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -95,7 +95,7 @@ func (b *branch) copy() node {
 		generation: b.generation,
 	}
 	copy(cpy.key, b.key)
-	copy(cpy.children[:], b.children[:])
+	copy(cpy.children[:], b.children[:]) // copy interface pointers
 
 	// nil and []byte{} are encoded differently, watch out!
 	if b.value != nil {
@@ -211,71 +211,93 @@ func (b *branch) setKey(key []byte) {
 	b.key = key
 }
 
-func (b *branch) encodeAndHash() ([]byte, []byte, error) {
+func (b *branch) encodeAndHash() (encoding, hash []byte, err error) {
 	if !b.dirty && b.encoding != nil && b.hash != nil {
 		return b.encoding, b.hash, nil
 	}
 
-	hasher := newHasher(false)
-	enc, err := hasher.encodeBranch(b)
+	buffer := encodingBufferPool.Get().(*bytes.Buffer)
+	buffer.Reset()
+	defer encodingBufferPool.Put(buffer)
+
+	err = encodeBranch(b, buffer, false)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if len(enc) < 32 {
-		b.encoding = enc
-		b.hash = enc
-		return enc, enc, nil
+	bufferBytes := buffer.Bytes()
+
+	b.encoding = make([]byte, len(bufferBytes))
+	encoding = make([]byte, len(bufferBytes))
+	copy(b.encoding, bufferBytes)
+	copy(encoding, bufferBytes)
+
+	if buffer.Len() < 32 {
+		b.hash = make([]byte, len(bufferBytes))
+		hash = make([]byte, len(bufferBytes))
+		copy(b.hash, bufferBytes)
+		copy(hash, bufferBytes)
+		return encoding, hash, nil
 	}
 
-	hash, err := common.Blake2bHash(enc)
+	// Note: using the sync.Pool's buffer is useful here.
+	hashArray, err := common.Blake2bHash(buffer.Bytes())
 	if err != nil {
 		return nil, nil, err
 	}
+	b.hash = hashArray[:]
+	hash = hashArray[:]
 
-	b.encoding = enc
-	b.hash = hash[:]
-	return enc, hash[:], nil
+	return encoding, hash, nil
 }
 
-func (l *leaf) encodeAndHash() ([]byte, []byte, error) {
+func (l *leaf) encodeAndHash() (encoding, hash []byte, err error) {
 	if !l.isDirty() && l.encoding != nil && l.hash != nil {
 		return l.encoding, l.hash, nil
 	}
-	hasher := newHasher(false)
-	enc, err := hasher.encodeLeaf(l)
 
+	buffer := encodingBufferPool.Get().(*bytes.Buffer)
+	buffer.Reset()
+	defer encodingBufferPool.Put(buffer)
+
+	err = encodeLeaf(l, buffer)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if len(enc) < 32 {
-		l.encoding = enc
-		l.hash = enc
-		return enc, enc, nil
+	bufferBytes := buffer.Bytes()
+
+	l.encoding = make([]byte, len(bufferBytes))
+	encoding = make([]byte, len(bufferBytes))
+	copy(l.encoding, bufferBytes)
+	copy(encoding, bufferBytes)
+
+	if len(bufferBytes) < 32 {
+		l.hash = make([]byte, len(bufferBytes))
+		hash = make([]byte, len(bufferBytes))
+		copy(l.hash, bufferBytes)
+		copy(hash, bufferBytes)
+		return encoding, hash, nil
 	}
 
-	hash, err := common.Blake2bHash(enc)
+	// Note: using the sync.Pool's buffer is useful here.
+	hashArray, err := common.Blake2bHash(buffer.Bytes())
 	if err != nil {
 		return nil, nil, err
 	}
 
-	l.encoding = enc
-	l.hash = hash[:]
-	return enc, hash[:], nil
+	l.hash = hashArray[:]
+	hash = hashArray[:]
+
+	return encoding, hash, nil
 }
 
 func decodeBytes(in []byte) (node, error) {
-	r := &bytes.Buffer{}
-	_, err := r.Write(in)
-	if err != nil {
-		return nil, err
-	}
-
-	return decode(r)
+	buffer := bytes.NewBuffer(in)
+	return decode(buffer)
 }
 
-// Decode wraps the decoding of different node types back into a node
+// decode wraps the decoding of different node types back into a node
 func decode(r io.Reader) (node, error) {
 	header, err := readByte(r)
 	if err != nil {
diff --git a/lib/trie/node_test.go b/lib/trie/node_test.go
index 754f662c5c..4bc452cfa3 100644
--- a/lib/trie/node_test.go
+++ b/lib/trie/node_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -136,44 +137,38 @@ func TestBranchEncode(t *testing.T) {
 
 	for i, testKey := range randKeys {
 		b := &branch{key: testKey, children: [16]node{}, value: randVals[i]}
-		expected := []byte{}
+		expected := bytes.NewBuffer(nil)
 
 		header, err := b.header()
 		if err != nil {
 			t.Fatalf("Error when encoding header: %s", err)
 		}
 
-		expected = append(expected, header...)
-		expected = append(expected, nibblesToKeyLE(b.key)...)
-		expected = append(expected, common.Uint16ToBytes(b.childrenBitmap())...)
+		expected.Write(header)
+		expected.Write(nibblesToKeyLE(b.key))
+		expected.Write(common.Uint16ToBytes(b.childrenBitmap()))
 
 		enc, err := scale.Marshal(b.value)
 		if err != nil {
 			t.Fatalf("Fail when encoding value with scale: %s", err)
 		}
 
-		expected = append(expected, enc...)
+		expected.Write(enc)
 
 		for _, child := range b.children {
-			if child != nil {
-				hasher := newHasher(false)
-				defer hasher.returnToPool()
-				encChild, er := hasher.Hash(child)
-				if er != nil {
-					t.Errorf("Fail when encoding branch child: %s", er)
-				}
-				expected = append(expected, encChild[:]...)
+			if child == nil {
+				continue
 			}
-		}
 
-		hasher := newHasher(false)
-		defer hasher.returnToPool()
-		res, err := hasher.encodeBranch(b)
-		if !bytes.Equal(res, expected) {
-			t.Errorf("Fail when encoding node: got %x expected %x", res, expected)
-		} else if err != nil {
-			t.Errorf("Fail when encoding node: %s", err)
+			err := hashNode(child, expected)
+			require.NoError(t, err)
 		}
+
+		buffer := bytes.NewBuffer(nil)
+		const parallel = false
+		err = encodeBranch(b, buffer, parallel)
+		require.NoError(t, err)
+		assert.Equal(t, expected.Bytes(), buffer.Bytes())
 	}
 }
 
@@ -199,14 +194,10 @@ func TestLeafEncode(t *testing.T) {
 
 		expected = append(expected, enc...)
 
-		hasher := newHasher(false)
-		defer hasher.returnToPool()
-		res, err := hasher.encodeLeaf(l)
-		if !bytes.Equal(res, expected) {
-			t.Errorf("Fail when encoding node: got %x expected %x", res, expected)
-		} else if err != nil {
-			t.Errorf("Fail when encoding node: %s", err)
-		}
+		buffer := bytes.NewBuffer(nil)
+		err = encodeLeaf(l, buffer)
+		require.NoError(t, err)
+		assert.Equal(t, expected, buffer.Bytes())
 	}
 }
 
@@ -223,12 +214,10 @@ func TestEncodeRoot(t *testing.T) {
 				t.Errorf("Fail to get key %x with value %x: got %x", test.key, test.value, val)
 			}
 
-			hasher := newHasher(false)
-			defer hasher.returnToPool()
-			_, err := hasher.encode(trie.root)
-			if err != nil {
-				t.Errorf("Fail to encode trie root: %s", err)
-			}
+			buffer := bytes.NewBuffer(nil)
+			const parallel = false
+			err := encodeNode(trie.root, buffer, parallel)
+			require.NoError(t, err)
 		}
 	}
 }
@@ -250,18 +239,16 @@ func TestBranchDecode(t *testing.T) {
 		{key: byteArray(573), children: [16]node{}, value: []byte{0x01}},
 	}
 
-	hasher := newHasher(false)
-	defer hasher.returnToPool()
+	buffer := bytes.NewBuffer(nil)
+	const parallel = false
+
 	for _, test := range tests {
-		enc, err := hasher.encodeBranch(test)
+		err := encodeBranch(test, buffer, parallel)
 		require.NoError(t, err)
 
 		res := new(branch)
-		r := &bytes.Buffer{}
-		_, err = r.Write(enc)
-		require.NoError(t, err)
+		err = res.decode(buffer, 0)
 
-		err = res.decode(r, 0)
 		require.NoError(t, err)
 		require.Equal(t, test.key, res.key)
 		require.Equal(t, test.childrenBitmap(), res.childrenBitmap())
@@ -281,18 +268,14 @@ func TestLeafDecode(t *testing.T) {
 		{key: byteArray(573), value: []byte{0x01}, dirty: true},
 	}
 
-	hasher := newHasher(false)
-	defer hasher.returnToPool()
+	buffer := bytes.NewBuffer(nil)
+
 	for _, test := range tests {
-		enc, err := hasher.encodeLeaf(test)
+		err := encodeLeaf(test, buffer)
 		require.NoError(t, err)
 
 		res := new(leaf)
-		r := &bytes.Buffer{}
-		_, err = r.Write(enc)
-		require.NoError(t, err)
-
-		err = res.decode(r, 0)
+		err = res.decode(buffer, 0)
 		require.NoError(t, err)
 
 		res.hash = nil
@@ -320,17 +303,14 @@ func TestDecode(t *testing.T) {
 		&leaf{key: byteArray(573), value: []byte{0x01}},
 	}
 
-	hasher := newHasher(false)
-	defer hasher.returnToPool()
-	for _, test := range tests {
-		enc, err := hasher.encode(test)
-		require.NoError(t, err)
+	buffer := bytes.NewBuffer(nil)
+	const parallel = false
 
-		r := &bytes.Buffer{}
-		_, err = r.Write(enc)
+	for _, test := range tests {
+		err := encodeNode(test, buffer, parallel)
 		require.NoError(t, err)
 
-		res, err := decode(r)
+		res, err := decode(buffer)
 		require.NoError(t, err)
 
 		switch n := test.(type) {
diff --git a/lib/trie/print.go b/lib/trie/print.go
index bba1f15123..2a3cbfeb15 100644
--- a/lib/trie/print.go
+++ b/lib/trie/print.go
@@ -4,6 +4,7 @@
 package trie
 
 import (
+	"bytes"
 	"fmt"
 
 	"github.com/ChainSafe/gossamer/lib/common"
@@ -25,15 +26,22 @@ func (t *Trie) String() string {
 func (t *Trie) string(tree gotree.Tree, curr node, idx int) {
 	switch c := curr.(type) {
 	case *branch:
-		hasher := newHasher(false)
-		defer hasher.returnToPool()
-		c.encoding, _ = hasher.encode(c)
+		buffer := encodingBufferPool.Get().(*bytes.Buffer)
+		buffer.Reset()
+
+		const parallel = false
+		_ = encodeBranch(c, buffer, parallel)
+		c.encoding = buffer.Bytes()
+
 		var bstr string
 		if len(c.encoding) > 1024 {
 			bstr = fmt.Sprintf("idx=%d %s hash=%x gen=%d", idx, c.String(), common.MustBlake2bHash(c.encoding), c.generation)
 		} else {
 			bstr = fmt.Sprintf("idx=%d %s encode=%x gen=%d", idx, c.String(), c.encoding, c.generation)
 		}
+
+		encodingBufferPool.Put(buffer)
+
 		sub := tree.Add(bstr)
 		for i, child := range c.children {
 			if child != nil {
@@ -41,22 +49,23 @@ func (t *Trie) string(tree gotree.Tree, curr node, idx int) {
 			}
 		}
 	case *leaf:
-		hasher := newHasher(false)
-		defer hasher.returnToPool()
-		c.encoding, _ = hasher.encode(c)
+		buffer := encodingBufferPool.Get().(*bytes.Buffer)
+		buffer.Reset()
+
+		_ = encodeLeaf(c, buffer)
+		c.encoding = buffer.Bytes()
+
 		var bstr string
 		if len(c.encoding) > 1024 {
 			bstr = fmt.Sprintf("idx=%d %s hash=%x gen=%d", idx, c.String(), common.MustBlake2bHash(c.encoding), c.generation)
 		} else {
 			bstr = fmt.Sprintf("idx=%d %s encode=%x gen=%d", idx, c.String(), c.encoding, c.generation)
 		}
+
+		encodingBufferPool.Put(buffer)
+
 		tree.Add(bstr)
 	default:
 		return
 	}
 }
-
-// Print prints the trie through pre-order traversal
-func (t *Trie) Print() {
-	fmt.Println(t.String())
-}
diff --git a/lib/trie/trie.go b/lib/trie/trie.go
index 3a1afd7147..0705b9ba7f 100644
--- a/lib/trie/trie.go
+++ b/lib/trie/trie.go
@@ -106,11 +106,9 @@ func (t *Trie) RootNode() node { //nolint
 	return t.root
 }
 
-// EncodeRoot returns the encoded root of the trie
-func (t *Trie) EncodeRoot() ([]byte, error) {
-	h := newHasher(t.parallel)
-	defer h.returnToPool()
-	return h.encode(t.RootNode())
+// encodeRoot returns the encoded root of the trie
+func (t *Trie) encodeRoot(buffer *bytes.Buffer) (err error) {
+	return encodeNode(t.RootNode(), buffer, t.parallel)
 }
 
 // MustHash returns the hashed root of the trie. It panics if it fails to hash the root node.
@@ -125,12 +123,16 @@ func (t *Trie) MustHash() common.Hash {
 
 // Hash returns the hashed root of the trie
 func (t *Trie) Hash() (common.Hash, error) {
-	encRoot, err := t.EncodeRoot()
+	buffer := encodingBufferPool.Get().(*bytes.Buffer)
+	buffer.Reset()
+	defer encodingBufferPool.Put(buffer)
+
+	err := t.encodeRoot(buffer)
 	if err != nil {
 		return [32]byte{}, err
 	}
 
-	return common.Blake2bHash(encRoot)
+	return common.Blake2bHash(buffer.Bytes())
 }
 
 // Entries returns all the key-value pairs in the trie as a map of keys to values

From 3f1350e03e306fdd326b0dc0b9b6fa47aaf2f0b8 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Mon, 15 Nov 2021 14:46:34 +0000
Subject: [PATCH 02/19] Do not copy when not needed

---
 lib/trie/node.go | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/lib/trie/node.go b/lib/trie/node.go
index 8b4728babd..434d771bf3 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -230,13 +230,13 @@ func (b *branch) encodeAndHash() (encoding, hash []byte, err error) {
 	b.encoding = make([]byte, len(bufferBytes))
 	encoding = make([]byte, len(bufferBytes))
 	copy(b.encoding, bufferBytes)
-	copy(encoding, bufferBytes)
+	encoding = b.encoding // no need to copy
 
 	if buffer.Len() < 32 {
 		b.hash = make([]byte, len(bufferBytes))
 		hash = make([]byte, len(bufferBytes))
 		copy(b.hash, bufferBytes)
-		copy(hash, bufferBytes)
+		hash = b.hash // no need to copy
 		return encoding, hash, nil
 	}
 
@@ -246,7 +246,7 @@ func (b *branch) encodeAndHash() (encoding, hash []byte, err error) {
 		return nil, nil, err
 	}
 	b.hash = hashArray[:]
-	hash = hashArray[:]
+	hash = b.hash // no need to copy
 
 	return encoding, hash, nil
 }
@@ -268,15 +268,13 @@ func (l *leaf) encodeAndHash() (encoding, hash []byte, err error) {
 	bufferBytes := buffer.Bytes()
 
 	l.encoding = make([]byte, len(bufferBytes))
-	encoding = make([]byte, len(bufferBytes))
 	copy(l.encoding, bufferBytes)
-	copy(encoding, bufferBytes)
+	encoding = l.encoding // no need to copy
 
 	if len(bufferBytes) < 32 {
 		l.hash = make([]byte, len(bufferBytes))
-		hash = make([]byte, len(bufferBytes))
 		copy(l.hash, bufferBytes)
-		copy(hash, bufferBytes)
+		hash = l.hash // no need to copy
 		return encoding, hash, nil
 	}
 
@@ -287,7 +285,7 @@ func (l *leaf) encodeAndHash() (encoding, hash []byte, err error) {
 	}
 
 	l.hash = hashArray[:]
-	hash = hashArray[:]
+	hash = l.hash // no need to copy
 
 	return encoding, hash, nil
 }

From 67abf17c64fc955c2338279a8cc3a6295a636ca2 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Mon, 15 Nov 2021 14:47:28 +0000
Subject: [PATCH 03/19] Run `go mod tidy`

---
 go.mod | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 83a0567d4b..ee23087213 100644
--- a/go.mod
+++ b/go.mod
@@ -40,7 +40,6 @@ require (
 	github.com/urfave/cli v1.22.5
 	github.com/wasmerio/go-ext-wasm v0.3.2-0.20200326095750-0a32be6068ec
 	golang.org/x/crypto v0.0.0-20210813211128-0a44fdfbc16e
-	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 	golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1
 	google.golang.org/protobuf v1.27.1
 )
@@ -174,6 +173,7 @@ require (
 	go.uber.org/multierr v1.7.0 // indirect
 	go.uber.org/zap v1.19.0 // indirect
 	golang.org/x/net v0.0.0-20210813160813-60bc85c4be6d // indirect
+	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c // indirect
 	golang.org/x/sys v0.0.0-20210816183151-1e6c022a8912 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect

From 212b091284b61880c0e72a7512b83aa8fc5eb0ae Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Tue, 16 Nov 2021 21:45:47 +0000
Subject: [PATCH 04/19] Simplify array copying

---
 lib/trie/node.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/trie/node.go b/lib/trie/node.go
index 434d771bf3..12eca341df 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -87,7 +87,7 @@ func (b *branch) copy() node {
 	defer b.Unlock()
 	cpy := &branch{
 		key:        make([]byte, len(b.key)),
-		children:   [16]node{},
+		children:   b.children, // copy interface pointers
 		value:      nil,
 		dirty:      b.dirty,
 		hash:       make([]byte, len(b.hash)),
@@ -95,7 +95,6 @@ func (b *branch) copy() node {
 		generation: b.generation,
 	}
 	copy(cpy.key, b.key)
-	copy(cpy.children[:], b.children[:]) // copy interface pointers
 
 	// nil and []byte{} are encoded differently, watch out!
 	if b.value != nil {

From e775110ce1a2f24fa47084840291bc9be911b4dc Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 11:27:17 +0000
Subject: [PATCH 05/19] Restore comment on `encodeBranch`

---
 lib/trie/hash.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index 439c3bd286..c5cdba8c3f 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -121,7 +121,7 @@ func encodeAndHash(n node) ([]byte, error) {
 	return scEncChild, nil
 }
 
-// encodeBranch encodes a branch with the encoding specified in hashedOrEncodedNode
+// encodeBranch encodes a branch with the encoding specified at the top of this package
 // to the buffer given.
 func encodeBranch(b *branch, buffer *bytes.Buffer, parallel bool) (err error) {
 	if !b.dirty && b.encoding != nil {

From 69e3df4fccbfd6320f74e59a15430d688d5810ba Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 11:27:40 +0000
Subject: [PATCH 06/19] Remove unneeded `make` for `hash`

---
 lib/trie/node.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/lib/trie/node.go b/lib/trie/node.go
index 12eca341df..98fd229897 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -233,7 +233,6 @@ func (b *branch) encodeAndHash() (encoding, hash []byte, err error) {
 
 	if buffer.Len() < 32 {
 		b.hash = make([]byte, len(bufferBytes))
-		hash = make([]byte, len(bufferBytes))
 		copy(b.hash, bufferBytes)
 		hash = b.hash // no need to copy
 		return encoding, hash, nil

From c3d036aedb392c4e39bd43cdc056b22f80ff64e6 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 11:30:29 +0000
Subject: [PATCH 07/19] Fix buffer not put back in pool

---
 lib/trie/hash.go | 47 +++++++++++++++++++++--------------------------
 1 file changed, 21 insertions(+), 26 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index c5cdba8c3f..0122d410d3 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -156,10 +156,10 @@ func encodeChildsInParallel(children [16]node, buffer *bytes.Buffer) (err error)
 	type result struct {
 		index  int
 		buffer *bytes.Buffer
+		err    error
 	}
 
 	resultsCh := make(chan result)
-	errorCh := make(chan error)
 
 	for i, child := range children {
 		go func(index int, child node) {
@@ -169,14 +169,11 @@ func encodeChildsInParallel(children [16]node, buffer *bytes.Buffer) (err error)
 			// data in the select block below.
 
 			err := encodeChild(child, buffer)
-			if err != nil {
-				errorCh <- err
-				return
-			}
 
 			resultsCh <- result{
 				index:  index,
 				buffer: buffer,
+				err:    err,
 			}
 		}(i, child)
 	}
@@ -184,28 +181,26 @@ func encodeChildsInParallel(children [16]node, buffer *bytes.Buffer) (err error)
 	currentIndex := 0
 	resultBuffers := make([]*bytes.Buffer, len(children))
 	for range children {
-		select {
-		case result := <-resultsCh:
-			resultBuffers[result.index] = result.buffer
-
-			// write as many completed buffers to the result buffer.
-			for currentIndex < len(children) &&
-				resultBuffers[currentIndex] != nil {
-				// note buffer.Write copies the byte slice given as argument
-				_, writeErr := buffer.Write(resultBuffers[currentIndex].Bytes())
-				if writeErr != nil && err == nil {
-					err = writeErr
-				}
-
-				encodingBufferPool.Put(resultBuffers[currentIndex])
-				resultBuffers[currentIndex] = nil
-
-				currentIndex++
-			}
-		case newErr := <-errorCh:
-			if err == nil { // only set the first error we get
-				err = newErr
+		result := <-resultsCh
+		if result.err != nil && err == nil { // only set the first error we get
+			err = result.err
+		}
+
+		resultBuffers[result.index] = result.buffer
+
+		// write as many completed buffers to the result buffer.
+		for currentIndex < len(children) &&
+			resultBuffers[currentIndex] != nil {
+			// note buffer.Write copies the byte slice given as argument
+			_, writeErr := buffer.Write(resultBuffers[currentIndex].Bytes())
+			if writeErr != nil && err == nil {
+				err = writeErr
 			}
+
+			encodingBufferPool.Put(resultBuffers[currentIndex])
+			resultBuffers[currentIndex] = nil
+
+			currentIndex++
 		}
 	}
 

From e28e6494eb58132cbbe0ed9e5dd63768224bfbe4 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 14:51:21 +0000
Subject: [PATCH 08/19] return buffer write errors with context

---
 lib/trie/hash.go | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index 0122d410d3..3c6e1c1914 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -240,22 +240,36 @@ func encodeChild(child node, buffer *bytes.Buffer) (err error) {
 func encodeLeaf(l *leaf, buffer *bytes.Buffer) (err error) {
 	if !l.dirty && l.encoding != nil {
 		_, err = buffer.Write(l.encoding)
-		return err
+		if err != nil {
+			return fmt.Errorf("cannot write stored encoding to buffer: %w", err)
+		}
+		return nil
 	}
 
 	encoding, err := l.header()
 	if err != nil {
-		return err
+		return fmt.Errorf("cannot encode header: %w", err)
 	}
-	_, _ = buffer.Write(encoding)
 
-	_, _ = buffer.Write(nibblesToKeyLE(l.key))
+	_, err = buffer.Write(encoding)
+	if err != nil {
+		return fmt.Errorf("cannot write encoded header to buffer: %w", err)
+	}
+
+	_, err = buffer.Write(nibblesToKeyLE(l.key))
+	if err != nil {
+		return fmt.Errorf("cannot write LE key to buffer: %w", err)
+	}
 
 	bytes, err := scale.Marshal(l.value) // TODO scale encoder to write to buffer
 	if err != nil {
 		return err
 	}
 
-	_, _ = buffer.Write(bytes)
+	_, err = buffer.Write(bytes)
+	if err != nil {
+		return fmt.Errorf("cannot write scale encoded value to buffer: %w", err)
+	}
+
 	return nil
 }

From e2654f0e7e2b7d39ca158d259fd64fa49e3f584b Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 15:58:59 +0000
Subject: [PATCH 09/19] Change copy locks to read locks

---
 lib/trie/node.go | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/lib/trie/node.go b/lib/trie/node.go
index 98fd229897..3372fb0889 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -83,8 +83,9 @@ func (l *leaf) setGeneration(generation uint64) {
 }
 
 func (b *branch) copy() node {
-	b.Lock()
-	defer b.Unlock()
+	b.RLock()
+	defer b.RUnlock()
+
 	cpy := &branch{
 		key:        make([]byte, len(b.key)),
 		children:   b.children, // copy interface pointers
@@ -108,8 +109,9 @@ func (b *branch) copy() node {
 }
 
 func (l *leaf) copy() node {
-	l.Lock()
-	defer l.Unlock()
+	l.RLock()
+	defer l.RUnlock()
+
 	cpy := &leaf{
 		key:        make([]byte, len(l.key)),
 		value:      make([]byte, len(l.value)),

From 84c0bdc936d881174e525de8496b62c1068e26ec Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 16:07:31 +0000
Subject: [PATCH 10/19] Add RW mutex for leaf encoding field

---
 lib/trie/hash.go  |  7 +++++++
 lib/trie/node.go  | 12 ++++++++++++
 lib/trie/print.go |  3 +++
 3 files changed, 22 insertions(+)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index 3c6e1c1914..f4f3151ce4 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -93,6 +93,10 @@ func encodeNode(n node, buffer *bytes.Buffer, parallel bool) (err error) {
 		if err != nil {
 			return fmt.Errorf("cannot encode leaf: %w", err)
 		}
+
+		n.encodingMu.Lock()
+		defer n.encodingMu.Unlock()
+
 		n.encoding = make([]byte, buffer.Len())
 		copy(n.encoding, buffer.Bytes())
 		return nil
@@ -238,13 +242,16 @@ func encodeChild(child node, buffer *bytes.Buffer) (err error) {
 // encodeLeaf encodes a leaf to the buffer given, with the encoding
 // specified at the top of this package.
 func encodeLeaf(l *leaf, buffer *bytes.Buffer) (err error) {
+	l.encodingMu.RLock()
 	if !l.dirty && l.encoding != nil {
 		_, err = buffer.Write(l.encoding)
+		l.encodingMu.RUnlock()
 		if err != nil {
 			return fmt.Errorf("cannot write stored encoding to buffer: %w", err)
 		}
 		return nil
 	}
+	l.encodingMu.RUnlock()
 
 	encoding, err := l.header()
 	if err != nil {
diff --git a/lib/trie/node.go b/lib/trie/node.go
index 3372fb0889..92eb4164a3 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -69,6 +69,7 @@ type (
 		dirty      bool
 		hash       []byte
 		encoding   []byte
+		encodingMu sync.RWMutex
 		generation uint64
 		sync.RWMutex
 	}
@@ -112,6 +113,9 @@ func (l *leaf) copy() node {
 	l.RLock()
 	defer l.RUnlock()
 
+	l.encodingMu.RLock()
+	defer l.encodingMu.RUnlock()
+
 	cpy := &leaf{
 		key:        make([]byte, len(l.key)),
 		value:      make([]byte, len(l.value)),
@@ -133,7 +137,10 @@ func (b *branch) setEncodingAndHash(enc, hash []byte) {
 }
 
 func (l *leaf) setEncodingAndHash(enc, hash []byte) {
+	l.encodingMu.Lock()
 	l.encoding = enc
+	l.encodingMu.Unlock()
+
 	l.hash = hash
 }
 
@@ -252,9 +259,12 @@ func (b *branch) encodeAndHash() (encoding, hash []byte, err error) {
 }
 
 func (l *leaf) encodeAndHash() (encoding, hash []byte, err error) {
+	l.encodingMu.RLock()
 	if !l.isDirty() && l.encoding != nil && l.hash != nil {
+		l.encodingMu.RUnlock()
 		return l.encoding, l.hash, nil
 	}
+	l.encodingMu.RUnlock()
 
 	buffer := encodingBufferPool.Get().(*bytes.Buffer)
 	buffer.Reset()
@@ -267,8 +277,10 @@ func (l *leaf) encodeAndHash() (encoding, hash []byte, err error) {
 
 	bufferBytes := buffer.Bytes()
 
+	l.encodingMu.Lock()
 	l.encoding = make([]byte, len(bufferBytes))
 	copy(l.encoding, bufferBytes)
+	l.encodingMu.Unlock()
 	encoding = l.encoding // no need to copy
 
 	if len(bufferBytes) < 32 {
diff --git a/lib/trie/print.go b/lib/trie/print.go
index 2a3cbfeb15..ba72fde4a5 100644
--- a/lib/trie/print.go
+++ b/lib/trie/print.go
@@ -53,6 +53,9 @@ func (t *Trie) string(tree gotree.Tree, curr node, idx int) {
 		buffer.Reset()
 
 		_ = encodeLeaf(c, buffer)
+
+		c.encodingMu.Lock()
+		defer c.encodingMu.Unlock()
 		c.encoding = buffer.Bytes()
 
 		var bstr string

From 9c4b1ebf369be527e2f1ba890d1c9b3ab0088ced Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 21:00:08 +0000
Subject: [PATCH 11/19] Fix memory leak on encoding error

---
 lib/trie/hash.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index f4f3151ce4..24b29b46ee 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -208,6 +208,13 @@ func encodeChildsInParallel(children [16]node, buffer *bytes.Buffer) (err error)
 		}
 	}
 
+	for _, buffer := range resultBuffers {
+		if buffer == nil { // already emptied and put back in pool
+			continue
+		}
+		encodingBufferPool.Put(buffer)
+	}
+
 	return err
 }
 

From 6e9ec9af03e11511e714baf34ec2fcf27c0e685e Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 22:48:47 +0100
Subject: [PATCH 12/19] Remove unneeded encoding `make` for `branch`

Co-authored-by: noot <36753753+noot@users.noreply.github.com>
---
 lib/trie/node.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/trie/node.go b/lib/trie/node.go
index 92eb4164a3..1da67bd9d9 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -236,9 +236,8 @@ func (b *branch) encodeAndHash() (encoding, hash []byte, err error) {
 	bufferBytes := buffer.Bytes()
 
 	b.encoding = make([]byte, len(bufferBytes))
-	encoding = make([]byte, len(bufferBytes))
 	copy(b.encoding, bufferBytes)
-	encoding = b.encoding // no need to copy
+	encoding := b.encoding // no need to copy
 
 	if buffer.Len() < 32 {
 		b.hash = make([]byte, len(bufferBytes))

From 361cf753f6062c1c0316bd2bc42f5504537155b2 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 21:35:05 +0000
Subject: [PATCH 13/19] `encodeChildsInParallel` to `encodeChildrenInParallel`

---
 lib/trie/hash.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index 24b29b46ee..dea96b8091 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -151,12 +151,12 @@ func encodeBranch(b *branch, buffer *bytes.Buffer, parallel bool) (err error) {
 	}
 
 	if parallel {
-		return encodeChildsInParallel(b.children, buffer)
+		return encodeChildrenInParallel(b.children, buffer)
 	}
 	return encodeChildsSequentially(b.children, buffer)
 }
 
-func encodeChildsInParallel(children [16]node, buffer *bytes.Buffer) (err error) {
+func encodeChildrenInParallel(children [16]node, buffer *bytes.Buffer) (err error) {
 	type result struct {
 		index  int
 		buffer *bytes.Buffer

From c2aaec480330247724cd99791b706e5bd8a267e8 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 21:35:26 +0000
Subject: [PATCH 14/19] `encodeChildsSequentially` to
 `encodeChildrenSequentially`

---
 lib/trie/hash.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index dea96b8091..f85fd58f52 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -153,7 +153,7 @@ func encodeBranch(b *branch, buffer *bytes.Buffer, parallel bool) (err error) {
 	if parallel {
 		return encodeChildrenInParallel(b.children, buffer)
 	}
-	return encodeChildsSequentially(b.children, buffer)
+	return encodeChildrenSequentially(b.children, buffer)
 }
 
 func encodeChildrenInParallel(children [16]node, buffer *bytes.Buffer) (err error) {
@@ -218,7 +218,7 @@ func encodeChildrenInParallel(children [16]node, buffer *bytes.Buffer) (err erro
 	return err
 }
 
-func encodeChildsSequentially(children [16]node, buffer *bytes.Buffer) (err error) {
+func encodeChildrenSequentially(children [16]node, buffer *bytes.Buffer) (err error) {
 	for _, child := range children {
 		err = encodeChild(child, buffer)
 		if err != nil {

From de287423c259b2a8117b2224d7c1c8c78c583d25 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Thu, 18 Nov 2021 21:53:14 +0000
Subject: [PATCH 15/19] Fix commit suggestion

---
 lib/trie/node.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/trie/node.go b/lib/trie/node.go
index 1da67bd9d9..b7ae35d9db 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -237,7 +237,7 @@ func (b *branch) encodeAndHash() (encoding, hash []byte, err error) {
 
 	b.encoding = make([]byte, len(bufferBytes))
 	copy(b.encoding, bufferBytes)
-	encoding := b.encoding // no need to copy
+	encoding = b.encoding // no need to copy
 
 	if buffer.Len() < 32 {
 		b.hash = make([]byte, len(bufferBytes))

From 095f21a045f304a3e7bf56adace9a6e8f210df5c Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Fri, 19 Nov 2021 15:21:56 +0000
Subject: [PATCH 16/19] More error checking in `encodeBranch`

---
 lib/trie/hash.go | 42 +++++++++++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 9 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index f85fd58f52..4b3d5d0a9b 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -130,30 +130,54 @@ func encodeAndHash(n node) ([]byte, error) {
 func encodeBranch(b *branch, buffer *bytes.Buffer, parallel bool) (err error) {
 	if !b.dirty && b.encoding != nil {
 		_, err = buffer.Write(b.encoding)
-		return err
+		if err != nil {
+			return fmt.Errorf("cannot write stored encoded branch to buffer: %w", err)
+		}
+		return nil
 	}
 
 	encoding, err := b.header()
 	if err != nil {
-		return err
+		return fmt.Errorf("cannot encode branch header: %w", err)
+	}
+
+	_, err = buffer.Write(encoding)
+	if err != nil {
+		return fmt.Errorf("cannot write encoded branch header to buffer: %w", err)
+	}
+
+	_, err = buffer.Write(nibblesToKeyLE(b.key))
+	if err != nil {
+		return fmt.Errorf("cannot write encoded branch key to buffer: %w", err)
 	}
 
-	buffer.Write(encoding)
-	buffer.Write(nibblesToKeyLE(b.key))
-	buffer.Write(common.Uint16ToBytes(b.childrenBitmap()))
+	_, err = buffer.Write(common.Uint16ToBytes(b.childrenBitmap()))
+	if err != nil {
+		return fmt.Errorf("cannot write branch children bitmap to buffer: %w", err)
+	}
 
 	if b.value != nil {
 		bytes, err := scale.Marshal(b.value)
 		if err != nil {
-			return err
+			return fmt.Errorf("cannot scale encode branch value: %w", err)
+		}
+
+		_, err = buffer.Write(bytes)
+		if err != nil {
+			return fmt.Errorf("cannot write encoded branch value to buffer: %w", err)
 		}
-		buffer.Write(bytes)
 	}
 
 	if parallel {
-		return encodeChildrenInParallel(b.children, buffer)
+		err = encodeChildrenInParallel(b.children, buffer)
+	} else {
+		err = encodeChildrenSequentially(b.children, buffer)
+	}
+	if err != nil {
+		return fmt.Errorf("cannot encode children of branch: %w", err)
 	}
-	return encodeChildrenSequentially(b.children, buffer)
+
+	return nil
 }
 
 func encodeChildrenInParallel(children [16]node, buffer *bytes.Buffer) (err error) {

From 330195b97f66170f8c3313002b4835582e89398b Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Fri, 19 Nov 2021 15:24:09 +0000
Subject: [PATCH 17/19] Use `io.Writer` interface when possible

---
 lib/trie/hash.go | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index 4b3d5d0a9b..c4752d4e2e 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"hash"
+	"io"
 	"sync"
 
 	"github.com/ChainSafe/gossamer/lib/common"
@@ -41,7 +42,7 @@ var hasherPool = &sync.Pool{
 	},
 }
 
-func hashNode(n node, digestBuffer *bytes.Buffer) (err error) {
+func hashNode(n node, digestBuffer io.Writer) (err error) {
 	encodingBuffer := encodingBufferPool.Get().(*bytes.Buffer)
 	encodingBuffer.Reset()
 	defer encodingBufferPool.Put(encodingBuffer)
@@ -127,7 +128,7 @@ func encodeAndHash(n node) ([]byte, error) {
 
 // encodeBranch encodes a branch with the encoding specified at the top of this package
 // to the buffer given.
-func encodeBranch(b *branch, buffer *bytes.Buffer, parallel bool) (err error) {
+func encodeBranch(b *branch, buffer io.Writer, parallel bool) (err error) {
 	if !b.dirty && b.encoding != nil {
 		_, err = buffer.Write(b.encoding)
 		if err != nil {
@@ -180,7 +181,7 @@ func encodeBranch(b *branch, buffer *bytes.Buffer, parallel bool) (err error) {
 	return nil
 }
 
-func encodeChildrenInParallel(children [16]node, buffer *bytes.Buffer) (err error) {
+func encodeChildrenInParallel(children [16]node, buffer io.Writer) (err error) {
 	type result struct {
 		index  int
 		buffer *bytes.Buffer
@@ -242,7 +243,7 @@ func encodeChildrenInParallel(children [16]node, buffer *bytes.Buffer) (err erro
 	return err
 }
 
-func encodeChildrenSequentially(children [16]node, buffer *bytes.Buffer) (err error) {
+func encodeChildrenSequentially(children [16]node, buffer io.Writer) (err error) {
 	for _, child := range children {
 		err = encodeChild(child, buffer)
 		if err != nil {
@@ -252,7 +253,7 @@ func encodeChildrenSequentially(children [16]node, buffer *bytes.Buffer) (err er
 	return nil
 }
 
-func encodeChild(child node, buffer *bytes.Buffer) (err error) {
+func encodeChild(child node, buffer io.Writer) (err error) {
 	if child == nil {
 		return nil
 	}
@@ -272,7 +273,7 @@ func encodeChild(child node, buffer *bytes.Buffer) (err error) {
 
 // encodeLeaf encodes a leaf to the buffer given, with the encoding
 // specified at the top of this package.
-func encodeLeaf(l *leaf, buffer *bytes.Buffer) (err error) {
+func encodeLeaf(l *leaf, buffer io.Writer) (err error) {
 	l.encodingMu.RLock()
 	if !l.dirty && l.encoding != nil {
 		_, err = buffer.Write(l.encoding)

From ba58bd5d40eefcc9356bacb9e152da4a34a71e7a Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Mon, 22 Nov 2021 10:51:30 +0000
Subject: [PATCH 18/19] Use `defer` for read mutex unlocking

---
 lib/trie/hash.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index c4752d4e2e..f798f594b8 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -275,15 +275,14 @@ func encodeChild(child node, buffer io.Writer) (err error) {
 // specified at the top of this package.
 func encodeLeaf(l *leaf, buffer io.Writer) (err error) {
 	l.encodingMu.RLock()
+	defer l.encodingMu.RUnlock()
 	if !l.dirty && l.encoding != nil {
 		_, err = buffer.Write(l.encoding)
-		l.encodingMu.RUnlock()
 		if err != nil {
 			return fmt.Errorf("cannot write stored encoding to buffer: %w", err)
 		}
 		return nil
 	}
-	l.encodingMu.RUnlock()
 
 	encoding, err := l.header()
 	if err != nil {

From 0ae2bf134db6f56cdf9c4bb3108ab1369ad4c05c Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Mon, 22 Nov 2021 12:44:58 +0000
Subject: [PATCH 19/19] Add TODOs

---
 lib/trie/hash.go | 2 ++
 lib/trie/node.go | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/lib/trie/hash.go b/lib/trie/hash.go
index f798f594b8..887ed83731 100644
--- a/lib/trie/hash.go
+++ b/lib/trie/hash.go
@@ -98,6 +98,8 @@ func encodeNode(n node, buffer *bytes.Buffer, parallel bool) (err error) {
 		n.encodingMu.Lock()
 		defer n.encodingMu.Unlock()
 
+		// TODO remove this copying since it defeats the purpose of `buffer`
+		// and the sync.Pool.
 		n.encoding = make([]byte, buffer.Len())
 		copy(n.encoding, buffer.Bytes())
 		return nil
diff --git a/lib/trie/node.go b/lib/trie/node.go
index b7ae35d9db..cc55da89c3 100644
--- a/lib/trie/node.go
+++ b/lib/trie/node.go
@@ -277,6 +277,8 @@ func (l *leaf) encodeAndHash() (encoding, hash []byte, err error) {
 	bufferBytes := buffer.Bytes()
 
 	l.encodingMu.Lock()
+	// TODO remove this copying since it defeats the purpose of `buffer`
+	// and the sync.Pool.
 	l.encoding = make([]byte, len(bufferBytes))
 	copy(l.encoding, bufferBytes)
 	l.encodingMu.Unlock()