Clear stateManager cache after setStateRoot

lib/stateManager.js

@@ -396,11 +396,17 @@ proto.setStateRoot = function (stateRoot, cb) {
 
   self._cache.flush(function (err) {
     if (err) { return cb(err) }
+    if (stateRoot === self._trie.EMPTY_TRIE_ROOT) {
+      self._trie.root = stateRoot
+      self._cache.clear()
+      return cb()
+    }
     self._trie.checkRoot(stateRoot, function (err, hasRoot) {
       if (err || !hasRoot) {
         cb(err || new Error('State trie does not contain state root'))
       } else {
         self._trie.root = stateRoot
+        self._cache.clear()
         cb()
       }
     })

tests/api/runBlock.js

@@ -64,7 +64,7 @@ tape('runBlock', async (t) => {
 
     // The mocked VM uses a mocked runTx
     // which always returns an error.
-    await suite.p.runBlock({ block, root: genesis.header.stateRoot, skipBlockValidation: true })
+    await suite.p.runBlock({ block, skipBlockValidation: true })
       .then(() => t.fail('should have returned error'))
       .catch((e) => t.equal(e.message, 'test'))
 
@@ -112,7 +112,7 @@ tape('should fail when tx gas limit higher than block gas limit', async (t) => {
 
   await suite.p.generateCanonicalGenesis()
 
-  await suite.p.runBlock({ block, root: genesis.header.stateRoot, skipBlockValidation: true })
+  await suite.p.runBlock({ block, skipBlockValidation: true })
     .then(() => t.fail('should have returned error'))
     .catch((e) => t.ok(e.message.includes('higher gas limit')))
 
@@ -137,7 +137,7 @@ tape('should fail when runCall fails', async (t) => {
   // runTx is a full implementation that works.
   suite.vm.runTx = runTx
 
-  await suite.p.runBlock({ block, root: suite.vm.stateManager._trie.root, skipBlockValidation: true })
+  await suite.p.runBlock({ block, skipBlockValidation: true })
 
     .then(() => t.fail('should have returned error'))
     .catch((e) => t.equal(e.message, 'test'))

tests/api/stateManager.js

@@ -18,6 +18,36 @@ tape('StateManager', (t) => {
     st.end()
   })
 
+  t.test('should clear the cache when the state root is set', async (st) => {
+    const stateManager = new StateManager()
+    const addressBuffer = Buffer.from('a94f5374fce5edbc8e2a8697c15331677e6ebf0b', "hex")
+    const account = createAccount()
+
+    const getStateRoot = promisify((...args) => stateManager.getStateRoot(...args))
+    const checkpoint = promisify((...args) => stateManager.checkpoint(...args))
+    const putAccount = promisify((...args) => stateManager.putAccount(...args))
+    const getAccount = promisify((...args) => stateManager.getAccount(...args))
+    const commit = promisify((...args) => stateManager.commit(...args))
+    const setStateRoot = promisify((...args) => stateManager.setStateRoot(...args))
+
+    const initialStateRoot = await getStateRoot()
+    await checkpoint()
+    await putAccount(addressBuffer, account)
+
+    const account0 = await getAccount(addressBuffer)
+    st.equal(account0.balance.toString("hex"), account.balance.toString("hex"), 'account value is set in the cache')
+
+    await commit()
+    const account1 = await getAccount(addressBuffer)
+    st.equal(account1.balance.toString("hex"), account.balance.toString("hex"), 'account value is set in the state trie')
+
+    await setStateRoot(initialStateRoot)
+    const account2 = await getAccount(addressBuffer)
+    st.equal(account2.balance.toString("hex"), "", 'account value is set to 0 in original state root')
+
+    st.end()
+  })
+
   t.test('should put and get account, and add to the underlying cache if the account is not found', async (st) => {
     const stateManager = new StateManager()
     const account = createAccount()