w3vm: Added State Setter to VM

w3vm/vm.go

@@ -202,6 +202,11 @@ func (vm *VM) Nonce(addr common.Address) (uint64, error) {
  return nonce, nil
 }
 
+// SetNonce sets the nonce of the given address.
+func (vm *VM) SetNonce(addr common.Address, nonce uint64) {
+ vm.db.SetNonce(addr, nonce)
+}
+
 // Balance returns the balance of the given address.
 func (vm *VM) Balance(addr common.Address) (*big.Int, error) {
  balance := vm.db.GetBalance(addr)
@@ -211,6 +216,11 @@ func (vm *VM) Balance(addr common.Address) (*big.Int, error) {
  return balance.ToBig(), nil
 }
 
+// SetBalance sets the balance of the given address.
+func (vm *VM) SetBalance(addr common.Address, balance *big.Int) {
+ vm.db.SetBalance(addr, uint256.MustFromBig(balance), tracing.BalanceChangeUnspecified)
+}
+
 // Code returns the code of the given address.
 func (vm *VM) Code(addr common.Address) ([]byte, error) {
  code := vm.db.GetCode(addr)
@@ -220,6 +230,11 @@ func (vm *VM) Code(addr common.Address) ([]byte, error) {
  return code, nil
 }
 
+// SetCode sets the code of the given address.
+func (vm *VM) SetCode(addr common.Address, code []byte) {
+ vm.db.SetCode(addr, code)
+}
+
 // StorageAt returns the state of the given address at the give storage slot.
 func (vm *VM) StorageAt(addr common.Address, slot common.Hash) (common.Hash, error) {
  val := vm.db.GetState(addr, slot)
@@ -229,6 +244,11 @@ func (vm *VM) StorageAt(addr common.Address, slot common.Hash) (common.Hash, err
  return val, nil
 }
 
+// SetStorageAt sets the state of the given address at the given storage slot.
+func (vm *VM) SetStorageAt(addr common.Address, slot, val common.Hash) {
+ vm.db.SetState(addr, slot, val)
+}
+
 // Snapshot the current state of the VM. The returned state can only be rolled
 // back to once. Use [state.StateDB.Copy] if you need to rollback multiple times.
 func (vm *VM) Snapshot() *state.StateDB { return vm.db.Copy() }

w3vm/vm_test.go

@@ -1,6 +1,7 @@
 package w3vm_test
 
 import (
+ "bytes"
  _ "embed"
  "errors"
  "fmt"
@@ -47,6 +48,66 @@ var (
  ))
 )
 
+func TestVMSetNonce(t *testing.T) {
+ vm, _ := w3vm.New()
+
+ if nonce, _ := vm.Nonce(addr0); nonce != 0 {
+ t.Fatalf("Nonce: want 0, got %d", nonce)
+ }
+
+ want := uint64(42)
+ vm.SetNonce(addr0, want)
+
+ if nonce, _ := vm.Nonce(addr0); want != nonce {
+ t.Fatalf("Nonce: want %d, got %d", want, nonce)
+ }
+}
+
+func TestVMSetBalance(t *testing.T) {
+ vm, _ := w3vm.New()
+
+ if balance, _ := vm.Balance(addr0); balance.Sign() != 0 {
+ t.Fatalf("Balance: want 0, got %s", balance)
+ }
+
+ want := w3.I("1 ether")
+ vm.SetBalance(addr0, want)
+
+ if balance, _ := vm.Balance(addr0); want.Cmp(balance) != 0 {
+ t.Fatalf("Balance: want %s ether, got %s ether", w3.FromWei(want, 18), w3.FromWei(balance, 18))
+ }
+}
+
+func TestVMSetCode(t *testing.T) {
+ vm, _ := w3vm.New()
+
+ if code, _ := vm.Code(addr0); len(code) != 0 {
+ t.Fatalf("Code: want empty, got %x", code)
+ }
+
+ want := []byte{0xc0, 0xfe}
+ vm.SetCode(addr0, want)
+
+ if code, _ := vm.Code(addr0); !bytes.Equal(want, code) {
+ t.Fatalf("Code: want %x, got %x", want, code)
+ }
+}
+
+func TestVMSetStorage(t *testing.T) {
+ vm, _ := w3vm.New()
+
+ if storage, _ := vm.StorageAt(addr0, common.Hash{}); storage != w3.Hash0 {
+ t.Fatalf("Storage: want empty, got %x", storage)
+ }
+
+ want := common.Hash{0xc0, 0xfe}
+ vm.SetStorageAt(addr0, common.Hash{}, want)
+
+ if storage, _ := vm.StorageAt(addr0, common.Hash{}); want != storage {
+ t.Fatalf("Storage: want %x, got %x", want, storage)
+ }
+}
+
 func TestVMApply(t *testing.T) {
  tests := []struct {
  PreState w3types.State
@@ -189,6 +250,25 @@ func TestVMApply(t *testing.T) {
  ContractAddress: ptr(crypto.CreateAddress(addr1, 1)),
  },
  },
+ { // EOA with storage
+ PreState: w3types.State{
+ addr0: {
+ Balance: w3.I("1 ether"),
+ Storage: w3types.Storage{
+ common.Hash{0x1}: common.Hash{0x2},
+ },
+ },
+ },
+ Message: &w3types.Message{
+ From: addr0,
+ To: &addr1,
+ Value: w3.I("1 ether"),
+ },
+ WantReceipt: &w3vm.Receipt{
+ GasUsed: 21_000,
+ GasLimit: 21_000,
+ },
+ },
  }
 
  for i, test := range tests {