tidy logic

core/trie/proof.go

@@ -63,11 +63,11 @@
 	proofs := [2][]ProofNode{}
 	leftProof, err := GetProof(leftBoundary, tri)
 	if err != nil {
 		return proofs, err
 	}
 	rightProof, err := GetProof(rightBoundary, tri)
 	if err != nil {
 		return proofs, err
 	}
 	proofs[0] = leftProof
 	proofs[1] = rightProof
@@ -219,18 +219,18 @@
 ) (bool, error) {
 	// Step 0: checks
 	if len(keys) != len(values) {
 		return false, fmt.Errorf("inconsistent proof data, number of keys: %d, number of values: %d", len(keys), len(values))
 	}
 
 	// Ensure all keys are monotonic increasing
 	if err := ensureMonotonicIncreasing(proofKeys, keys); err != nil {
 		return false, err
 	}
 
 	// Ensure the inner values contain no deletions
 	for _, value := range values {
 		if value.Equal(&felt.Zero) {
 			return false, errors.New("range contains deletion")
 		}
 	}
 
@@ -240,12 +240,12 @@
 	for i := 0; i < 2; i++ {
 		if proofs[i] != nil {
 			if !VerifyProof(root, proofKeys[i], proofValues[i], proofs[i], hash) {
 				return false, fmt.Errorf("invalid proof for key %x", proofKeys[i].String())
 			}
 
 			proofPaths[i], err = ProofToPath(proofs[i], proofKeys[i], hash)
 			if err != nil {
 				return false, err
 			}
 		}
 	}
@@ -253,16 +253,16 @@
 	// Step 2: Build trie from proofPaths and keys
 	tmpTrie, err := BuildTrie(proofPaths[0], proofPaths[1], keys, values)
 	if err != nil {
 		return false, err
 	}
 
 	// Verify that the recomputed root hash matches the provided root hash
 	recomputedRoot, err := tmpTrie.Root()
 	if err != nil {
 		return false, err
 	}
 	if !recomputedRoot.Equal(root) {
 		return false, errors.New("root hash mismatch")
 	}
 
 	return true, nil
@@ -272,19 +272,19 @@
 	if proofKeys[0] != nil {
 		leftProofFelt := proofKeys[0].Felt()
 		if leftProofFelt.Cmp(keys[0]) >= 0 {
 			return errors.New("range is not monotonically increasing")
 		}
 	}
 	if proofKeys[1] != nil {
 		rightProofFelt := proofKeys[1].Felt()
 		if keys[len(keys)-1].Cmp(&rightProofFelt) >= 0 {
 			return errors.New("range is not monotonically increasing")
 		}
 	}
 	if len(keys) >= 2 {
 		for i := 0; i < len(keys)-1; i++ {
 			if keys[i].Cmp(keys[i+1]) >= 0 {
 				return errors.New("range is not monotonically increasing")
 			}
 		}
 	}
@@ -317,7 +317,7 @@
 		} else if parent.Binary.RightHash.Equal(childHash) {
 			return 1, child.Edge.Path.len, nil
 		} else {
 			return 0, 0, errors.New("can't determine the child hash from the parent and child")
 		}
 	}
 
@@ -325,17 +325,23 @@
 		return 0, 1, nil
 	}
 
 	return 0, 1, nil
 }
 
-func assignChild(crntNode *Node, nilKey, childKey *Key, isRight bool) {
-	if isRight {
+func assignChild(i, squishedParent int, crntNode *Node, nilKey, leafKey, crntKey *Key, proofNodes []ProofNode, hashF hashFunc) (*Key, error) {
+	childInd := i + squishedParent + 1
+	childKey, err := getChildKey(childInd, crntKey, leafKey, nilKey, proofNodes, hashF)
+	if err != nil {
+		return nil, err
+	}
+	if leafKey.Test(leafKey.len - crntKey.len - 1) {
 		crntNode.Right = childKey
 		crntNode.Left = nilKey
 	} else {
 		crntNode.Right = nilKey
 		crntNode.Left = childKey
 	}
+	return childKey, nil
 }
 
 // ProofToPath returns a set of storage nodes from the root to the end of the proof path.
@@ -352,32 +358,22 @@
 	for i, pNode := range proofNodes {
 		// Keep moving along the path (may need to skip nodes that were compressed into the last path node)
 		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 {
+			if skipNode(pNode, pathNodes, hashF) {
 				continue
 			}
 		}
 
 		var crntKey *Key
 		crntNode := Node{}
 
-		height := getHeight(i, pathNodes, proofNodes)
-
 		// Set the key of the current node
 		squishedParent, squishParentOffset, err := shouldSquish(i, proofNodes, hashF)
 		if err != nil {
 			return nil, err
 		}
-		crntKey, err = getCrntKey(height, squishParentOffset, leafKey, pNode)
+		crntKey, err = getCrntKey(i, squishParentOffset, leafKey, pNode, pathNodes, proofNodes)
 		if err != nil {
 			return nil, err
 		}
 
 		// Set the value of the current node
@@ -385,17 +381,14 @@
 
 		// End of the line, no children
 		if crntKey.len == 251 { //nolint:gomnd
 			break
 		}
 
 		// Set the child key of the current node.
-		childInd := i + squishedParent + 1
-		childKey, err := getChildKey(childInd, crntKey, leafKey, &nilKey, proofNodes, hashF)
+		childKey, err := assignChild(i, squishedParent, &crntNode, &nilKey, leafKey, crntKey, proofNodes, hashF)
 		if err != nil {
 			return nil, err
 		}
-		childIsRight := leafKey.Test(leafKey.len - crntKey.len - 1)
-		assignChild(&crntNode, &nilKey, childKey, childIsRight)
 
 		// Set the LeftHash and RightHash values
 		crntNode.LeftHash, crntNode.RightHash = getLeftRightHash(i, proofNodes)
@@ -410,6 +403,21 @@
 	return pathNodes, nil
 }
 
+func skipNode(pNode ProofNode, pathNodes []StorageNode, hashF hashFunc) bool {
+	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 {
+		return true
+	}
+	return false
+}
+
 func getLeftRightHash(parentInd int, proofNodes []ProofNode) (*felt.Felt, *felt.Felt) {
 	var leftHash, rightHash *felt.Felt
 	parent := &proofNodes[parentInd]
@@ -426,9 +434,21 @@
 	return leftHash, rightHash
 }
 
-func getCrntKey(height, squishParentOffset uint8, leafKey *Key, pNode ProofNode) (*Key, error) {
+func getCrntKey(idx int, squishParentOffset uint8, leafKey *Key, pNode ProofNode, pathNodes []StorageNode, proofNodes []ProofNode) (*Key, error) {
 	var crntKey *Key
 	var err error
+
+	var height uint8
+	if len(pathNodes) > 0 {
+		if proofNodes[idx].Edge != nil {
+			height = pathNodes[len(pathNodes)-1].key.len + proofNodes[idx].Edge.Path.len
+		} else {
+			height = pathNodes[len(pathNodes)-1].key.len + 1
+		}
+	} else {
+		height = 0
+	}
+
 	if pNode.Binary != nil {
 		crntKey, err = leafKey.SubKey(height)
 	} else {
@@ -446,7 +466,7 @@
 	} else {
 		squishChild, squishChildOffset, err = shouldSquish(childIdx, proofNodes, hashF)
 		if err != nil {
 			return nil, err
 		}
 	}
 	if crntKey.len+uint8(squishChild)+squishChildOffset == 251 { //nolint:gomnd
@@ -456,15 +476,15 @@
 	return key, err
 }
 
 func getHeight(idx int, pathNodes []StorageNode, proofNodes []ProofNode) uint8 {
 	if len(pathNodes) > 0 {
 		if proofNodes[idx].Edge != nil {
 			return pathNodes[len(pathNodes)-1].key.len + proofNodes[idx].Edge.Path.len
 		} else {
 			return pathNodes[len(pathNodes)-1].key.len + 1
 		}
 	} else {
 		return 0
 	}
 }
 
@@ -472,7 +492,7 @@
 func BuildTrie(leftProofPath, rightProofPath []StorageNode, keys, values []*felt.Felt) (*Trie, error) { //nolint:gocyclo
 	tempTrie, err := NewTriePedersen(newMemStorage(), 251) //nolint:gomnd
 	if err != nil {
 		return nil, err
 	}
 
 	// merge proof paths
@@ -480,40 +500,40 @@
 		// Can't store nil keys so stop merging
 		if leftProofPath[i].node.Left == nil || leftProofPath[i].node.Right == nil ||
 			rightProofPath[i].node.Left == nil || rightProofPath[i].node.Right == nil {
 			break
 		}
 		if leftProofPath[i].key.Equal(rightProofPath[i].key) {
 			leftProofPath[i].node.Right = rightProofPath[i].node.Right
 			rightProofPath[i].node.Left = leftProofPath[i].node.Left
 		} else {
 			break
 		}
 	}
 
 	for _, sNode := range leftProofPath {
 		if sNode.node.Left == nil || sNode.node.Right == nil {
 			break
 		}
 		_, err := tempTrie.PutInner(sNode.key, sNode.node)
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	for _, sNode := range rightProofPath {
 		if sNode.node.Left == nil || sNode.node.Right == nil {
 			break
 		}
 		_, err := tempTrie.PutInner(sNode.key, sNode.node)
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	for i := range len(keys) {
 		_, err := tempTrie.PutWithProof(keys[i], values[i], leftProofPath, rightProofPath)
 		if err != nil {
 			return nil, err
 		}
 	}
 	return tempTrie, nil