reconstructed trie has correct structure (but incorrect root)

core/trie/node.go

@@ -9,9 +9,10 @@ import (
 
 // A Node represents a node in the [Trie]
 type Node struct {
-	Value *felt.Felt
-	Left  *Key
-	Right *Key
+	Value   *felt.Felt
+	Left    *Key
+	Right   *Key
+	IsProof bool
 }
 
 // Hash calculates the hash of a [Node]

core/trie/proof.go

@@ -431,13 +431,14 @@ func BuildTrie(leftProof, rightProof []StorageNode, keys []*felt.Felt, values []
 	builtRightNode, err := tempTrie.GetNodeFromKey(builtRootNode.Right) // correct
 
 	for i := range len(keys) {
-		_, err := tempTrie.PutWithProof(keys[i], values[i], leftProof, rightProof) // This is where the issues arises
+		_, err := tempTrie.PutWithProof(keys[i], values[i], leftProof, rightProof) // This is where the issues arises (right gets overwritten)
 		if err != nil {
 			return nil, err
 		}
 	}
+	builtRootKey = tempTrie.RootKey()
 	builtRootNode, err = tempTrie.GetNodeFromKey(builtRootKey)         // correct
-	builtLeftNode, err = tempTrie.GetNodeFromKey(builtRootNode.Left)   // leftPath was getting overwritten... now the right is (parent sibling)
+	builtLeftNode, err = tempTrie.GetNodeFromKey(builtRootNode.Left)   // correct
 	builtRightNode, err = tempTrie.GetNodeFromKey(builtRootNode.Right) // correct
 	fmt.Println(builtLeftNode, builtRightNode)
 

core/trie/proof_test.go

@@ -147,11 +147,10 @@ func build3KeyTrie(t *testing.T) *trie.Trie {
 	_, err = tempTrie.Put(key1, value1)
 	require.NoError(t, err)
 
-	_, err = tempTrie.Put(key2, value2)
-	require.NoError(t, err)
-
 	_, err = tempTrie.Put(key3, value3)
 	require.NoError(t, err)
+	_, err = tempTrie.Put(key2, value2)
+	require.NoError(t, err)
 
 	require.NoError(t, tempTrie.Commit())
 	return tempTrie
@@ -557,15 +556,15 @@ func TestBuildTrie(t *testing.T) {
 		builtRootKey := builtTrie.RootKey()
 		builtRootNode, err := builtTrie.GetNodeFromKey(builtRootKey) // looks correct
 		require.NoError(t, err)
-		builtLeftNode, err := builtTrie.GetNodeFromKey(builtRootNode.Left) // left len gets incorrectly overwritten
+		builtLeftNode, err := builtTrie.GetNodeFromKey(builtRootNode.Left) //  looks correct
 		require.NoError(t, err)
 		builtRightNode, err := builtTrie.GetNodeFromKey(builtRootNode.Right) // looks correct
 		require.NoError(t, err)
 
 		require.Equal(t, rootKey, builtRootKey)
 		require.Equal(t, rootNode.Left, builtRootNode.Left)
 		require.Equal(t, rootNode.Right, builtRootNode.Right)
-		require.Equal(t, rootNode.Value.String(), builtRootNode.Value.String())
+		require.Equal(t, rootNode.Value.String(), builtRootNode.Value.String()) // Incorrect
 		require.Equal(t, leftNode, builtLeftNode)
 		require.Equal(t, rightNode, builtRightNode)
 

core/trie/trie.go

@@ -267,7 +267,7 @@ func (t *Trie) insertOrUpdateValue(nodeKey *Key, node *Node, nodes []StorageNode
 		newParent.Left, newParent.Right = sibling.key, nodeKey
 		leftChild, rightChild = sibling.node, node
 	} else {
-		newParent.Left, newParent.Right = nodeKey, sibling.key
+		newParent.Left, newParent.Right = nodeKey, sibling.key //
 		leftChild, rightChild = node, sibling.node
 	}
 
@@ -279,15 +279,16 @@ func (t *Trie) insertOrUpdateValue(nodeKey *Key, node *Node, nodes []StorageNode
 		return err
 	}
 
-	if len(nodes) > 1 { // sibling has a parent
+	// Don't modify the structure outlined by the proof paths
+	if len(nodes) > 1 && !sibling.node.IsProof { // sibling has a parent
 		siblingParent := (nodes)[len(nodes)-2]
 
-		t.replaceLinkWithNewParent(sibling.key, commonKey, siblingParent)
+		t.replaceLinkWithNewParent(sibling.key, commonKey, siblingParent) // error with overwritting right arises here
 		if err := t.storage.Put(siblingParent.key, siblingParent.node); err != nil {
 			return err
 		}
 		t.dirtyNodes = append(t.dirtyNodes, &commonKey)
-	} else {
+	} else if !sibling.node.IsProof {
 		t.setRootKey(&commonKey)
 	}
 
@@ -396,12 +397,14 @@ func (t *Trie) PutWithProof(key, value *felt.Felt, lProofPath, rProofPath []Stor
 		for i, proof := range lProofPath {
 			if proof.key.Equal(sibling.key) {
 				sibling = lProofPath[i+1]
+				sibling.node.IsProof = true
 				break
 			}
 		}
 		for i, proof := range rProofPath {
 			if proof.key.Equal(sibling.key) {
 				sibling = rProofPath[i+1]
+				sibling.node.IsProof = true
 				break
 			}
 		}