From 55edee556590a1f4df96f1dad4c909f6d6c01ced Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Mon, 4 Nov 2024 17:00:57 +0100
Subject: [PATCH] core/vm/program: minor UX tweaks

---
 core/vm/program/program.go      |  29 ++++++-
 core/vm/program/program_test.go | 121 ++++++++++++++++++++++++++++-
 core/vm/runtime/runtime_test.go | 133 +++++++++-----------------------
 3 files changed, 178 insertions(+), 105 deletions(-)

diff --git a/core/vm/program/program.go b/core/vm/program/program.go
index 19fb861332fa..e8d667fa9d98 100644
--- a/core/vm/program/program.go
+++ b/core/vm/program/program.go
@@ -127,14 +127,14 @@ func (p *Program) Push0() *Program {
 	return p.Op(vm.PUSH0)
 }
 
-// Bytecode returns the Program bytecode
-func (p *Program) Bytecode() []byte {
+// Bytes returns the Program bytecode
+func (p *Program) Bytes() []byte {
 	return p.code
 }
 
 // Hex returns the Program bytecode as a hex string
 func (p *Program) Hex() string {
-	return fmt.Sprintf("%02x", p.Bytecode())
+	return fmt.Sprintf("%02x", p.Bytes())
 }
 
 // ExtcodeCopy performsa an extcodecopy invocation
@@ -280,6 +280,29 @@ func (p *Program) Mstore(data []byte, memStart uint32) *Program {
 	return p
 }
 
+// MstorePadded stores the provided data (into the memory area starting at memStart).
+// If data does not align on 32 bytes, it will be LHS-padded.
+// For example, providing data 0x1122, it will do a PUSH2:
+// PUSH2 0x1122, resulting in
+// stack: 0x0000000000000000000000000000000000000000000000000000000000001122
+// followed by MSTORE(0,0)
+// And thus, the resulting memory will be
+// [ 0000000000000000000000000000000000000000000000000000000000001122 ]
+func (p *Program) MstorePadded(data []byte, memStart uint32) *Program {
+	var idx = 0
+	// We need to store it in chunks of 32 bytes
+	for ; idx < len(data); idx += 32 {
+		end := min(len(data), idx+32)
+		chunk := data[idx:end]
+		// push the value
+		p.Push(chunk)
+		// push the memory index
+		p.Push(uint32(idx) + memStart)
+		p.Op(vm.MSTORE)
+	}
+	return p
+}
+
 // MemToStorage copies the given memory area into SSTORE slots,
 // It expects data to be aligned to 32 byte, and does not zero out
 // remainders if some data is not
diff --git a/core/vm/program/program_test.go b/core/vm/program/program_test.go
index 5291e034effd..352b59ce960f 100644
--- a/core/vm/program/program_test.go
+++ b/core/vm/program/program_test.go
@@ -143,7 +143,7 @@ func TestCreateAndCall(t *testing.T) {
 	deployed.Return(0, 32)
 
 	// Pack them
-	ctor.ReturnData(deployed.Bytecode())
+	ctor.ReturnData(deployed.Bytes())
 	// Verify constructor + runtime code
 	{
 		want := "6005600055606060005360006001536054600253606060035360006004536052600553606060065360206007536060600853600060095360f3600a53600b6000f3"
@@ -155,21 +155,134 @@ func TestCreateAndCall(t *testing.T) {
 
 func TestCreate2Call(t *testing.T) {
 	// Some runtime code
-	runtime := New().Ops(vm.ADDRESS, vm.SELFDESTRUCT).Bytecode()
+	runtime := New().Ops(vm.ADDRESS, vm.SELFDESTRUCT).Bytes()
 	want := common.FromHex("0x30ff")
 	if !bytes.Equal(want, runtime) {
 		t.Fatalf("runtime code error\nwant: %x\nhave: %x\n", want, runtime)
 	}
 	// A constructor returning the runtime code
-	initcode := New().ReturnData(runtime).Bytecode()
+	initcode := New().ReturnData(runtime).Bytes()
 	want = common.FromHex("603060005360ff60015360026000f3")
 	if !bytes.Equal(want, initcode) {
 		t.Fatalf("initcode error\nwant: %x\nhave: %x\n", want, initcode)
 	}
 	// A factory invoking the constructor
-	outer := New().Create2AndCall(initcode, nil).Bytecode()
+	outer := New().Create2AndCall(initcode, nil).Bytes()
 	want = common.FromHex("60606000536030600153606060025360006003536053600453606060055360ff6006536060600753600160085360536009536060600a536002600b536060600c536000600d5360f3600e536000600f60006000f560006000600060006000855af15050")
 	if !bytes.Equal(want, outer) {
 		t.Fatalf("factory error\nwant: %x\nhave: %x\n", want, outer)
 	}
 }
+
+func TestGenerator(t *testing.T) {
+	for i, tc := range []struct {
+		want   []byte
+		haveFn func() []byte
+	}{
+		{ // CREATE
+			want: []byte{
+				// Store initcode in memory at 0x00 (5 bytes left-padded to 32 bytes)
+				byte(vm.PUSH5),
+				// Init code: PUSH1 0, PUSH1 0, RETURN (3 steps)
+				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.RETURN),
+				byte(vm.PUSH1), 0,
+				byte(vm.MSTORE),
+				// length, offset, value
+				byte(vm.PUSH1), 5, byte(vm.PUSH1), 27, byte(vm.PUSH1), 0,
+				byte(vm.CREATE),
+				byte(vm.POP),
+			},
+			haveFn: func() []byte {
+				initcode := New().Return(0, 0).Bytes()
+				return New().MstorePadded(initcode, 0).
+					Push(len(initcode)).      // length
+					Push(32 - len(initcode)). // offset
+					Push(0).                  // value
+					Op(vm.CREATE).
+					Op(vm.POP).Bytes()
+			},
+		},
+		{ // CREATE2
+			want: []byte{
+				// Store initcode in memory at 0x00 (5 bytes left-padded to 32 bytes)
+				byte(vm.PUSH5),
+				// Init code: PUSH1 0, PUSH1 0, RETURN (3 steps)
+				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.RETURN),
+				byte(vm.PUSH1), 0,
+				byte(vm.MSTORE),
+				// salt, length, offset, value
+				byte(vm.PUSH1), 1, byte(vm.PUSH1), 5, byte(vm.PUSH1), 27, byte(vm.PUSH1), 0,
+				byte(vm.CREATE2),
+				byte(vm.POP),
+			},
+			haveFn: func() []byte {
+				initcode := New().Return(0, 0).Bytes()
+				return New().MstorePadded(initcode, 0).
+					Push(1).                  // salt
+					Push(len(initcode)).      // length
+					Push(32 - len(initcode)). // offset
+					Push(0).                  // value
+					Op(vm.CREATE2).
+					Op(vm.POP).Bytes()
+			},
+		},
+		{ // CALL
+			want: []byte{
+				// outsize, outoffset, insize, inoffset
+				byte(vm.PUSH1), 0, byte(vm.DUP1), byte(vm.DUP1), byte(vm.DUP1),
+				byte(vm.DUP1),        // value
+				byte(vm.PUSH1), 0xbb, //address
+				byte(vm.GAS), // gas
+				byte(vm.CALL),
+				byte(vm.POP),
+			},
+			haveFn: func() []byte {
+				return New().Call(nil, 0xbb, 0, 0, 0, 0, 0).Op(vm.POP).Bytes()
+			},
+		},
+		{ // CALLCODE
+			want: []byte{
+				// outsize, outoffset, insize, inoffset
+				byte(vm.PUSH1), 0, byte(vm.DUP1), byte(vm.DUP1), byte(vm.DUP1),
+				byte(vm.PUSH1), 0, // value
+				byte(vm.PUSH1), 0xcc, //address
+				byte(vm.GAS), // gas
+				byte(vm.CALLCODE),
+				byte(vm.POP),
+			},
+			haveFn: func() []byte {
+				return New().CallCode(nil, 0xcc, 0, 0, 0, 0, 0).Op(vm.POP).Bytes()
+			},
+		},
+		{ // STATICCALL
+			want: []byte{
+				// outsize, outoffset, insize, inoffset
+				byte(vm.PUSH1), 0, byte(vm.DUP1), byte(vm.DUP1), byte(vm.DUP1),
+				byte(vm.PUSH1), 0xdd, //address
+				byte(vm.GAS), // gas
+				byte(vm.STATICCALL),
+				byte(vm.POP),
+			},
+			haveFn: func() []byte {
+				return New().StaticCall(nil, 0xdd, 0, 0, 0, 0).Op(vm.POP).Bytes()
+			},
+		},
+		{ // DELEGATECALL
+			want: []byte{
+				// outsize, outoffset, insize, inoffset
+				byte(vm.PUSH1), 0, byte(vm.DUP1), byte(vm.DUP1), byte(vm.DUP1),
+				byte(vm.PUSH1), 0xee, //address
+				byte(vm.GAS), // gas
+				byte(vm.DELEGATECALL),
+				byte(vm.POP),
+			},
+			haveFn: func() []byte {
+				return New().DelegateCall(nil, 0xee, 0, 0, 0, 0).Op(vm.POP).Bytes()
+			},
+		},
+	} {
+		if have := tc.haveFn(); !bytes.Equal(have, tc.want) {
+			t.Fatalf("test %d error\nhave: %x\nwant: %x\n", i, have, tc.want)
+		}
+	}
+}
diff --git a/core/vm/runtime/runtime_test.go b/core/vm/runtime/runtime_test.go
index dcdcb674d762..c8064e85abbc 100644
--- a/core/vm/runtime/runtime_test.go
+++ b/core/vm/runtime/runtime_test.go
@@ -442,22 +442,22 @@ func BenchmarkSimpleLoop(b *testing.B) {
 	// Call identity, and pop return value
 	staticCallIdentity := p.
 		StaticCall(nil, 0x4, 0, 0, 0, 0).
-		Op(vm.POP).Jump(lbl).Bytecode() // pop return value and jump to label
+		Op(vm.POP).Jump(lbl).Bytes() // pop return value and jump to label
 
 	p, lbl = program.New().Jumpdest()
 	callIdentity := p.
 		Call(nil, 0x4, 0, 0, 0, 0, 0).
-		Op(vm.POP).Jump(lbl).Bytecode() // pop return value and jump to label
+		Op(vm.POP).Jump(lbl).Bytes() // pop return value and jump to label
 
 	p, lbl = program.New().Jumpdest()
 	callInexistant := p.
 		Call(nil, 0xff, 0, 0, 0, 0, 0).
-		Op(vm.POP).Jump(lbl).Bytecode() // pop return value and jump to label
+		Op(vm.POP).Jump(lbl).Bytes() // pop return value and jump to label
 
 	p, lbl = program.New().Jumpdest()
 	callEOA := p.
 		Call(nil, 0xE0, 0, 0, 0, 0, 0). // call addr of EOA
-		Op(vm.POP).Jump(lbl).Bytecode() // pop return value and jump to label
+		Op(vm.POP).Jump(lbl).Bytes() // pop return value and jump to label
 
 	p, lbl = program.New().Jumpdest()
 	// Push as if we were making call, then pop it off again, and loop
@@ -465,19 +465,19 @@ func BenchmarkSimpleLoop(b *testing.B) {
 		Ops(vm.DUP1, vm.DUP1, vm.DUP1).
 		Push(0x4).
 		Ops(vm.GAS, vm.POP, vm.POP, vm.POP, vm.POP, vm.POP, vm.POP).
-		Jump(lbl).Bytecode()
+		Jump(lbl).Bytes()
 
 	p, lbl = program.New().Jumpdest()
 	loopingCode2 := p.
 		Push(0x01020304).Push(uint64(0x0102030405)).
 		Ops(vm.POP, vm.POP).
 		Op(vm.PUSH6).Append(make([]byte, 6)).Op(vm.JUMP). // Jumpdest zero expressed in 6 bytes
-		Bytecode()
+		Bytes()
 
 	p, lbl = program.New().Jumpdest()
 	callRevertingContractWithInput := p.
 		Call(nil, 0xee, 0, 0, 0x20, 0x0, 0x0).
-		Op(vm.POP).Jump(lbl).Bytecode() // pop return value and jump to label
+		Op(vm.POP).Jump(lbl).Bytes() // pop return value and jump to label
 
 	//tracer := logger.NewJSONLogger(nil, os.Stdout)
 	//Execute(loopingCode, nil, &Config{
@@ -716,104 +716,49 @@ func TestRuntimeJSTracer(t *testing.T) {
 		this.exits++;
 		this.gasUsed = res.getGasUsed();
 	}}`}
+	initcode := program.New().Return(0, 0).Bytes()
 	tests := []struct {
 		code []byte
 		// One result per tracer
 		results []string
 	}{
-		{
-			// CREATE
-			code: []byte{
-				// Store initcode in memory at 0x00 (5 bytes left-padded to 32 bytes)
-				byte(vm.PUSH5),
-				// Init code: PUSH1 0, PUSH1 0, RETURN (3 steps)
-				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.RETURN),
-				byte(vm.PUSH1), 0,
-				byte(vm.MSTORE),
-				// length, offset, value
-				byte(vm.PUSH1), 5, byte(vm.PUSH1), 27, byte(vm.PUSH1), 0,
-				byte(vm.CREATE),
-				byte(vm.POP),
-			},
+		{ // CREATE
+			code: program.New().MstorePadded(initcode, 0).
+				Push(len(initcode)). // length
+				Push(32 - len(initcode)). // offset
+				Push(0). // value
+				Op(vm.CREATE).
+				Op(vm.POP).Bytes(),
 			results: []string{`"1,1,952853,6,12"`, `"1,1,952853,6,0"`},
 		},
-		{
-			// CREATE2
-			code: []byte{
-				// Store initcode in memory at 0x00 (5 bytes left-padded to 32 bytes)
-				byte(vm.PUSH5),
-				// Init code: PUSH1 0, PUSH1 0, RETURN (3 steps)
-				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.RETURN),
-				byte(vm.PUSH1), 0,
-				byte(vm.MSTORE),
-				// salt, length, offset, value
-				byte(vm.PUSH1), 1, byte(vm.PUSH1), 5, byte(vm.PUSH1), 27, byte(vm.PUSH1), 0,
-				byte(vm.CREATE2),
-				byte(vm.POP),
-			},
+		{ // CREATE2
+			code: program.New().MstorePadded(initcode, 0).
+				Push(1). // salt
+				Push(len(initcode)). // length
+				Push(32 - len(initcode)). // offset
+				Push(0). // value
+				Op(vm.CREATE2).
+				Op(vm.POP).Bytes(),
 			results: []string{`"1,1,952844,6,13"`, `"1,1,952844,6,0"`},
 		},
-		{
-			// CALL
-			code: []byte{
-				// outsize, outoffset, insize, inoffset
-				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0,
-				byte(vm.PUSH1), 0, // value
-				byte(vm.PUSH1), 0xbb, //address
-				byte(vm.GAS), // gas
-				byte(vm.CALL),
-				byte(vm.POP),
-			},
+		{ // CALL
+			code:    program.New().Call(nil, 0xbb, 0, 0, 0, 0, 0).Op(vm.POP).Bytes(),
 			results: []string{`"1,1,981796,6,13"`, `"1,1,981796,6,0"`},
 		},
-		{
-			// CALLCODE
-			code: []byte{
-				// outsize, outoffset, insize, inoffset
-				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0,
-				byte(vm.PUSH1), 0, // value
-				byte(vm.PUSH1), 0xcc, //address
-				byte(vm.GAS), // gas
-				byte(vm.CALLCODE),
-				byte(vm.POP),
-			},
+		{ // CALLCODE
+			code:    program.New().CallCode(nil, 0xcc, 0, 0, 0, 0, 0).Op(vm.POP).Bytes(),
 			results: []string{`"1,1,981796,6,13"`, `"1,1,981796,6,0"`},
 		},
-		{
-			// STATICCALL
-			code: []byte{
-				// outsize, outoffset, insize, inoffset
-				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0,
-				byte(vm.PUSH1), 0xdd, //address
-				byte(vm.GAS), // gas
-				byte(vm.STATICCALL),
-				byte(vm.POP),
-			},
+		{ // STATICCALL
+			code:    program.New().StaticCall(nil, 0xdd, 0, 0, 0, 0).Op(vm.POP).Bytes(),
 			results: []string{`"1,1,981799,6,12"`, `"1,1,981799,6,0"`},
 		},
-		{
-			// DELEGATECALL
-			code: []byte{
-				// outsize, outoffset, insize, inoffset
-				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0,
-				byte(vm.PUSH1), 0xee, //address
-				byte(vm.GAS), // gas
-				byte(vm.DELEGATECALL),
-				byte(vm.POP),
-			},
+		{ // DELEGATECALL
+			code:    program.New().DelegateCall(nil, 0xee, 0, 0, 0, 0).Op(vm.POP).Bytes(),
 			results: []string{`"1,1,981799,6,12"`, `"1,1,981799,6,0"`},
 		},
-		{
-			// CALL self-destructing contract
-			code: []byte{
-				// outsize, outoffset, insize, inoffset
-				byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0, byte(vm.PUSH1), 0,
-				byte(vm.PUSH1), 0, // value
-				byte(vm.PUSH1), 0xff, //address
-				byte(vm.GAS), // gas
-				byte(vm.CALL),
-				byte(vm.POP),
-			},
+		{ // CALL self-destructing contract
+			code:    program.New().Call(nil, 0xff, 0, 0, 0, 0, 0).Op(vm.POP).Bytes(),
 			results: []string{`"2,2,0,5003,12"`, `"2,2,0,5003,0"`},
 		},
 	}
@@ -896,16 +841,8 @@ func TestJSTracerCreateTx(t *testing.T) {
 
 func BenchmarkTracerStepVsCallFrame(b *testing.B) {
 	// Simply pushes and pops some values in a loop
-	code := []byte{
-		byte(vm.JUMPDEST),
-		byte(vm.PUSH1), 0,
-		byte(vm.PUSH1), 0,
-		byte(vm.POP),
-		byte(vm.POP),
-		byte(vm.PUSH1), 0, // jumpdestination
-		byte(vm.JUMP),
-	}
-
+	p, lbl := program.New().Jumpdest()
+	code := p.Push(0).Push(0).Ops(vm.POP, vm.POP).Jump(lbl).Bytes()
 	stepTracer := `
 	{
 	step: function() {},