vm: reintroduce modifyAccountFields (#1763)

* add modifyAccountFields method and tests
* use new method in eei
* update usage in various tests

Author: ryanio

packages/client/test/miner/miner.spec.ts

@@ -30,9 +30,8 @@ const B = {
 }
 
 const setBalance = async (stateManager: StateManager, address: Address, balance: BN) => {
-  // this fn can be replaced with modifyAccountFields() when #1369 is available
   await stateManager.checkpoint()
-  await stateManager.putAccount(address, new Account(new BN(0), balance))
+  await stateManager.modifyAccountFields(address, { balance })
   await stateManager.commit()
 }
 

packages/client/test/miner/pendingBlock.spec.ts

@@ -27,9 +27,8 @@ const B = {
 }
 
 const setBalance = async (stateManager: StateManager, address: Address, balance: BN) => {
-  // this fn can be replaced with modifyAccountFields() when #1369 is available
   await stateManager.checkpoint()
-  await stateManager.putAccount(address, new Account(new BN(0), balance))
+  await stateManager.modifyAccountFields(address, { balance })
   await stateManager.commit()
 }
 

packages/vm/src/evm/eei.ts

@@ -452,9 +452,9 @@ export default class EEI {
     await this._state.putAccount(toAddress, toAccount)
 
     // Subtract from contract balance
-    const account = await this._state.getAccount(this._env.address)
-    account.balance = new BN(0)
-    await this._state.putAccount(this._env.address, account)
+    await this._state.modifyAccountFields(this._env.address, {
+      balance: new BN(0),
+    })
 
     trap(ERROR.STOP)
   }

packages/vm/src/state/interface.ts

@@ -9,12 +9,15 @@ export interface StorageDump {
   [key: string]: string
 }
 
+export type AccountFields = Partial<Pick<Account, 'nonce' | 'balance' | 'stateRoot' | 'codeHash'>>
+
 export interface StateManager {
   copy(): StateManager
   getAccount(address: Address): Promise<Account>
   putAccount(address: Address, account: Account): Promise<void>
   deleteAccount(address: Address): Promise<void>
   touchAccount(address: Address): void
+  modifyAccountFields(address: Address, accountFields: AccountFields): Promise<void>
   putContractCode(address: Address, value: Buffer): Promise<void>
   getContractCode(address: Address): Promise<Buffer>
   getContractStorage(address: Address, key: Buffer): Promise<Buffer>

packages/vm/src/state/stateManager.ts

@@ -15,7 +15,7 @@ import {
   setLengthLeft,
 } from 'ethereumjs-util'
 import Common from '@ethereumjs/common'
-import { StateManager, StorageDump } from './interface'
+import { StateManager, StorageDump, AccountFields } from './interface'
 import Cache, { getCb, putCb } from './cache'
 import { BaseStateManager } from './'
 import { short } from '../evm/opcodes'
@@ -132,12 +132,10 @@ export default class DefaultStateManager extends BaseStateManager implements Sta
     const key = Buffer.concat([CODEHASH_PREFIX, codeHash])
     await this._trie.db.put(key, value)
 
-    const account = await this.getAccount(address)
     if (this.DEBUG) {
       this._debug(`Update codeHash (-> ${short(codeHash)}) for account ${address}`)
     }
-    account.codeHash = codeHash
-    await this.putAccount(address, account)
+    await this.modifyAccountFields(address, { codeHash })
   }
 
   /**
@@ -520,4 +518,20 @@ export default class DefaultStateManager extends BaseStateManager implements Sta
     }
     return false
   }
+
+  /**
+   * Gets the account associated with `address`, modifies the given account
+   * fields, then saves the account into state. Account fields can include
+   * `nonce`, `balance`, `stateRoot`, and `codeHash`.
+   * @param address - Address of the account to modify
+   * @param accountFields - Object containing account fields and values to modify
+   */
+  async modifyAccountFields(address: Address, accountFields: AccountFields): Promise<void> {
+    const account = await this.getAccount(address)
+    account.nonce = accountFields.nonce ?? account.nonce
+    account.balance = accountFields.balance ?? account.balance
+    account.stateRoot = accountFields.stateRoot ?? account.stateRoot
+    account.codeHash = accountFields.codeHash ?? account.codeHash
+    await this.putAccount(address, account)
+  }
 }

packages/vm/tests/api/EIPs/eip-1559-FeeMarket.spec.ts

@@ -109,11 +109,8 @@ tape('EIP1559 tests', (t) => {
       { common }
     )
     const block2 = makeBlock(GWEI, tx2, 1)
-    account = await vm.stateManager.getAccount(sender)
-    account.balance = balance
-    await vm.stateManager.putAccount(sender, account)
-    miner.balance = new BN(0)
-    await vm.stateManager.putAccount(coinbase, miner)
+    await vm.stateManager.modifyAccountFields(sender, { balance })
+    await vm.stateManager.modifyAccountFields(coinbase, { balance: new BN(0) })
     const results2 = await vm.runTx({
       tx: block2.transactions[0],
       block: block2,
@@ -140,11 +137,8 @@ tape('EIP1559 tests', (t) => {
       { common }
     )
     const block3 = makeBlock(GWEI, tx3, 0)
-    account = await vm.stateManager.getAccount(sender)
-    account.balance = balance
-    await vm.stateManager.putAccount(sender, account)
-    miner.balance = new BN(0)
-    await vm.stateManager.putAccount(coinbase, miner)
+    await vm.stateManager.modifyAccountFields(sender, { balance })
+    await vm.stateManager.modifyAccountFields(coinbase, { balance: new BN(0) })
     const results3 = await vm.runTx({
       tx: block3.transactions[0],
       block: block3,
@@ -180,10 +174,8 @@ tape('EIP1559 tests', (t) => {
     )
     const block = makeBlock(GWEI, tx, 2)
     const vm = new VM({ common })
-    const account = await vm.stateManager.getAccount(sender)
     const balance = GWEI.muln(210000).muln(10)
-    account.balance = balance
-    await vm.stateManager.putAccount(sender, account)
+    await vm.stateManager.modifyAccountFields(sender, { balance })
 
     /**
      * GASPRICE

packages/vm/tests/api/EIPs/eip-3198-BaseFee.spec.ts

@@ -74,9 +74,7 @@ tape('EIP3198 tests', (t) => {
     )
     const block = makeBlock(fee, tx, 2)
     const vm = new VM({ common })
-    const account = await vm.stateManager.getAccount(sender)
-    account.balance = ETHER
-    await vm.stateManager.putAccount(sender, account)
+    await vm.stateManager.modifyAccountFields(sender, { balance: ETHER })
 
     // Track stack
 

packages/vm/tests/api/state/stateManager.spec.ts

@@ -168,6 +168,84 @@ tape('StateManager', (t) => {
       st.end()
     }
   )
+
+  t.test('should modify account fields correctly', async (st) => {
+    const stateManager = new DefaultStateManager()
+    const account = createAccount()
+    const address = new Address(Buffer.from('a94f5374fce5edbc8e2a8697c15331677e6ebf0b', 'hex'))
+    await stateManager.putAccount(address, account)
+
+    await stateManager.modifyAccountFields(address, { balance: new BN(1234) })
+
+    const res1 = await stateManager.getAccount(address)
+
+    st.equal(res1.balance.toString('hex'), '4d2')
+
+    await stateManager.modifyAccountFields(address, { nonce: new BN(1) })
+
+    const res2 = await stateManager.getAccount(address)
+
+    st.equal(res2.nonce.toNumber(), 1)
+
+    await stateManager.modifyAccountFields(address, {
+      codeHash: Buffer.from(
+        'd748bf26ab37599c944babfdbeecf6690801bd61bf2670efb0a34adfc6dca10b',
+        'hex'
+      ),
+      stateRoot: Buffer.from(
+        'cafd881ab193703b83816c49ff6c2bf6ba6f464a1be560c42106128c8dbc35e7',
+        'hex'
+      ),
+    })
+
+    const res3 = await stateManager.getAccount(address)
+
+    st.equal(
+      res3.codeHash.toString('hex'),
+      'd748bf26ab37599c944babfdbeecf6690801bd61bf2670efb0a34adfc6dca10b'
+    )
+    st.equal(
+      res3.stateRoot.toString('hex'),
+      'cafd881ab193703b83816c49ff6c2bf6ba6f464a1be560c42106128c8dbc35e7'
+    )
+
+    st.end()
+  })
+
+  t.test(
+    'should modify account fields correctly on previously non-existent account',
+    async (st) => {
+      const stateManager = new DefaultStateManager()
+      const address = new Address(Buffer.from('a94f5374fce5edbc8e2a8697c15331677e6ebf0b', 'hex'))
+
+      await stateManager.modifyAccountFields(address, { balance: new BN(1234) })
+      const res1 = await stateManager.getAccount(address)
+      st.equal(res1.balance.toString('hex'), '4d2')
+
+      await stateManager.modifyAccountFields(address, { nonce: new BN(1) })
+      const res2 = await stateManager.getAccount(address)
+      st.equal(res2.nonce.toNumber(), 1)
+
+      const newCodeHash = Buffer.from(
+        'd748bf26ab37599c944babfdbeecf6690801bd61bf2670efb0a34adfc6dca10b',
+        'hex'
+      )
+      const newStateRoot = Buffer.from(
+        'cafd881ab193703b83816c49ff6c2bf6ba6f464a1be560c42106128c8dbc35e7',
+        'hex'
+      )
+      await stateManager.modifyAccountFields(address, {
+        codeHash: newCodeHash,
+        stateRoot: newStateRoot,
+      })
+
+      const res3 = await stateManager.getAccount(address)
+      st.ok(res3.codeHash.equals(newCodeHash))
+      st.ok(res3.stateRoot.equals(newStateRoot))
+      st.end()
+    }
+  )
+
   t.test(
     'should generate the genesis state root correctly for mainnet from ethereum/tests data',
     async (st) => {