From c8b54e557f8fac8c3f9cb3ee5e1dde43e2553c65 Mon Sep 17 00:00:00 2001
From: ryanmorphl2 <163962984+ryanmorphl2@users.noreply.github.com>
Date: Wed, 26 Jun 2024 14:50:46 +0800
Subject: [PATCH] ethdb/memorydb reduced allocations

---
 ethdb/memorydb/memorydb.go      | 14 +++++++-------
 ethdb/memorydb/memorydb_test.go | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/ethdb/memorydb/memorydb.go b/ethdb/memorydb/memorydb.go
index 902c78172..986738eb4 100644
--- a/ethdb/memorydb/memorydb.go
+++ b/ethdb/memorydb/memorydb.go
@@ -206,7 +206,7 @@ func (db *Database) Len() int {
 // keyvalue is a key-value tuple tagged with a deletion field to allow creating
 // memory-database write batches.
 type keyvalue struct {
-	key    []byte
+	key    string
 	value  []byte
 	delete bool
 }
@@ -221,14 +221,14 @@ type batch struct {
 
 // Put inserts the given value into the batch for later committing.
 func (b *batch) Put(key, value []byte) error {
-	b.writes = append(b.writes, keyvalue{common.CopyBytes(key), common.CopyBytes(value), false})
+	b.writes = append(b.writes, keyvalue{string(key), common.CopyBytes(value), false})
 	b.size += len(key) + len(value)
 	return nil
 }
 
 // Delete inserts the a key removal into the batch for later committing.
 func (b *batch) Delete(key []byte) error {
-	b.writes = append(b.writes, keyvalue{common.CopyBytes(key), nil, true})
+	b.writes = append(b.writes, keyvalue{string(key), nil, true})
 	b.size += len(key)
 	return nil
 }
@@ -245,10 +245,10 @@ func (b *batch) Write() error {
 
 	for _, keyvalue := range b.writes {
 		if keyvalue.delete {
-			delete(b.db.db, string(keyvalue.key))
+			delete(b.db.db, keyvalue.key)
 			continue
 		}
-		b.db.db[string(keyvalue.key)] = keyvalue.value
+		b.db.db[keyvalue.key] = keyvalue.value
 	}
 	return nil
 }
@@ -263,12 +263,12 @@ func (b *batch) Reset() {
 func (b *batch) Replay(w ethdb.KeyValueWriter) error {
 	for _, keyvalue := range b.writes {
 		if keyvalue.delete {
-			if err := w.Delete(keyvalue.key); err != nil {
+			if err := w.Delete([]byte(keyvalue.key)); err != nil {
 				return err
 			}
 			continue
 		}
-		if err := w.Put(keyvalue.key, keyvalue.value); err != nil {
+		if err := w.Put([]byte(keyvalue.key), keyvalue.value); err != nil {
 			return err
 		}
 	}
diff --git a/ethdb/memorydb/memorydb_test.go b/ethdb/memorydb/memorydb_test.go
index 5219aaa5e..58f1b489c 100644
--- a/ethdb/memorydb/memorydb_test.go
+++ b/ethdb/memorydb/memorydb_test.go
@@ -17,6 +17,7 @@
 package memorydb
 
 import (
+	"encoding/binary"
 	"testing"
 
 	"github.com/scroll-tech/go-ethereum/ethdb"
@@ -30,3 +31,20 @@ func TestMemoryDB(t *testing.T) {
 		})
 	})
 }
+
+// BenchmarkBatchAllocs measures the time/allocs for storing 120 kB of data
+func BenchmarkBatchAllocs(b *testing.B) {
+	b.ReportAllocs()
+	var key = make([]byte, 20)
+	var val = make([]byte, 100)
+	// 120 * 1_000 -> 120_000 == 120kB
+	for i := 0; i < b.N; i++ {
+		batch := New().NewBatch()
+		for j := uint64(0); j < 1000; j++ {
+			binary.BigEndian.PutUint64(key, j)
+			binary.BigEndian.PutUint64(val, j)
+			batch.Put(key, val)
+		}
+		batch.Write()
+	}
+}