core: optimise (*Billet).Traverse and (*Trie).Find

(*Billet).Traverse changes:
1. Get rid of the `offset` argument. We can cut `from` and pass just the
part that remains. This implies that node with path matching `from` will
also be included in the result, so additional check needs to be added to
the callback function.
2. Pass `path` and `from` without search prefix. Append prefix to the
result inside the callback.
3. Remove duplicating code.

(*Trie).Find changes:
1. Properly prepare `from` argument for traversing function. It closly
depends on the `path` argument.

pkg/core/mpt/billet.go

@@ -200,7 +200,7 @@ func (b *Billet) incrementRefAndStore(h util.Uint256, bs []byte) {
 // returned from `process` function. It also replaces all HashNodes to their
 // "unhashed" counterparts until the stop condition is satisfied.
 func (b *Billet) Traverse(process func(pathToNode []byte, node Node, nodeBytes []byte) bool, ignoreStorageErr bool) error {
-	r, err := b.traverse(b.root, []byte{}, []byte{}, 0, process, ignoreStorageErr)
+	r, err := b.traverse(b.root, []byte{}, []byte{}, process, ignoreStorageErr)
 	if err != nil && !errors.Is(err, errStop) {
 		return err
 	}
@@ -208,7 +208,7 @@ func (b *Billet) Traverse(process func(pathToNode []byte, node Node, nodeBytes [
 	return nil
 }
 
-func (b *Billet) traverse(curr Node, path, from []byte, offset int, process func(pathToNode []byte, node Node, nodeBytes []byte) bool, ignoreStorageErr bool) (Node, error) {
+func (b *Billet) traverse(curr Node, path, from []byte, process func(pathToNode []byte, node Node, nodeBytes []byte) bool, ignoreStorageErr bool) (Node, error) {
 	if _, ok := curr.(EmptyNode); ok {
 		// We're not interested in EmptyNodes, and they do not affect the
 		// traversal process, thus remain them untouched.
@@ -222,9 +222,9 @@ func (b *Billet) traverse(curr Node, path, from []byte, offset int, process func
 			}
 			return nil, err
 		}
-		return b.traverse(r, path, from, offset, process, ignoreStorageErr)
+		return b.traverse(r, path, from, process, ignoreStorageErr)
 	}
-	if _, ok := curr.(*LeafNode); !ok || len(from) <= offset && (len(from) == 0 || !bytes.Equal(path, from)) {
+	if len(from) == 0 {
 		bytes := slice.Copy(curr.Bytes())
 		if process(fromNibbles(path), curr, bytes) {
 			return curr, errStop
@@ -234,33 +234,25 @@ func (b *Billet) traverse(curr Node, path, from []byte, offset int, process func
 	case *LeafNode:
 		return b.tryCollapseLeaf(n), nil
 	case *BranchNode:
-		if len(from) > offset {
+		var (
-			startIndex := from[offset]
+			startIndex byte
-			for i := startIndex; i < lastChild; i++ {
+			endIndex   byte = childrenCount
-				newOffset := len(from)
+		)
-				if i == startIndex {
+		if len(from) != 0 {
-					newOffset = offset + 1
+			endIndex = lastChild
-				}
+			startIndex, from = splitPath(from)
-				r, err := b.traverse(n.Children[i], append(path, i), from, newOffset, process, ignoreStorageErr)
-				if err != nil {
-					if !errors.Is(err, errStop) {
-						return nil, err
-					}
-					n.Children[i] = r
-					return n, err
-				}
-				n.Children[i] = r
-			}
-			return b.tryCollapseBranch(n), nil
 		}
-		for i, child := range n.Children {
+		for i := startIndex; i < endIndex; i++ {
 			var newPath []byte
 			if i == lastChild {
 				newPath = path
 			} else {
-				newPath = append(path, byte(i))
+				newPath = append(path, i)
 			}
-			r, err := b.traverse(child, newPath, from, offset, process, ignoreStorageErr)
+			if i != startIndex {
+				from = []byte{}
+			}
+			r, err := b.traverse(n.Children[i], newPath, from, process, ignoreStorageErr)
 			if err != nil {
 				if !errors.Is(err, errStop) {
 					return nil, err
@@ -272,15 +264,14 @@ func (b *Billet) traverse(curr Node, path, from []byte, offset int, process func
 		}
 		return b.tryCollapseBranch(n), nil
 	case *ExtensionNode:
-		var newOffset int
+		if len(from) != 0 && bytes.HasPrefix(from, n.key) {
-		if len(from) > offset && bytes.HasPrefix(from[offset:], n.key) {
+			from = from[len(n.key):]
-			newOffset = offset + len(n.key)
+		} else if len(from) == 0 || bytes.Compare(n.key, from) > 0 {
-		} else if len(from) <= offset || bytes.Compare(n.key, from[offset:]) > 0 {
+			from = []byte{}
-			newOffset = len(from)
 		} else {
 			return b.tryCollapseExtension(n), nil
 		}
-		r, err := b.traverse(n.next, append(path, n.key...), from, newOffset, process, ignoreStorageErr)
+		r, err := b.traverse(n.next, append(path, n.key...), from, process, ignoreStorageErr)
 		if err != nil && !errors.Is(err, errStop) {
 			return nil, err
 		}

pkg/core/mpt/trie.go

@@ -541,12 +541,29 @@ func (t *Trie) Find(prefix, from []byte, max int) ([]storage.KeyValue, error) {
 		if !bytes.HasPrefix(from, prefix) {
 			return nil, errors.New("`from` argument doesn't match specified prefix")
 		}
-		fromP = toNibbles(from)
+		fromP = toNibbles(from)[len(prefixP):]
 	}
 	_, start, path, err := t.getWithPath(t.root, prefixP, false)
 	if err != nil {
 		return nil, fmt.Errorf("failed to determine the start node: %w", err)
 	}
+	path = path[len(prefixP):]
+
+	if len(fromP) > 0 {
+		if len(path) <= len(fromP) && bytes.HasPrefix(fromP, path) {
+			fromP = fromP[len(path):]
+		} else if len(path) > len(fromP) && bytes.HasPrefix(path, fromP) {
+			fromP = []byte{}
+		} else {
+			cmp := bytes.Compare(path, fromP)
+			switch {
+			case cmp < 0:
+				return []storage.KeyValue{}, nil
+			case cmp > 0:
+				fromP = []byte{}
+			}
+		}
+	}
 
 	var (
 		res   []storage.KeyValue
@@ -555,15 +572,18 @@ func (t *Trie) Find(prefix, from []byte, max int) ([]storage.KeyValue, error) {
 	b := NewBillet(t.root.Hash(), false, t.Store)
 	process := func(pathToNode []byte, node Node, _ []byte) bool {
 		if leaf, ok := node.(*LeafNode); ok {
-			res = append(res, storage.KeyValue{
+			key := append(prefix, pathToNode...)
-				Key:   pathToNode,
+			if !bytes.Equal(key, from) { // (*Billet).traverse includes `from` path into result if so. Need to filter out manually.
-				Value: slice.Copy(leaf.value),
+				res = append(res, storage.KeyValue{
-			})
+					Key:   key,
-			count++
+					Value: slice.Copy(leaf.value),
+				})
+				count++
+			}
 		}
 		return count >= max
 	}
-	_, err = b.traverse(start, path, fromP, len(path), process, false)
+	_, err = b.traverse(start, path, fromP, process, false)
 	if err != nil && !errors.Is(err, errStop) {
 		return nil, err
 	}