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,67 +272,58 @@
 	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")
 			}
 		}
 	}
 	return nil
 }
 
-// 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, error) {
+// compressNode determines if the node needs compressed, and if so, the len needed to arrive at the next key
+func compressNode(idx int, proofNodes []ProofNode, hashF hashFunc) (int, uint8, error) {
 	parent := &proofNodes[idx]
-	if idx == len(proofNodes)-1 { // The node may have children, but we can only derive their hashes here
-		var isEdge int
+
+	if idx == len(proofNodes)-1 {
 		if parent.Edge != nil {
-			isEdge = 1
-		} else if parent.Binary != nil {
-			isEdge = 0
+			return 1, parent.Len(), nil
 		}
-		return isEdge, parent.Len(), nil
+		return 0, parent.Len(), nil
 	}
 
 	child := &proofNodes[idx+1]
 
-	if parent.Edge != nil && child.Binary != nil {
+	switch {
+	case parent.Edge != nil && child.Binary != nil:
 		return 1, parent.Edge.Path.len, nil
-	}
-
-	if parent.Binary != nil && child.Edge != nil {
+	case parent.Binary != nil && child.Edge != nil:
 		childHash := child.Hash(hashF)
-		if parent.Binary.LeftHash.Equal(childHash) {
-			return 1, child.Edge.Path.len, nil
-		} else if parent.Binary.RightHash.Equal(childHash) {
+		if parent.Binary.LeftHash.Equal(childHash) || 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")
 		}
 	}
 
-	if parent.Binary != nil && child.Binary != nil {
-		return 0, 1, nil
-	}
-
 	return 0, 1, nil
 }
 
 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
@@ -363,39 +354,39 @@
 			}
 		}
 
-		var crntKey *Key
-		crntNode := Node{}
+		var parentKey *Key
+		parentNode := Node{}
 
 		// Set the key of the current node
-		squishedParent, squishParentOffset, err := shouldSquish(i, proofNodes, hashF)
+		compressParent, compressParentOffset, err := compressNode(i, proofNodes, hashF)
 		if err != nil {
 			return nil, err
 		}
-		crntKey, err = getCrntKey(i, squishParentOffset, leafKey, pNode, pathNodes, proofNodes)
+		parentKey, err = getParentKey(i, compressParentOffset, leafKey, pNode, pathNodes, proofNodes)
 		if err != nil {
 			return nil, err
 		}
 
-		// Set the value of the current node
-		crntNode.Value = pNode.Hash(hashF)
-
-		// End of the line, no children
-		if crntKey.len == 251 { //nolint:gomnd
+		// Don't store leafs along proof paths
+		if parentKey.len == 251 { //nolint:gomnd
 			break
 		}
 
+		// Set the value of the current node
+		parentNode.Value = pNode.Hash(hashF)
+
 		// Set the child key of the current node.
-		childKey, err := assignChild(i, squishedParent, &crntNode, &nilKey, leafKey, crntKey, proofNodes, hashF)
+		childKey, err := assignChild(i, compressParent, &parentNode, &nilKey, leafKey, parentKey, proofNodes, hashF)
 		if err != nil {
 			return nil, err
 		}
 
 		// Set the LeftHash and RightHash values
-		crntNode.LeftHash, crntNode.RightHash = getLeftRightHash(i, proofNodes)
+		parentNode.LeftHash, parentNode.RightHash = getLeftRightHash(i, proofNodes)
 
-		pathNodes = append(pathNodes, StorageNode{key: crntKey, node: &crntNode})
+		pathNodes = append(pathNodes, StorageNode{key: parentKey, node: &parentNode})
 
-		// break early
+		// break early since we don't store leafs along proof paths, or if no more nodes exist along the proof paths
 		if childKey.len == 0 || childKey.len == 251 {
 			break
 		}
@@ -434,14 +425,14 @@
 	return leftHash, rightHash
 }
 
-func getCrntKey(idx int, squishParentOffset uint8, leafKey *Key, pNode ProofNode, pathNodes []StorageNode, proofNodes []ProofNode) (*Key, error) {
+func getParentKey(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
 		}
@@ -458,41 +449,27 @@
 }
 
 func getChildKey(childIdx int, crntKey, leafKey, nilKey *Key, proofNodes []ProofNode, hashF hashFunc) (*Key, error) {
-	var squishChildOffset uint8
-	var squishChild int
-	var err error
 	if childIdx > len(proofNodes)-1 {
 		return nilKey, nil
-	} else {
-		squishChild, squishChildOffset, err = shouldSquish(childIdx, proofNodes, hashF)
-		if err != nil {
-			return nil, err
-		}
 	}
-	if crntKey.len+uint8(squishChild)+squishChildOffset == 251 { //nolint:gomnd
-		return nilKey, nil
+
+	compressChild, compressChildOffset, err := compressNode(childIdx, proofNodes, hashF)
+	if err != nil {
+		return nil, err
 	}
-	key, err := leafKey.SubKey(crntKey.len + uint8(squishChild) + squishChildOffset)
-	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
+	if crntKey.len+uint8(compressChild)+compressChildOffset == 251 { //nolint:gomnd
+		return nilKey, nil
 	}
+
+	return leafKey.SubKey(crntKey.len + uint8(compressChild) + compressChildOffset)
 }
 
 // BuildTrie builds a trie using the proof paths (including inner nodes), and then sets all the keys-values (leaves)
 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
@@ -500,40 +477,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