From fc88cea64864f61ed3ae24d502ade785deb5d077 Mon Sep 17 00:00:00 2001
From: lmittmann <3458786+lmittmann@users.noreply.github.com>
Date: Tue, 20 Aug 2024 14:31:06 +0200
Subject: [PATCH] core/vm: reuse Memory instances (#30137)

This PR adds a sync.Pool to reuse instances of Memory in EVMInterpreter.
---
 core/vm/instructions.go         |  4 ++--
 core/vm/interpreter.go          |  1 +
 core/vm/memory.go               | 21 ++++++++++++++++++-
 core/vm/runtime/runtime_test.go | 37 +++++++++++++++++++++++++++++++++
 4 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/core/vm/instructions.go b/core/vm/instructions.go
index 2e0f4c40ab..35d6393fba 100644
--- a/core/vm/instructions.go
+++ b/core/vm/instructions.go
@@ -871,14 +871,14 @@ func opStaticCall(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext)
 
 func opReturn(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]byte, error) {
 	offset, size := scope.Stack.pop(), scope.Stack.pop()
-	ret := scope.Memory.GetPtr(offset.Uint64(), size.Uint64())
+	ret := scope.Memory.GetCopy(offset.Uint64(), size.Uint64())
 
 	return ret, errStopToken
 }
 
 func opRevert(pc *uint64, interpreter *EVMInterpreter, scope *ScopeContext) ([]byte, error) {
 	offset, size := scope.Stack.pop(), scope.Stack.pop()
-	ret := scope.Memory.GetPtr(offset.Uint64(), size.Uint64())
+	ret := scope.Memory.GetCopy(offset.Uint64(), size.Uint64())
 
 	interpreter.returnData = ret
 	return ret, ErrExecutionReverted
diff --git a/core/vm/interpreter.go b/core/vm/interpreter.go
index 2b1ea38483..6e7d28a0ce 100644
--- a/core/vm/interpreter.go
+++ b/core/vm/interpreter.go
@@ -198,6 +198,7 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 	// they are returned to the pools
 	defer func() {
 		returnStack(stack)
+		mem.Free()
 	}()
 	contract.Input = input
 
diff --git a/core/vm/memory.go b/core/vm/memory.go
index 33203879ae..1ddd8d1ead 100644
--- a/core/vm/memory.go
+++ b/core/vm/memory.go
@@ -17,9 +17,17 @@
 package vm
 
 import (
+	"sync"
+
 	"github.com/holiman/uint256"
 )
 
+var memoryPool = sync.Pool{
+	New: func() any {
+		return &Memory{}
+	},
+}
+
 // Memory implements a simple memory model for the ethereum virtual machine.
 type Memory struct {
 	store       []byte
@@ -28,7 +36,18 @@ type Memory struct {
 
 // NewMemory returns a new memory model.
 func NewMemory() *Memory {
-	return &Memory{}
+	return memoryPool.Get().(*Memory)
+}
+
+// Free returns the memory to the pool.
+func (m *Memory) Free() {
+	// To reduce peak allocation, return only smaller memory instances to the pool.
+	const maxBufferSize = 16 << 10
+	if cap(m.store) <= maxBufferSize {
+		m.store = m.store[:0]
+		m.lastGasCost = 0
+		memoryPool.Put(m)
+	}
 }
 
 // Set sets offset + size to value
diff --git a/core/vm/runtime/runtime_test.go b/core/vm/runtime/runtime_test.go
index f52484606b..f56d4a7452 100644
--- a/core/vm/runtime/runtime_test.go
+++ b/core/vm/runtime/runtime_test.go
@@ -17,9 +17,11 @@
 package runtime
 
 import (
+	"encoding/binary"
 	"fmt"
 	"math/big"
 	"os"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -241,6 +243,41 @@ func BenchmarkEVM_SWAP1(b *testing.B) {
 	})
 }
 
+func BenchmarkEVM_RETURN(b *testing.B) {
+	// returns a contract that returns a zero-byte slice of len size
+	returnContract := func(size uint64) []byte {
+		contract := []byte{
+			byte(vm.PUSH8), 0, 0, 0, 0, 0, 0, 0, 0, // PUSH8 0xXXXXXXXXXXXXXXXX
+			byte(vm.PUSH0),  // PUSH0
+			byte(vm.RETURN), // RETURN
+		}
+		binary.BigEndian.PutUint64(contract[1:], size)
+		return contract
+	}
+
+	state, _ := state.New(types.EmptyRootHash, state.NewDatabase(rawdb.NewMemoryDatabase()), nil)
+	contractAddr := common.BytesToAddress([]byte("contract"))
+
+	for _, n := range []uint64{1_000, 10_000, 100_000, 1_000_000} {
+		b.Run(strconv.FormatUint(n, 10), func(b *testing.B) {
+			b.ReportAllocs()
+
+			contractCode := returnContract(n)
+			state.SetCode(contractAddr, contractCode)
+
+			for i := 0; i < b.N; i++ {
+				ret, _, err := Call(contractAddr, []byte{}, &Config{State: state})
+				if err != nil {
+					b.Fatal(err)
+				}
+				if uint64(len(ret)) != n {
+					b.Fatalf("expected return size %d, got %d", n, len(ret))
+				}
+			}
+		})
+	}
+}
+
 func fakeHeader(n uint64, parentHash common.Hash) *types.Header {
 	header := types.Header{
 		Coinbase:   common.HexToAddress("0x00000000000000000000000000000000deadbeef"),