From c5a394b659f46d0d37447c75cfe249da738d76d1 Mon Sep 17 00:00:00 2001
From: John Adler <john.adler@protonmail.ch>
Date: Tue, 2 Mar 2021 14:59:46 -0500
Subject: [PATCH] Fix common prefix handling default nodes.

---
 smt.go      | 18 ++++++++----------
 smt_test.go | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/smt.go b/smt.go
index c3b14e5..04c57ea 100644
--- a/smt.go
+++ b/smt.go
@@ -286,8 +286,12 @@ func (smt *SparseMerkleTree) updateWithSideNodes(path []byte, value []byte, side
 	for i := 0; i < smt.depth(); i++ {
 		sideNode := make([]byte, smt.th.pathSize())
 
-		if i >= len(sideNodes) || sideNodes[i] == nil {
-			if commonPrefixCount != smt.depth() && commonPrefixCount > i {
+		// The offset from the bottom of the tree to the start of the side nodes
+		// i-offsetOfSideNodes is the index into sideNodes[]
+		offsetOfSideNodes := smt.depth() - len(sideNodes)
+
+		if i-offsetOfSideNodes < 0 || sideNodes[i-offsetOfSideNodes] == nil {
+			if commonPrefixCount != smt.depth() && commonPrefixCount > smt.depth()-1-i {
 				// If there are no sidenodes at this height, but the number of
 				// bits that the paths of the two leaf nodes share in common is
 				// greater than this height, then we need to build up the tree
@@ -297,16 +301,10 @@ func (smt *SparseMerkleTree) updateWithSideNodes(path []byte, value []byte, side
 				continue
 			}
 		} else {
-			copy(sideNode, sideNodes[i])
+			copy(sideNode, sideNodes[i-offsetOfSideNodes])
 		}
 
-		var index int
-		if commonPrefixCount == smt.depth() {
-			index = len(sideNodes) - 1 - i
-		} else {
-			index = commonPrefixCount - 1 - i
-		}
-		if getBitAtFromMSB(path, index) == right {
+		if getBitAtFromMSB(path, smt.depth()-1-i) == right {
 			currentHash, currentData = smt.th.digestNode(sideNode, currentData)
 		} else {
 			currentHash, currentData = smt.th.digestNode(currentData, sideNode)
diff --git a/smt_test.go b/smt_test.go
index 6faefda..decba87 100644
--- a/smt_test.go
+++ b/smt_test.go
@@ -189,6 +189,9 @@ func TestSparseMerkleTreeKnown(t *testing.T) {
 	key4 := make([]byte, h.Size()+4)
 	copy(key4, baseKey)
 	key4[4] = byte(0b11000000)
+	key5 := make([]byte, h.Size()+4)
+	copy(key5, baseKey)
+	key5[4] = byte(0b11010000)
 
 	_, err = smt.Update(key1, []byte("testValue1"))
 	if err != nil {
@@ -206,6 +209,10 @@ func TestSparseMerkleTreeKnown(t *testing.T) {
 	if err != nil {
 		t.Errorf("returned error when updating empty key: %v", err)
 	}
+	_, err = smt.Update(key5, []byte("testValue5"))
+	if err != nil {
+		t.Errorf("returned error when updating empty key: %v", err)
+	}
 
 	value, err = smt.Get(key1)
 	if err != nil {
@@ -235,6 +242,40 @@ func TestSparseMerkleTreeKnown(t *testing.T) {
 	if !bytes.Equal([]byte("testValue4"), value) {
 		t.Error("did not get correct value when getting non-empty key")
 	}
+	value, err = smt.Get(key5)
+	if err != nil {
+		t.Errorf("returned error when getting non-empty key: %v", err)
+	}
+	if !bytes.Equal([]byte("testValue5"), value) {
+		t.Error("did not get correct value when getting non-empty key")
+	}
+
+	proof1, _ := smt.Prove(key1)
+	proof2, _ := smt.Prove(key2)
+	proof3, _ := smt.Prove(key3)
+	proof4, _ := smt.Prove(key4)
+	proof5, _ := smt.Prove(key5)
+	dsmst := NewDeepSparseMerkleSubTree(NewSimpleMap(), h, smt.Root())
+	err = dsmst.AddBranch(proof1, key1, []byte("testValue1"))
+	if err != nil {
+		t.Errorf("returned error when adding branch to deep subtree: %v", err)
+	}
+	err = dsmst.AddBranch(proof2, key2, []byte("testValue2"))
+	if err != nil {
+		t.Errorf("returned error when adding branch to deep subtree: %v", err)
+	}
+	err = dsmst.AddBranch(proof3, key3, []byte("testValue3"))
+	if err != nil {
+		t.Errorf("returned error when adding branch to deep subtree: %v", err)
+	}
+	err = dsmst.AddBranch(proof4, key4, []byte("testValue4"))
+	if err != nil {
+		t.Errorf("returned error when adding branch to deep subtree: %v", err)
+	}
+	err = dsmst.AddBranch(proof5, key5, []byte("testValue5"))
+	if err != nil {
+		t.Errorf("returned error when adding branch to deep subtree: %v", err)
+	}
 }
 
 // Test tree operations when two leafs are immediate neighbours.