ProoftoPath update to handle unset proof key

core/trie/proof.go

@@ -242,7 +242,7 @@ func VerifyRangeProof(root *felt.Felt, keys, values []*felt.Felt, proofKeys [2]*
 				return false, fmt.Errorf("invalid proof for key %x", proofKeys[i].String())
 			}
 
-			proofPaths[i], err = ProofToPath(proofs[i], proofKeys[i], hash) // Todo: incorrect for both paths :(
+			proofPaths[i], err = ProofToPath(proofs[i], proofKeys[i], proofValues[i], hash) // Todo: incorrect for both paths :(
 			if err != nil {
 				return false, err
 			}
@@ -292,6 +292,9 @@ 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, hashF hashFunc) (int, uint8, *felt.Felt, *felt.Felt) {
+	if idx > len(proofNodes)-1 {
+		return 0, 1, nil, nil
+	}
 	parent := &proofNodes[idx]
 	var child *ProofNode
 	// The child is nil of the current node is a leaf
@@ -300,7 +303,7 @@ func shouldSquish(idx int, proofNodes []ProofNode, hashF hashFunc) (int, uint8,
 	}
 
 	if child == nil {
-		return 0, 0, nil, nil
+		return 0, 1, nil, nil
 	}
 
 	if parent.Edge != nil && child.Binary != nil {
@@ -320,10 +323,10 @@ func shouldSquish(idx int, proofNodes []ProofNode, hashF hashFunc) (int, uint8,
 	}
 
 	if parent.Binary != nil && child.Binary != nil {
-		return 0, 0, parent.Binary.LeftHash, parent.Binary.RightHash
+		return 0, 1, parent.Binary.LeftHash, parent.Binary.RightHash
 	}
 
-	return 0, 0, nil, nil
+	return 0, 1, nil, nil
 }
 
 func assignChild(crntNode *Node, nilKey, childKey *Key, isRight bool) {
@@ -338,15 +341,29 @@ func assignChild(crntNode *Node, nilKey, childKey *Key, isRight bool) {
 
 // ProofToPath returns the set of storage nodes along the proofNodes towards the leaf.
 // Note that only the nodes and children along this path will be set correctly.
-func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF hashFunc) ([]StorageNode, error) {
+func ProofToPath(proofNodes []ProofNode, leafKey *Key, leafValue *felt.Felt, hashF hashFunc) ([]StorageNode, error) {
 	pathNodes := []StorageNode{}
 
 	// Hack: this allows us to store a right without an existing left node.
 	zeroFeltBytes := new(felt.Felt).Bytes()
 	nilKey := NewKey(0, zeroFeltBytes[:])
 
-	i, offset := 0, 0
-	for i <= len(proofNodes)-1 {
+	for i, pNode := range proofNodes {
+		// Skip this node if it has been squished
+		if i != 0 {
+			lastNode := pathNodes[len(pathNodes)-1].node
+			noLeftMatch, noRightMatch := false, false
+			if lastNode.LeftHash != nil && !pNode.Hash(hashF).Equal(lastNode.LeftHash) {
+				noLeftMatch = true
+			}
+			if lastNode.RightHash != nil && !pNode.Hash(hashF).Equal(lastNode.RightHash) {
+				noRightMatch = true
+			}
+			if noLeftMatch && noRightMatch {
+				continue
+			}
+		}
+
 		var crntKey *Key
 		crntNode := Node{}
 
@@ -355,51 +372,42 @@ func ProofToPath(proofNodes []ProofNode, leafKey *Key, hashF hashFunc) ([]Storag
 		// Set the key of the current node
 		var err error
 		squishedParent, squishParentOffset, leftHash, rightHash := shouldSquish(i, proofNodes, hashF)
-		if proofNodes[i].Binary != nil {
+		if pNode.Binary != nil {
 			crntKey, err = leafKey.SubKey(height)
 		} else {
 			crntKey, err = leafKey.SubKey(height + squishParentOffset)
 		}
 		if err != nil {
 			return nil, err
 		}
-		offset += squishedParent
 
 		// Set the value,left hash, and right hash of the current node
-		crntNode.Value = proofNodes[i].Hash(hashF)
+		crntNode.Value = pNode.Hash(hashF)
 		crntNode.LeftHash = leftHash
 		crntNode.RightHash = rightHash
 
+		// If current key is a leaf then don't set the children
+		if crntKey.len == 251 {
+			break
+		}
+
 		// Set the child key of the current node.
 		childIdx := i + squishedParent + 1
-		childIsRight := leafKey.Test(leafKey.len - crntKey.len - 1)
-		// 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
-			if proofNodes[childIdx].Edge != nil {
-				assignChild(&crntNode, &nilKey, leafKey, childIsRight)
-			} else {
-				childKey, err := leafKey.SubKey(crntKey.len + proofNodes[childIdx].Len())
-				if err != nil {
-					return nil, err
-				}
-				assignChild(&crntNode, &nilKey, childKey, childIsRight)
-			}
-		} else { // The child is a leaf
-			if proofNodes[i].Edge != nil && len(pathNodes) > 0 {
-				break
-			}
-			assignChild(&crntNode, &nilKey, leafKey, childIsRight)
+		_, squishChildOffset, _, _ := shouldSquish(childIdx, proofNodes, hashF)
+		childKey, err := leafKey.SubKey(crntKey.len + squishChildOffset)
+		if err != nil {
+			return nil, err
 		}
+		childIsRight := leafKey.Test(leafKey.len - crntKey.len - 1)
+		assignChild(&crntNode, &nilKey, childKey, childIsRight)
 
 		pathNodes = append(pathNodes, StorageNode{key: crntKey, node: &crntNode})
-		i += 1 + offset
+	}
+	// If the proof node is unset, then ensure that the child keys are also
+	// unset (they are erroneously set above)
+	if leafValue == nil {
+		pathNodes[len(pathNodes)-1].node.Left = nil
+		pathNodes[len(pathNodes)-1].node.Right = nil
 	}
 	pathNodes = addLeafNode(proofNodes, pathNodes, leafKey)
 	return pathNodes, nil

core/trie/proof_test.go

@@ -204,7 +204,7 @@ func build4KeyTrie(t *testing.T) *trie.Trie {
 	// 		249			// Note we cant derive the right key, but need to store it's hash
 	//		/ \
 	//	 250   \
-	//   / \   /\
+	//   / \   / (Left hash set, no key)
 	//  0
 
 	//			Pathfinder (???)
@@ -592,7 +592,7 @@ func TestProofToPath(t *testing.T) {
 		}
 		zeroFeltBytes := new(felt.Felt).SetUint64(0).Bytes()
 		leafkey := trie.NewKey(251, zeroFeltBytes[:])
-		sns, err := trie.ProofToPath(proofNodes, &leafkey, crypto.Pedersen) // Todo : we should be able to set the leaf as well
+		sns, err := trie.ProofToPath(proofNodes, &leafkey, new(felt.Felt).SetUint64(2), crypto.Pedersen) // Todo : we should be able to set the leaf as well
 		require.NoError(t, err)
 
 		rootKey := tempTrie.RootKey()
@@ -628,7 +628,7 @@ func TestProofToPath(t *testing.T) {
 
 		zeroFeltBytes := new(felt.Felt).SetUint64(0).Bytes()
 		leafkey := trie.NewKey(251, zeroFeltBytes[:])
-		sns, err := trie.ProofToPath(proofNodes, &leafkey, crypto.Pedersen)
+		sns, err := trie.ProofToPath(proofNodes, &leafkey, utils.HexToFelt(t, "0xcc"), crypto.Pedersen)
 		require.NoError(t, err)
 
 		rootKey := tempTrie.RootKey()
@@ -654,7 +654,7 @@ func TestProofToPath(t *testing.T) {
 		bProofs, err := trie.GetBoundaryProofs(&zeroLeafkey, &twoLeafkey, tri)
 		require.NoError(t, err)
 
-		leftProofPath, err := trie.ProofToPath(bProofs[0], &zeroLeafkey, crypto.Pedersen)
+		leftProofPath, err := trie.ProofToPath(bProofs[0], &zeroLeafkey, new(felt.Felt).SetUint64(4), crypto.Pedersen)
 		require.Equal(t, 3, len(leftProofPath))
 		require.NoError(t, err)
 		require.Equal(t, rootKey, leftProofPath[0].Key())
@@ -667,7 +667,7 @@ func TestProofToPath(t *testing.T) {
 		require.Equal(t, leftNode.Left, leftProofPath[1].Node().Left)
 		require.NotEqual(t, leftNode.Right, leftProofPath[0].Node().Right)
 
-		rightProofPath, err := trie.ProofToPath(bProofs[1], &twoLeafkey, crypto.Pedersen)
+		rightProofPath, err := trie.ProofToPath(bProofs[1], &twoLeafkey, new(felt.Felt).SetUint64(6), crypto.Pedersen)
 		require.Equal(t, 2, len(rightProofPath))
 		require.NoError(t, err)
 		require.Equal(t, rootKey, rightProofPath[0].Key())
@@ -711,10 +711,10 @@ func TestBuildTrie(t *testing.T) {
 		bProofs, err := trie.GetBoundaryProofs(&zeroLeafkey, &twoLeafkey, tri)
 		require.NoError(t, err)
 
-		leftProof, err := trie.ProofToPath(bProofs[0], &zeroLeafkey, crypto.Pedersen)
+		leftProof, err := trie.ProofToPath(bProofs[0], &zeroLeafkey, new(felt.Felt).SetUint64(4), crypto.Pedersen)
 		require.NoError(t, err)
 
-		rightProof, err := trie.ProofToPath(bProofs[1], &twoLeafkey, crypto.Pedersen)
+		rightProof, err := trie.ProofToPath(bProofs[1], &twoLeafkey, new(felt.Felt).SetUint64(6), crypto.Pedersen)
 		require.NoError(t, err)
 
 		keys := []*felt.Felt{new(felt.Felt).SetUint64(1)}
@@ -866,9 +866,9 @@ func TestVerifyRangeProof(t *testing.T) {
 		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
+		leftProofPath, err := trie.ProofToPath(leftProof, proofKeys[0], proofValues[0], crypto.Pedersen) // correct
 		require.NoError(t, err)
-		rightProofPath, err := trie.ProofToPath(rightProof, proofKeys[1], crypto.Pedersen) // todo: last node should not have a value?
+		rightProofPath, err := trie.ProofToPath(rightProof, proofKeys[1], proofValues[1], crypto.Pedersen) // gives the correct structure
 		require.NoError(t, err)
 		fmt.Println(leftProofPath, rightProofPath)
 		proofs := [2][]trie.ProofNode{leftProof, rightProof}