store the hashes of children in ProofToPath

core/trie/node.go

@@ -9,9 +9,11 @@ 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
+	LeftHash  *felt.Felt // To verify proofs, we need to store the hash of children
+	RightHash *felt.Felt // even when we can't derive their key
 }
 
 // Hash calculates the hash of a [Node]

core/trie/proof.go

@@ -291,7 +291,7 @@ func ensureMonotonicIncreasing(proofKeys [2]*Key, keys []*felt.Felt) error {
 }
 
 // shouldSquish determines if the node needs compressed, and if so, the len needed to arrive at the next key
-func shouldSquish(idx int, proofNodes []ProofNode) (int, uint8) {
+func shouldSquish(idx int, proofNodes []ProofNode, hashF hashFunc) (int, uint8, *felt.Felt, *felt.Felt) {
 	parent := &proofNodes[idx]
 	var child *ProofNode
 	// The child is nil of the current node is a leaf
@@ -300,18 +300,30 @@ func shouldSquish(idx int, proofNodes []ProofNode) (int, uint8) {
 	}
 
 	if child == nil {
-		return 0, 0
+		return 0, 0, nil, nil
 	}
 
 	if parent.Edge != nil && child.Binary != nil {
-		return 1, parent.Edge.Path.len
+		return 1, parent.Edge.Path.len, child.Binary.LeftHash, child.Binary.RightHash
 	}
 
 	if parent.Binary != nil && child.Edge != nil {
-		return 1, child.Edge.Path.len
+		childHash := child.Hash(hashF)
+		if parent.Binary.LeftHash.Equal(childHash) {
+			return 1, child.Edge.Path.len, child.Edge.Value, parent.Binary.RightHash
+		} else if parent.Binary.RightHash.Equal(childHash) {
+			return 1, child.Edge.Path.len, parent.Binary.LeftHash, child.Edge.Value
+		} else {
+			panic("can't determine the child hash from the parent and child") // Todo: pass error up
+		}
+
+	}
+
+	if parent.Binary != nil && child.Binary != nil {
+		return 0, 0, parent.Binary.LeftHash, parent.Binary.RightHash
 	}
 
-	return 0, 0
+	return 0, 0, nil, nil
 }
 
 func assignChild(crntNode *Node, nilKey, childKey *Key, isRight bool) {
@@ -342,7 +354,7 @@ func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF hashFunc) ([]Storag
 
 		// Set the key of the current node
 		var err error
-		squishedParent, squishParentOffset := shouldSquish(i, proofNodes)
+		squishedParent, squishParentOffset, leftHash, rightHash := shouldSquish(i, proofNodes, hashF)
 		if proofNodes[i].Binary != nil {
 			crntKey, err = leafKey.SubKey(height)
 		} else {
@@ -353,20 +365,23 @@ func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF hashFunc) ([]Storag
 		}
 		offset += squishedParent
 
-		// Set the value of the current node
+		// Set the value,left hash, and right hash of the current node
 		crntNode.Value = proofNodes[i].Hash(hashF)
+		crntNode.LeftHash = leftHash
+		crntNode.RightHash = rightHash
 
-		// Set the children of the current node
+		// Set the child key of the current node.
 		childIdx := i + squishedParent + 1
 		childIsRight := leafKey.Test(leafKey.len - crntKey.len - 1)
-		if i+2+squishedParent < len(proofNodes)-1 { // The child will be compressed, so point to its compressed form
-			squishedChild, squishChildOffset := shouldSquish(childIdx, proofNodes)
+		// There are two+ inner nodes between (compressed) current node and leaf. May be able to compress the child.
+		if i+2+squishedParent < len(proofNodes)-1 {
+			squishedChild, squishChildOffset, _, _ := shouldSquish(childIdx, proofNodes, hashF)
 			childKey, err := leafKey.SubKey(height + squishParentOffset + squishChildOffset + uint8(squishedChild))
 			if err != nil {
 				return nil, err
 			}
 			assignChild(&crntNode, &nilKey, childKey, childIsRight)
-		} else if i+1+offset == len(proofNodes)-1 { // The child points to a leaf, keep it as is
+		} else if i+1+offset == len(proofNodes)-1 { // The child points to a leaf
 			if proofNodes[childIdx].Edge != nil {
 				assignChild(&crntNode, &nilKey, leafKey, childIsRight)
 			} else {
@@ -376,7 +391,7 @@ func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF hashFunc) ([]Storag
 				}
 				assignChild(&crntNode, &nilKey, childKey, childIsRight)
 			}
-		} else { // Current node points directly to leaf
+		} else { // The child is a leaf
 			if proofNodes[i].Edge != nil && len(pathNodes) > 0 {
 				break
 			}

core/trie/proof_test.go

@@ -1,6 +1,7 @@
 package trie_test
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/NethermindEth/juno/core/crypto"
@@ -200,7 +201,7 @@ func build4KeyTrie(t *testing.T) *trie.Trie {
 	//			Juno - should be able to reconstruct this from proofs
 	//			248
 	// 			/  \
-	// 		249
+	// 		249			// Note we cant derive the right key, but need to store it's hash
 	//		/ \
 	//	 250   \
 	//   / \   /\
@@ -211,9 +212,9 @@ func build4KeyTrie(t *testing.T) *trie.Trie {
 	// 			  |
 	// 			 248	Binary
 	//			 / \
-	//		   249	250		Binary Edge		??
+	//		   249	\		Binary Edge		??
 	//	   	   / \	 \
-	//		250  250  \		Binary Edge		??
+	//		250  250  250		Binary Edge		??
 	//	    / \   \    \
 	// 	   0   1   2    4
 
@@ -249,6 +250,7 @@ func build4KeyTrie(t *testing.T) *trie.Trie {
 
 	return tempTrie
 }
+
 func TestGetProof(t *testing.T) {
 	t.Run("Simple Trie - simple binary", func(t *testing.T) {
 		tempTrie := buildSimpleTrie(t)
@@ -863,6 +865,12 @@ func TestVerifyRangeProof(t *testing.T) {
 		require.NoError(t, err)
 		rightProof, err := trie.GetProof(proofKeys[1], tri) // Looks correct in terms of binary and edges
 		require.NoError(t, err)
+
+		leftProofPath, err := trie.ProofToPath(leftProof, proofKeys[0], crypto.Pedersen) // todo: missing 250 node
+		require.NoError(t, err)
+		rightProofPath, err := trie.ProofToPath(rightProof, proofKeys[1], crypto.Pedersen) // todo: last node should not have a value?
+		require.NoError(t, err)
+		fmt.Println(leftProofPath, rightProofPath)
 		proofs := [2][]trie.ProofNode{leftProof, rightProof}
 		rootCommitment, err := tri.Root()
 		require.NoError(t, err)

core/trie/trie.go

@@ -255,7 +255,7 @@ func (t *Trie) replaceLinkWithNewParent(key *Key, commonKey Key, siblingParent S
 	}
 }
 
-func (t *Trie) insertOrUpdateValue(nodeKey *Key, node *Node, nodes []StorageNode, sibling StorageNode, isProof bool) error {
+func (t *Trie) insertOrUpdateValue(nodeKey *Key, node *Node, nodes []StorageNode, sibling StorageNode, siblingIsProof bool) error {
 	commonKey, _ := findCommonKey(nodeKey, sibling.key)
 
 	newParent := &Node{}
@@ -265,28 +265,41 @@ 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
 	}
 
 	leftPath := path(newParent.Left, &commonKey)
 	rightPath := path(newParent.Right, &commonKey)
 
-	newParent.Value = t.hash(leftChild.Hash(&leftPath, t.hash), rightChild.Hash(&rightPath, t.hash))
+	// It is not possible to derive both child keys from the proofs, but we have both
+	// child hashes, so use them in place of the missing key when computing this nodes value
+	if siblingIsProof {
+		if node.LeftHash != nil {
+			newParent.Value = t.hash(node.LeftHash, rightChild.Hash(&rightPath, t.hash))
+		} else if node.RightHash != nil {
+			newParent.Value = t.hash(leftChild.Hash(&leftPath, t.hash), node.RightHash)
+		} else {
+			return errors.New("proof node has neither right/left child hash")
+		}
+	} else {
+		newParent.Value = t.hash(leftChild.Hash(&leftPath, t.hash), rightChild.Hash(&rightPath, t.hash))
+	}
+
 	if err := t.storage.Put(&commonKey, newParent); err != nil {
 		return err
 	}
 
 	// Don't modify the structure outlined by the proof paths
-	if len(nodes) > 1 && !isProof { // sibling has a parent
+	if len(nodes) > 1 && !siblingIsProof { // sibling has a parent
 		siblingParent := (nodes)[len(nodes)-2]
 
 		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 if !isProof {
+	} else if !siblingIsProof {
 		t.setRootKey(&commonKey)
 	} else {
 		t.dirtyNodes = append(t.dirtyNodes, &commonKey)
@@ -393,11 +406,11 @@ func (t *Trie) PutWithProof(key, value *felt.Felt, lProofPath, rProofPath []Stor
 			return nil, nil // no-op
 		}
 
-		IsProof, found := false, false
+		siblingIsProof, found := false, false
 		for i, proof := range lProofPath {
 			if proof.key.Equal(sibling.key) {
 				sibling = lProofPath[i+1]
-				IsProof = true
+				siblingIsProof = true
 				found = true
 				break
 			}
@@ -406,13 +419,13 @@ func (t *Trie) PutWithProof(key, value *felt.Felt, lProofPath, rProofPath []Stor
 			for i, proof := range rProofPath {
 				if proof.key.Equal(sibling.key) {
 					sibling = rProofPath[i+1]
-					IsProof = true
+					siblingIsProof = true
 					break
 				}
 			}
 		}
 
-		err := t.insertOrUpdateValue(&nodeKey, node, nodes, sibling, IsProof)
+		err := t.insertOrUpdateValue(&nodeKey, node, nodes, sibling, siblingIsProof)
 		if err != nil {
 			return nil, err
 		}