From 92ce9a05248d861f5da82b2df36221137c7f67d0 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Tue, 8 Nov 2022 11:58:30 +0000
Subject: [PATCH 1/2] fix(trie): do not get buffer for nil child

---
 internal/trie/node/branch_encode.go      | 39 ++++++++++++++++--------
 internal/trie/node/branch_encode_test.go |  1 -
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/internal/trie/node/branch_encode.go b/internal/trie/node/branch_encode.go
index eebfef7dcc..927f727c0d 100644
--- a/internal/trie/node/branch_encode.go
+++ b/internal/trie/node/branch_encode.go
@@ -50,7 +50,12 @@ func encodeChildrenOpportunisticParallel(children []*Node, buffer io.Writer) (er
 	resultsCh := make(chan encodingAsyncResult, ChildrenCapacity)
 
 	for i, child := range children {
-		if child == nil || child.Kind() == Leaf {
+		if child == nil {
+			resultsCh <- encodingAsyncResult{index: i}
+			continue
+		}
+
+		if child.Kind() == Leaf {
 			runEncodeChild(child, i, resultsCh, nil)
 			continue
 		}
@@ -69,19 +74,30 @@ func encodeChildrenOpportunisticParallel(children []*Node, buffer io.Writer) (er
 	}
 
 	currentIndex := 0
-	resultBuffers := make([]*bytes.Buffer, ChildrenCapacity)
+	indexToBuffer := make(map[int]*bytes.Buffer, ChildrenCapacity)
 	for range children {
 		result := <-resultsCh
 		if result.err != nil && err == nil { // only set the first error we get
 			err = result.err
 		}
 
-		resultBuffers[result.index] = result.buffer
+		indexToBuffer[result.index] = result.buffer
 
 		// write as many completed buffers to the result buffer.
-		for currentIndex < ChildrenCapacity &&
-			resultBuffers[currentIndex] != nil {
-			bufferSlice := resultBuffers[currentIndex].Bytes()
+		for currentIndex < ChildrenCapacity {
+			resultBuffer, done := indexToBuffer[currentIndex]
+			if !done {
+				break
+			}
+
+			nilChildNode := resultBuffer == nil
+			if nilChildNode {
+				delete(indexToBuffer, currentIndex)
+				currentIndex++
+				continue
+			}
+
+			bufferSlice := resultBuffer.Bytes()
 			if err == nil && len(bufferSlice) > 0 {
 				// note buffer.Write copies the byte slice given as argument
 				_, writeErr := buffer.Write(bufferSlice)
@@ -92,8 +108,7 @@ func encodeChildrenOpportunisticParallel(children []*Node, buffer io.Writer) (er
 				}
 			}
 
-			resultBuffers[currentIndex] = nil
-
+			delete(indexToBuffer, currentIndex)
 			currentIndex++
 		}
 	}
@@ -103,6 +118,10 @@ func encodeChildrenOpportunisticParallel(children []*Node, buffer io.Writer) (er
 
 func encodeChildrenSequentially(children []*Node, buffer io.Writer) (err error) {
 	for i, child := range children {
+		if child == nil {
+			continue
+		}
+
 		err = encodeChild(child, buffer)
 		if err != nil {
 			return fmt.Errorf("encoding child at index %d: %w", i, err)
@@ -114,10 +133,6 @@ func encodeChildrenSequentially(children []*Node, buffer io.Writer) (err error)
 // encodeChild computes the Merkle value of the node
 // and then SCALE encodes it to the given buffer.
 func encodeChild(child *Node, buffer io.Writer) (err error) {
-	if child == nil {
-		return nil
-	}
-
 	merkleValue, err := child.CalculateMerkleValue()
 	if err != nil {
 		return fmt.Errorf("computing %s Merkle value: %w", child.Kind(), err)
diff --git a/internal/trie/node/branch_encode_test.go b/internal/trie/node/branch_encode_test.go
index b749225b60..ddb83911b3 100644
--- a/internal/trie/node/branch_encode_test.go
+++ b/internal/trie/node/branch_encode_test.go
@@ -283,7 +283,6 @@ func Test_encodeChild(t *testing.T) {
 		wrappedErr error
 		errMessage string
 	}{
-		"nil node": {},
 		"empty branch child": {
 			child: &Node{
 				Children: make([]*Node, ChildrenCapacity),

From 753189d0a68beb73a449d846d8e3c251c7cd6e8c Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Mon, 14 Nov 2022 16:46:07 +0000
Subject: [PATCH 2/2] Apply Eclesio's suggestion

---
 internal/trie/node/branch_encode.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/internal/trie/node/branch_encode.go b/internal/trie/node/branch_encode.go
index 927f727c0d..75d00cc6d8 100644
--- a/internal/trie/node/branch_encode.go
+++ b/internal/trie/node/branch_encode.go
@@ -90,9 +90,10 @@ func encodeChildrenOpportunisticParallel(children []*Node, buffer io.Writer) (er
 				break
 			}
 
+			delete(indexToBuffer, currentIndex)
+
 			nilChildNode := resultBuffer == nil
 			if nilChildNode {
-				delete(indexToBuffer, currentIndex)
 				currentIndex++
 				continue
 			}
@@ -108,7 +109,6 @@ func encodeChildrenOpportunisticParallel(children []*Node, buffer io.Writer) (er
 				}
 			}
 
-			delete(indexToBuffer, currentIndex)
 			currentIndex++
 		}
 	}