Finished is needed internally

core/trie/snap_support.go

@@ -1,11 +1,11 @@
 package trie
 
 import (
 	"fmt"
 	"github.com/NethermindEth/juno/core/felt"
 )
 
-func (t *Trie) IterateAndGenerateProof(startValue *felt.Felt, consumer func(key, value *felt.Felt) (bool, error)) ([]ProofNode, error) {
+func (t *Trie) IterateAndGenerateProof(startValue *felt.Felt, consumer func(key, value *felt.Felt) (bool, error)) ([]ProofNode, bool, error) {
 	var lastKey *felt.Felt
 
 	finished, err := t.Iterate(startValue, func(key, value *felt.Felt) (bool, error) {
@@ -14,7 +14,7 @@
 		return consumer(key, value)
 	})
 	if err != nil {
-		return nil, err
+		return nil, false, err
 	}
 
 	proofset := map[felt.Felt]ProofNode{}
@@ -27,21 +27,21 @@
 		// actually check that the server did not skip leafs.
 		leftProof, err := GetProof(&startKey, t)
 		if err != nil {
-			return nil, err
+			return nil, false, err
 		}
 		for _, proof := range leftProof {
 			// Well.. using the trie hash here is kinda slow... but I just need it to work right now.
 			proofset[*proof.Hash(t.hash)] = proof
 		}
 
 	}
 
 	if !finished && lastKey != nil {
 		feltBts := lastKey.Bytes()
 		lastKey := NewKey(t.height, feltBts[:])
 		rightProof, err := GetProof(&lastKey, t)
 		if err != nil {
-			return nil, err
+			return nil, false, err
 		}
 
 		for _, proof := range rightProof {
@@ -54,10 +54,10 @@
 		proofs = append(proofs, node)
 	}
 
-	return proofs, nil
+	return proofs, finished, nil
 }
 
 func VerifyRange(root, startKey *felt.Felt, keys, values []*felt.Felt, proofs []ProofNode, hash hashFunc) (hasMore bool, valid bool, err error) {
 	proofMap := map[felt.Felt]ProofNode{}
 	for _, proof := range proofs {
 		proofHash := proof.Hash(hash)
@@ -68,21 +68,21 @@
 		// Special case where the whole trie is sent in one go.
 		// We just need to completely reconstruct the trie.
 
 		tempTrie, err := newTrie(newMemStorage(), 251, hash) //nolint:gomnd
 		if err != nil {
 			return false, false, err
 		}
 
 		for i, key := range keys {
 			_, err := tempTrie.Put(key, values[i])
 			if err != nil {
 				return false, false, err
 			}
 		}
 
 		recalculatedRoot, err := tempTrie.Root()
 		if err != nil {
 			return false, false, err
 		}
 
 		if !root.Equal(recalculatedRoot) {
@@ -101,11 +101,11 @@
 	proofKeys := map[felt.Felt]Key{}
 	err = buildKeys(NewKey(0, []byte{}), root, proofMap, proofKeys, 0)
 	if err != nil {
 		return false, false, err
 	}
 
 	// No idea how this work
 	/*
 		proofValuesArray := []*felt.Felt{}
 		proofKeysArray := []*Key{}
 		for f, key := range proofKeys {
@@ -122,7 +122,7 @@
 	}
 
 	feltBytes := hasMoreKeyCheck.Bytes()
 	hasMoreKeyCheckKey := NewKey(251, feltBytes[:])
 
 	// does this actually work on all case?
 	hasMore = false
@@ -148,7 +148,7 @@
 		ch := proofNode.Edge.Child
 		err := buildKeys(chKey, ch, proofMap, keys, depth+1)
 		if err != nil {
 			return err
 		}
 	} else {
 		binary := proofNode.Binary
@@ -157,14 +157,14 @@
 		ch := binary.LeftHash
 		err := buildKeys(chKey, ch, proofMap, keys, depth+1)
 		if err != nil {
 			return err
 		}
 
 		chKey = currentKey.AppendBit(true)
 		ch = binary.RightHash
 		err = buildKeys(chKey, ch, proofMap, keys, depth+1)
 		if err != nil {
 			return err
 		}
 	}
 

core/trie/snap_support_test.go

@@ -106,7 +106,7 @@ func TestRangeAndVerify(t *testing.T) {
 			var keys []*felt.Felt
 			var values []*felt.Felt
 
-			proofs, err := testTrie.IterateAndGenerateProof(startQuery, func(key, value *felt.Felt) (bool, error) {
+			proofs, _, err := testTrie.IterateAndGenerateProof(startQuery, func(key, value *felt.Felt) (bool, error) {
 				keys = append(keys, key)
 				values = append(values, value)
 				if scenario.maxNode > 0 && len(keys) >= scenario.maxNode {
@@ -255,7 +255,7 @@ func TestRangeAndVerifyReject(t *testing.T) {
 			var keys []*felt.Felt
 			var values []*felt.Felt
 
-			proofs, err := testTrie.IterateAndGenerateProof(startQuery, func(key, value *felt.Felt) (bool, error) {
+			proofs, _, err := testTrie.IterateAndGenerateProof(startQuery, func(key, value *felt.Felt) (bool, error) {
 				keys = append(keys, key)
 				values = append(values, value)
 				if scenario.maxNode > 0 && len(keys) >= scenario.maxNode {