core/vm: optimize MSTORE and SLOAD (#16939)

* vm/test: add tests+benchmarks for mstore

* core/vm: less alloc and copying for mstore

* core/vm: less allocs in sload

* vm: check for errors more correctly

3 files changed, 68 insertions, 11 deletions

diff --git a/core/vm/instructions.go b/core/vm/instructions.go
index 3a67e1865..9475322cd 100644
--- a/core/vm/instructions.go
+++ b/core/vm/instructions.go
@@ -556,7 +556,7 @@ func opMload(pc *uint64, evm *EVM, contract *Contract, memory *Memory, stack *St
 func opMstore(pc *uint64, evm *EVM, contract *Contract, memory *Memory, stack *Stack) ([]byte, error) {
 	// pop value of the stack
 	mStart, val := stack.pop(), stack.pop()
-	memory.Set(mStart.Uint64(), 32, math.PaddedBigBytes(val, 32))
+	memory.Set32(mStart.Uint64(), val)
 
 	evm.interpreter.intPool.put(mStart, val)
 	return nil, nil
@@ -570,9 +570,9 @@ func opMstore8(pc *uint64, evm *EVM, contract *Contract, memory *Memory, stack *
 }
 
 func opSload(pc *uint64, evm *EVM, contract *Contract, memory *Memory, stack *Stack) ([]byte, error) {
-	loc := common.BigToHash(stack.pop())
-	val := evm.StateDB.GetState(contract.Address(), loc).Big()
-	stack.push(val)
+	loc := stack.peek()
+	val := evm.StateDB.GetState(contract.Address(), common.BigToHash(loc))
+	loc.SetBytes(val.Bytes())
 	return nil, nil
 }
 
diff --git a/core/vm/instructions_test.go b/core/vm/instructions_test.go
index 0de558612..f51e6363f 100644
--- a/core/vm/instructions_test.go
+++ b/core/vm/instructions_test.go
@@ -425,3 +425,42 @@ func BenchmarkOpIsZero(b *testing.B) {
 	x := "FBCDEF090807060504030201ffffffffFBCDEF090807060504030201ffffffff"
 	opBenchmark(b, opIszero, x)
 }
+
+func TestOpMstore(t *testing.T) {
+	var (
+		env   = NewEVM(Context{}, nil, params.TestChainConfig, Config{})
+		stack = newstack()
+		mem   = NewMemory()
+	)
+	mem.Resize(64)
+	pc := uint64(0)
+	v := "abcdef00000000000000abba000000000deaf000000c0de00100000000133700"
+	stack.pushN(new(big.Int).SetBytes(common.Hex2Bytes(v)), big.NewInt(0))
+	opMstore(&pc, env, nil, mem, stack)
+	if got := common.Bytes2Hex(mem.Get(0, 32)); got != v {
+		t.Fatalf("Mstore fail, got %v, expected %v", got, v)
+	}
+	stack.pushN(big.NewInt(0x1), big.NewInt(0))
+	opMstore(&pc, env, nil, mem, stack)
+	if common.Bytes2Hex(mem.Get(0, 32)) != "0000000000000000000000000000000000000000000000000000000000000001" {
+		t.Fatalf("Mstore failed to overwrite previous value")
+	}
+}
+
+func BenchmarkOpMstore(bench *testing.B) {
+	var (
+		env   = NewEVM(Context{}, nil, params.TestChainConfig, Config{})
+		stack = newstack()
+		mem   = NewMemory()
+	)
+	mem.Resize(64)
+	pc := uint64(0)
+	memStart := big.NewInt(0)
+	value := big.NewInt(0x1337)
+
+	bench.ResetTimer()
+	for i := 0; i < bench.N; i++ {
+		stack.pushN(value, memStart)
+		opMstore(&pc, env, nil, mem, stack)
+	}
+}
diff --git a/core/vm/memory.go b/core/vm/memory.go
index d5cc7870b..43f6ff5bf 100644
--- a/core/vm/memory.go
+++ b/core/vm/memory.go
@@ -16,7 +16,12 @@
 
 package vm
 
-import "fmt"
+import (
+	"fmt"
+	"math/big"
+
+	"github.com/ethereum/go-ethereum/common/math"
+)
 
 // Memory implements a simple memory model for the ethereum virtual machine.
 type Memory struct {
@@ -30,19 +35,32 @@ func NewMemory() *Memory {
 
 // Set sets offset + size to value
 func (m *Memory) Set(offset, size uint64, value []byte) {
-	// length of store may never be less than offset + size.
-	// The store should be resized PRIOR to setting the memory
-	if size > uint64(len(m.store)) {
-		panic("INVALID memory: store empty")
-	}
-
 	// It's possible the offset is greater than 0 and size equals 0. This is because
 	// the calcMemSize (common.go) could potentially return 0 when size is zero (NO-OP)
 	if size > 0 {
+		// length of store may never be less than offset + size.
+		// The store should be resized PRIOR to setting the memory
+		if offset+size > uint64(len(m.store)) {
+			panic("invalid memory: store empty")
+		}
 		copy(m.store[offset:offset+size], value)
 	}
 }
 
+// Set32 sets the 32 bytes starting at offset to the value of val, left-padded with zeroes to
+// 32 bytes.
+func (m *Memory) Set32(offset uint64, val *big.Int) {
+	// length of store may never be less than offset + size.
+	// The store should be resized PRIOR to setting the memory
+	if offset+32 > uint64(len(m.store)) {
+		panic("invalid memory: store empty")
+	}
+	// Zero the memory area
+	copy(m.store[offset:offset+32], []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})
+	// Fill in relevant bits
+	math.ReadBits(val, m.store[offset:offset+32])
+}
+
 // Resize resizes the memory to size
 func (m *Memory) Resize(size uint64) {
 	if uint64(m.Len()) < size {