feat: optimize the fast key format and encoding (cosmos#923)

Kava-Labs · Mar 27, 2024 · 8967633 · 8967633
1 parent 2894221
commit 8967633
Show file tree

Hide file tree

Showing 6 changed files with 152 additions and 87 deletions.
diff --git a/internal/encoding/encoding.go b/internal/encoding/encoding.go
@@ -30,6 +30,7 @@ var uvarintPool = &sync.Pool{
 
 // decodeBytes decodes a varint length-prefixed byte slice, returning it along with the number
 // of input bytes read.
+// Assumes bz will not be mutated.
 func DecodeBytes(bz []byte) ([]byte, int, error) {
 	s, n, err := DecodeUvarint(bz)
 	if err != nil {
@@ -51,9 +52,7 @@ func DecodeBytes(bz []byte) ([]byte, int, error) {
 	if len(bz) < end {
 		return nil, n, fmt.Errorf("insufficient bytes decoding []byte of length %v", size)
 	}
-	bz2 := make([]byte, size)
-	copy(bz2, bz[n:end])
-	return bz2, end, nil
+	return bz[n:end], end, nil
 }
 
 // decodeUvarint decodes a varint-encoded unsigned integer from a byte slice, returning it and the
@@ -97,6 +96,23 @@ func EncodeBytes(w io.Writer, bz []byte) error {
 	return err
 }
 
+var hashLenBz []byte
+
+func init() {
+	hashLenBz = make([]byte, 1)
+	binary.PutUvarint(hashLenBz, 32)
+}
+
+// Encode 32 byte long hash
+func Encode32BytesHash(w io.Writer, bz []byte) error {
+	_, err := w.Write(hashLenBz)
+	if err != nil {
+		return err
+	}
+	_, err = w.Write(bz)
+	return err
+}
+
 // encodeBytesSlice length-prefixes the byte slice and returns it.
 func EncodeBytesSlice(bz []byte) ([]byte, error) {
 	buf := bufPool.Get().(*bytes.Buffer)

diff --git a/keyformat/prefix_formatter.go b/keyformat/prefix_formatter.go
@@ -0,0 +1,42 @@
+package keyformat
+
+import "encoding/binary"
+
+// This file builds some dedicated key formatters for what appears in benchmarks.
+
+// Prefixes a single byte before a 32 byte hash.
+type FastPrefixFormatter struct {
+	prefix      byte
+	length      int
+	prefixSlice []byte
+}
+
+func NewFastPrefixFormatter(prefix byte, length int) *FastPrefixFormatter {
+	return &FastPrefixFormatter{prefix: prefix, length: length, prefixSlice: []byte{prefix}}
+}
+
+func (f *FastPrefixFormatter) Key(bz []byte) []byte {
+	key := make([]byte, 1+f.length)
+	key[0] = f.prefix
+	copy(key[1:], bz)
+	return key
+}
+
+func (f *FastPrefixFormatter) Scan(key []byte, a interface{}) {
+	scan(a, key[1:])
+}
+
+func (f *FastPrefixFormatter) KeyInt64(bz int64) []byte {
+	key := make([]byte, 1+f.length)
+	key[0] = f.prefix
+	binary.BigEndian.PutUint64(key[1:], uint64(bz))
+	return key
+}
+
+func (f *FastPrefixFormatter) Prefix() []byte {
+	return f.prefixSlice
+}
+
+func (f *FastPrefixFormatter) Length() int {
+	return 1 + f.length
+}
diff --git a/mutable_tree.go b/mutable_tree.go
@@ -263,64 +263,68 @@ func (tree *MutableTree) set(key []byte, value []byte) (updated bool, err error)
 func (tree *MutableTree) recursiveSet(node *Node, key []byte, value []byte) (
 	newSelf *Node, updated bool, err error,
 ) {
-	version := tree.version + 1
-
 	if node.isLeaf() {
-		if !tree.skipFastStorageUpgrade {
-			tree.addUnsavedAddition(key, fastnode.NewNode(key, value, version))
-		}
-		switch bytes.Compare(key, node.key) {
-		case -1: // setKey < leafKey
-			return &Node{
-				key:           node.key,
-				subtreeHeight: 1,
-				size:          2,
-				nodeKey:       nil,
-				leftNode:      NewNode(key, value),
-				rightNode:     node,
-			}, false, nil
-		case 1: // setKey > leafKey
-			return &Node{
-				key:           key,
-				subtreeHeight: 1,
-				size:          2,
-				nodeKey:       nil,
-				leftNode:      node,
-				rightNode:     NewNode(key, value),
-			}, false, nil
-		default:
-			return NewNode(key, value), true, nil
+		return tree.recursiveSetLeaf(node, key, value)
+	}
+	node, err = node.clone(tree)
+	if err != nil {
+		return nil, false, err
+	}
+
+	if bytes.Compare(key, node.key) < 0 {
+		node.leftNode, updated, err = tree.recursiveSet(node.leftNode, key, value)
+		if err != nil {
+			return nil, updated, err
 		}
 	} else {
-		node, err = node.clone(tree)
+		node.rightNode, updated, err = tree.recursiveSet(node.rightNode, key, value)
 		if err != nil {
-			return nil, false, err
+			return nil, updated, err
 		}
+	}
 
-		if bytes.Compare(key, node.key) < 0 {
-			node.leftNode, updated, err = tree.recursiveSet(node.leftNode, key, value)
-			if err != nil {
-				return nil, updated, err
-			}
-		} else {
-			node.rightNode, updated, err = tree.recursiveSet(node.rightNode, key, value)
-			if err != nil {
-				return nil, updated, err
-			}
-		}
+	if updated {
+		return node, updated, nil
+	}
+	err = node.calcHeightAndSize(tree.ImmutableTree)
+	if err != nil {
+		return nil, false, err
+	}
+	newNode, err := tree.balance(node)
+	if err != nil {
+		return nil, false, err
+	}
+	return newNode, updated, err
+}
 
-		if updated {
-			return node, updated, nil
-		}
-		err = node.calcHeightAndSize(tree.ImmutableTree)
-		if err != nil {
-			return nil, false, err
-		}
-		newNode, err := tree.balance(node)
-		if err != nil {
-			return nil, false, err
-		}
-		return newNode, updated, err
+func (tree *MutableTree) recursiveSetLeaf(node *Node, key []byte, value []byte) (
+	newSelf *Node, updated bool, err error,
+) {
+	version := tree.version + 1
+	if !tree.skipFastStorageUpgrade {
+		tree.addUnsavedAddition(key, fastnode.NewNode(key, value, version))
+	}
+	switch bytes.Compare(key, node.key) {
+	case -1: // setKey < leafKey
+		return &Node{
+			key:           node.key,
+			subtreeHeight: 1,
+			size:          2,
+			nodeKey:       nil,
+			leftNode:      NewNode(key, value),
+			rightNode:     node,
+		}, false, nil
+	case 1: // setKey > leafKey
+		return &Node{
+			key:           key,
+			subtreeHeight: 1,
+			size:          2,
+			nodeKey:       nil,
+			leftNode:      node,
+			rightNode:     NewNode(key, value),
+		}, false, nil
+	default:
+		return NewNode(key, value), true, nil
 	}
 }
 

diff --git a/node.go b/node.go
@@ -57,11 +57,15 @@ func GetRootKey(version int64) []byte {
 
 // Node represents a node in a Tree.
 type Node struct {
-	key           []byte
-	value         []byte
-	hash          []byte
-	nodeKey       *NodeKey
-	leftNodeKey   []byte
+	key     []byte
+	value   []byte
+	hash    []byte
+	nodeKey *NodeKey
+	// Legacy: LeftNodeHash
+	// v1: Left node ptr via Version/key
+	leftNodeKey []byte
+	// Legacy: RightNodeHash
+	// v1: Right node ptr via Version/key
 	rightNodeKey  []byte
 	size          int64
 	leftNode      *Node
@@ -522,19 +526,19 @@ func (node *Node) writeHashBytes(w io.Writer, version int64) error {
 		// (e.g. ProofLeafNode.ValueHash)
 		valueHash := sha256.Sum256(node.value)
 
-		err = encoding.EncodeBytes(w, valueHash[:])
+		err = encoding.Encode32BytesHash(w, valueHash[:])
 		if err != nil {
 			return fmt.Errorf("writing value, %w", err)
 		}
 	} else {
 		if node.leftNode == nil || node.rightNode == nil {
 			return ErrEmptyChild
 		}
-		err = encoding.EncodeBytes(w, node.leftNode.hash)
+		err = encoding.Encode32BytesHash(w, node.leftNode.hash)
 		if err != nil {
 			return fmt.Errorf("writing left hash, %w", err)
 		}
-		err = encoding.EncodeBytes(w, node.rightNode.hash)
+		err = encoding.Encode32BytesHash(w, node.rightNode.hash)
 		if err != nil {
 			return fmt.Errorf("writing right hash, %w", err)
 		}
@@ -600,7 +604,7 @@ func (node *Node) writeBytes(w io.Writer) error {
 			return fmt.Errorf("writing value, %w", err)
 		}
 	} else {
-		err = encoding.EncodeBytes(w, node.hash)
+		err = encoding.Encode32BytesHash(w, node.hash)
 		if err != nil {
 			return fmt.Errorf("writing hash, %w", err)
 		}
@@ -620,7 +624,7 @@ func (node *Node) writeBytes(w io.Writer) error {
 			return fmt.Errorf("writing mode, %w", err)
 		}
 		if mode&ModeLegacyLeftNode != 0 { // legacy leftNodeKey
-			err = encoding.EncodeBytes(w, node.leftNodeKey)
+			err = encoding.Encode32BytesHash(w, node.leftNodeKey)
 			if err != nil {
 				return fmt.Errorf("writing the legacy left node key, %w", err)
 			}
@@ -639,7 +643,7 @@ func (node *Node) writeBytes(w io.Writer) error {
 			return ErrRightNodeKeyEmpty
 		}
 		if mode&ModeLegacyRightNode != 0 { // legacy rightNodeKey
-			err = encoding.EncodeBytes(w, node.rightNodeKey)
+			err = encoding.Encode32BytesHash(w, node.rightNodeKey)
 			if err != nil {
 				return fmt.Errorf("writing the legacy right node key, %w", err)
 			}

diff --git a/node_test.go b/node_test.go
@@ -61,8 +61,8 @@ func TestNode_encode_decode(t *testing.T) {
 			},
 			leftNodeKey:  childNodeKey.GetKey(),
 			rightNodeKey: childNodeKey.GetKey(),
-			hash:         []byte{0x70, 0x80, 0x90, 0xa0},
-		}, "060e036b657904708090a00002020202", false},
+			hash:         []byte{0x7f, 0x68, 0x90, 0xca, 0x16, 0xde, 0xa6, 0xe8, 0x89, 0x3d, 0x96, 0xf0, 0xa3, 0xd, 0xa, 0x14, 0xe5, 0x55, 0x59, 0xfc, 0x9b, 0x83, 0x4, 0x91, 0xe3, 0xd2, 0x45, 0x1c, 0x81, 0xf6, 0xd1, 0xe},
+		}, "060e036b6579207f6890ca16dea6e8893d96f0a30d0a14e55559fc9b830491e3d2451c81f6d10e0002020202", false},
 		"inner hybrid": {&Node{
 			subtreeHeight: 3,
 			size:          7,
@@ -73,8 +73,8 @@ func TestNode_encode_decode(t *testing.T) {
 			},
 			leftNodeKey:  childNodeKey.GetKey(),
 			rightNodeKey: childNodeHash,
-			hash:         []byte{0x70, 0x80, 0x90, 0xa0},
-		}, "060e036b657904708090a0040202207f6890ca16dea6e8893d96f0a30d0a14e55559fc9b830491e3d2451c81f6d10e", false},
+			hash:         []byte{0x7f, 0x68, 0x90, 0xca, 0x16, 0xde, 0xa6, 0xe8, 0x89, 0x3d, 0x96, 0xf0, 0xa3, 0xd, 0xa, 0x14, 0xe5, 0x55, 0x59, 0xfc, 0x9b, 0x83, 0x4, 0x91, 0xe3, 0xd2, 0x45, 0x1c, 0x81, 0xf6, 0xd1, 0xe},
+		}, "060e036b6579207f6890ca16dea6e8893d96f0a30d0a14e55559fc9b830491e3d2451c81f6d10e040202207f6890ca16dea6e8893d96f0a30d0a14e55559fc9b830491e3d2451c81f6d10e", false},
 		"leaf": {&Node{
 			subtreeHeight: 0,
 			size:          1,