Merkle db intermediate node key compression

x/merkledb/db.go

@@ -1187,21 +1187,26 @@ func (db *merkleDB) getNode(key path, hasValue bool) (*node, error) {
 // The returned []byte is taken from [bufferPool] and
 // should be returned to it when the caller is done with it.
 func addPrefixToKey(bufferPool *sync.Pool, prefix []byte, key []byte) []byte {
-	prefixedKey := bufferPool.Get().([]byte)
 	prefixLen := len(prefix)
 	keyLen := prefixLen + len(key)
-	if cap(prefixedKey) >= keyLen {
+	prefixedKey := getBufferFromPool(bufferPool, keyLen)
+	copy(prefixedKey, prefix)
+	copy(prefixedKey[prefixLen:], key)
+	return prefixedKey
+}
+
+func getBufferFromPool(bufferPool *sync.Pool, size int) []byte {
+	buffer := bufferPool.Get().([]byte)
+	if cap(buffer) >= size {
 		// The [] byte we got from the pool is big enough to hold the prefixed key
-		prefixedKey = prefixedKey[:keyLen]
+		buffer = buffer[:size]
 	} else {
 		// The []byte from the pool wasn't big enough.
 		// Put it back and allocate a new, bigger one
-		bufferPool.Put(prefixedKey)
-		prefixedKey = make([]byte, keyLen)
+		bufferPool.Put(buffer)
+		buffer = make([]byte, size)
 	}
-	copy(prefixedKey, prefix)
-	copy(prefixedKey[prefixLen:], key)
-	return prefixedKey
+	return buffer
 }
 
 // cacheEntrySize returns a rough approximation of the memory consumed by storing the path and node

x/merkledb/intermediate_node_db.go

@@ -89,13 +89,13 @@ func (db *intermediateNodeDB) onEviction(key path, n *node) error {
 }
 
 func (db *intermediateNodeDB) addToBatch(b database.Batch, key path, n *node) error {
-	prefixedKey := addPrefixToKey(db.bufferPool, intermediateNodePrefix, key.Bytes())
-	defer db.bufferPool.Put(prefixedKey)
+	dbKey := db.constructDBKey(key)
+	defer db.bufferPool.Put(dbKey)
 	db.metrics.DatabaseNodeWrite()
 	if n == nil {
-		return b.Delete(prefixedKey)
+		return b.Delete(dbKey)
 	}
-	return b.Put(prefixedKey, n.bytes())
+	return b.Put(dbKey, n.bytes())
 }
 
 func (db *intermediateNodeDB) Get(key path) (*node, error) {
@@ -108,17 +108,43 @@ func (db *intermediateNodeDB) Get(key path) (*node, error) {
 	}
 	db.metrics.IntermediateNodeCacheMiss()
 
-	prefixedKey := addPrefixToKey(db.bufferPool, intermediateNodePrefix, key.Bytes())
+	dbKey := db.constructDBKey(key)
 	db.metrics.DatabaseNodeRead()
-	nodeBytes, err := db.baseDB.Get(prefixedKey)
+	nodeBytes, err := db.baseDB.Get(dbKey)
 	if err != nil {
 		return nil, err
 	}
-	db.bufferPool.Put(prefixedKey)
+	db.bufferPool.Put(dbKey)
 
 	return parseNode(key, nodeBytes)
 }
 
+// constructDBKey creates a key that can be used in the base db
+// We need to be able to differentiate between two paths of equal byte length but different token length,
+// so it adds a padding that can be used to differentiate.
+// Additionally, it adds the prefix indicating it is part of the intermediateDB
+func (db *intermediateNodeDB) constructDBKey(key path) []byte {
+	compressedKey := key.Serialize()
+
+	// add one additional byte to store padding when the path has a length that fits into a whole number of bytes
+	remainder := compressedKey.NibbleLength % 2
+	keyLen := len(compressedKey.Value)
+	if remainder == 0 {
+		keyLen++
+	}
+	dbKey := getBufferFromPool(db.bufferPool, keyLen)
+	defer db.bufferPool.Put(dbKey)
+
+	copy(dbKey, compressedKey.Value)
+	if remainder == 0 {
+		dbKey[keyLen-1] = 0b1000_0000
+	} else {
+		dbKey[keyLen-1] += 0b0000_1000
+	}
+
+	return addPrefixToKey(db.bufferPool, intermediateNodePrefix, dbKey)
+}
+
 func (db *intermediateNodeDB) Put(key path, n *node) error {
 	return db.nodeCache.Put(key, n)
 }

x/merkledb/intermediate_node_db_test.go

@@ -70,7 +70,7 @@ func TestIntermediateNodeDB(t *testing.T) {
 	added := 0
 	for {
 		key := newPath([]byte{byte(added)})
-		node := newNode(nil, EmptyPath)
+		node := newNode(nil, key)
 		node.setValue(maybe.Some([]byte{byte(added)}))
 		newExpectedSize := expectedSize + cacheEntrySize(key, node)
 		if newExpectedSize > cacheSize {
@@ -90,7 +90,7 @@ func TestIntermediateNodeDB(t *testing.T) {
 	// of all but 2 elements. 2 elements remain rather than 1 element because of
 	// the added key prefix increasing the size tracked by the batch.
 	key := newPath([]byte{byte(added)})
-	node := newNode(nil, EmptyPath)
+	node := newNode(nil, key)
 	node.setValue(maybe.Some([]byte{byte(added)}))
 	require.NoError(db.Put(key, node))