Got some tests working for the trie libs

packages/contracts/contracts/optimistic-ethereum/OVM/verification/OVM_StateTransitioner.sol

@@ -173,7 +173,12 @@ contract OVM_StateTransitioner is iOVM_StateTransitioner, Lib_AddressResolver {
         require(
             Lib_EthMerkleTrie.proveAccountState(
                 _ovmContractAddress,
-                _account,
+                Lib_OVMCodec.EVMAccount({
+                    balance: _account.balance,
+                    nonce: _account.nonce,
+                    storageRoot: _account.storageRoot,
+                    codeHash: _account.codeHash
+                }),
                 _stateTrieWitness,
                 preStateRoot
             ),
@@ -284,7 +289,12 @@ contract OVM_StateTransitioner is iOVM_StateTransitioner, Lib_AddressResolver {
 
         postStateRoot = Lib_EthMerkleTrie.updateAccountState(
             _ovmContractAddress,
-            _account,
+            Lib_OVMCodec.EVMAccount({
+                balance: _account.balance,
+                nonce: _account.nonce,
+                storageRoot: _account.storageRoot,
+                codeHash: _account.codeHash
+            }),
             _stateTrieWitness,
             postStateRoot
         );

packages/contracts/contracts/optimistic-ethereum/libraries/codec/Lib_OVMCodec.sol

@@ -68,13 +68,6 @@ library Lib_OVMCodec {
         bytes data;
     }
 
-    struct ProofMatrix {
-        bool checkNonce;
-        bool checkBalance;
-        bool checkStorageRoot;
-        bool checkCodeHash;
-    }
-
     struct QueueElement {
         uint256 timestamp;
         bytes32 batchRoot;

packages/contracts/contracts/optimistic-ethereum/libraries/trie/Lib_MerkleTrie.sol

@@ -6,6 +6,8 @@ import { Lib_BytesUtils } from "../utils/Lib_BytesUtils.sol";
 import { Lib_RLPReader } from "../rlp/Lib_RLPReader.sol";
 import { Lib_RLPWriter } from "../rlp/Lib_RLPWriter.sol";
 
+import { console } from "@nomiclabs/buidler/console.sol";
+
 /**
  * @title Lib_MerkleTrie
  */
@@ -355,6 +357,8 @@ library Lib_MerkleTrie {
                         continue;
                     }
                 }
+            } else {
+                revert("Received an unparseable node.");
             }
         }
 

packages/contracts/contracts/test-libraries/trie/TestLib_EthMerkleTrie.sol

@@ -9,7 +9,7 @@ import { Lib_OVMCodec } from "../../optimistic-ethereum/libraries/codec/Lib_OVMC
 /**
  * @title TestLib_EthMerkleTrie
  */
-library TestLib_EthMerkleTrie {
+contract TestLib_EthMerkleTrie {
 
     function proveAccountStorageSlotValue(
         address _address,
@@ -73,78 +73,6 @@ library TestLib_EthMerkleTrie {
         );
     }
 
-    function proveAccountNonce(
-        address _address,
-        uint256 _nonce,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bool)
-    {
-        return Lib_EthMerkleTrie.proveAccountNonce(
-            _address,
-            _nonce,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
-
-    function proveAccountBalance(
-        address _address,
-        uint256 _balance,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bool)
-    {
-        return Lib_EthMerkleTrie.proveAccountBalance(
-            _address,
-            _balance,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
-
-    function proveAccountStorageRoot(
-        address _address,
-        bytes32 _storageRoot,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bool)
-    {
-        return Lib_EthMerkleTrie.proveAccountStorageRoot(
-            _address,
-            _storageRoot,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
-
-    function proveAccountCodeHash(
-        address _address,
-        bytes32 _codeHash,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bool)
-    {
-        return Lib_EthMerkleTrie.proveAccountCodeHash(
-            _address,
-            _codeHash,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
-
     function updateAccountState(
         address _address,
         Lib_OVMCodec.EVMAccount memory _accountState,
@@ -162,76 +90,4 @@ library TestLib_EthMerkleTrie {
             _stateTrieRoot
         );
     }
-
-    function updateAccountNonce(
-        address _address,
-        uint256 _nonce,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bytes32)
-    {
-        return Lib_EthMerkleTrie.updateAccountNonce(
-            _address,
-            _nonce,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
-
-    function updateAccountBalance(
-        address _address,
-        uint256 _balance,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bytes32)
-    {
-        return Lib_EthMerkleTrie.updateAccountBalance(
-            _address,
-            _balance,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
-
-    function updateAccountStorageRoot(
-        address _address,
-        bytes32 _storageRoot,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bytes32)
-    {
-        return Lib_EthMerkleTrie.updateAccountStorageRoot(
-            _address,
-            _storageRoot,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
-
-    function updateAccountCodeHash(
-        address _address,
-        bytes32 _codeHash,
-        bytes memory _stateTrieWitness,
-        bytes32 _stateTrieRoot
-    )
-        public
-        view
-        returns (bytes32)
-    {
-        return Lib_EthMerkleTrie.updateAccountCodeHash(
-            _address,
-            _codeHash,
-            _stateTrieWitness,
-            _stateTrieRoot
-        );
-    }
 }

packages/contracts/package.json

@@ -10,7 +10,7 @@
     "build:dump": "ts-node \"bin/take-dump.ts\"",
     "build:copy": "copyfiles -u 2 \"contracts/optimistic-ethereum/**/*.sol\" \"build/contracts\"",
     "test": "yarn run test:contracts",
-    "test:contracts": "buidler test --show-stack-traces",
+    "test:contracts": "buidler test \"test/contracts/libraries/trie/Lib_EthMerkleTrie.spec.ts\" --show-stack-traces",
     "lint": "yarn run lint:typescript",
     "lint:typescript": "tslint --format stylish --project .",
     "lint:fix": "yarn run lint:fix:typescript",
@@ -18,7 +18,7 @@
     "clean": "rm -rf ./artifacts ./build ./cache"
   },
   "devDependencies": {
-    "@eth-optimism/smock": "^0.1.0",
+    "@eth-optimism/smock": "^0.0.1",
     "@nomiclabs/buidler": "^1.4.4",
     "@nomiclabs/buidler-ethers": "^2.0.0",
     "@nomiclabs/buidler-waffle": "^2.0.0",
@@ -35,8 +35,12 @@
     "fs-extra": "^9.0.1",
     "ganache-core": "^2.12.1",
     "lodash": "^4.17.20",
+    "merkle-patricia-tree": "git+https://github.com/kfichter/merkle-patricia-tree",
     "mocha": "^8.1.1",
     "prettier": "^2.1.2",
+    "random-bytes-seed": "^1.0.3",
+    "rlp": "^2.2.6",
+    "seedrandom": "^3.0.5",
     "ts-node": "^9.0.0",
     "tslint": "^6.1.3",
     "tslint-config-prettier": "^1.18.0",

packages/contracts/test/contracts/libraries/trie/Lib_EthMerkleTrie.spec.ts

+import { expect } from '../../../setup'
+
+/* External Imports */
+import { ethers } from '@nomiclabs/buidler'
+import { Contract } from 'ethers'
+
+/* Internal Imports */
+import { TrieTestGenerator, NON_NULL_BYTES32 } from '../../../helpers'
+import { keccak256 } from 'ethers/lib/utils'
+
+const makeDummyAccounts = (count: number): any[] => {
+  return [...Array(count)].map((x, idx) => {
+    return {
+      address: '0xc0de' + `${idx.toString(16)}`.padStart(36, '0'),
+      nonce: 0,
+      balance: 0,
+      codeHash: null,
+      storage: [
+        {
+          key: keccak256('0x1234'),
+          val: keccak256('0x5678'),
+        },
+      ],
+    }
+  })
+}
+
+const NODE_COUNTS = [1, 2, 128, 256, 512, 1024, 2048, 4096]
+
+describe('Lib_EthMerkleTrie', () => {
+  let Lib_EthMerkleTrie: Contract
+  before(async () => {
+    Lib_EthMerkleTrie = await (
+      await ethers.getContractFactory('TestLib_EthMerkleTrie')
+    ).deploy()
+  })
+
+  describe('proveAccountStorageSlotValue', () => {})
+
+  describe('updateAccountStorageSlotValue', () => {})
+
+  describe('proveAccountState', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromAccounts({
+            accounts: makeDummyAccounts(nodeCount),
+            secure: true,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly prove inclusion for node #${i}`, async () => {
+            const test = await generator.makeAccountProofTest(i)
+
+            expect(
+              await Lib_EthMerkleTrie.proveAccountState(
+                test.address,
+                test.account,
+                test.accountTrieWitness,
+                test.accountTrieRoot
+              )
+            ).to.equal(true)
+          })
+        }
+      })
+    }
+  })
+
+  describe.only('updateAccountState', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromAccounts({
+            accounts: makeDummyAccounts(nodeCount),
+            secure: true,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly update node #${i}`, async () => {
+            const test = await generator.makeAccountUpdateTest(i, {
+              nonce: 1234,
+              balance: 5678,
+              codeHash: NON_NULL_BYTES32,
+              storageRoot: NON_NULL_BYTES32,
+            })
+
+            expect(
+              await Lib_EthMerkleTrie.updateAccountState(
+                test.address,
+                test.account,
+                test.accountTrieWitness,
+                test.accountTrieRoot
+              )
+            ).to.equal(test.newAccountTrieRoot)
+          })
+        }
+      })
+    }
+  })
+})

packages/contracts/test/contracts/libraries/trie/Lib_MerkleTrie.spec.ts

+import { expect } from '../../../setup'
+
+/* External Imports */
+import { ethers } from '@nomiclabs/buidler'
+import { Contract } from 'ethers'
+
 /* Internal Imports */
-import { Lib_MerkleTrie_TEST_JSON } from '../../../data'
-import { runJsonTest } from '../../../helpers'
+import { TrieTestGenerator } from '../../../helpers'
+
+const NODE_COUNTS = [1, 2, 128, 256, 512, 1024, 2048, 4096]
 
 describe('Lib_MerkleTrie', () => {
-  describe('JSON tests', () => {
-    runJsonTest('TestLib_MerkleTrie', Lib_MerkleTrie_TEST_JSON)
+  let Lib_MerkleTrie: Contract
+  before(async () => {
+    Lib_MerkleTrie = await (
+      await ethers.getContractFactory('TestLib_MerkleTrie')
+    ).deploy()
+  })
+
+  describe('verifyInclusionProof', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromRandom({
+            seed: `seed.incluson.${nodeCount}`,
+            nodeCount,
+            secure: false,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly prove inclusion for node #${i}`, async () => {
+            const test = await generator.makeInclusionProofTest(i)
+
+            expect(
+              await Lib_MerkleTrie.verifyInclusionProof(
+                test.key,
+                test.val,
+                test.proof,
+                test.root
+              )
+            ).to.equal(true)
+          })
+        }
+      })
+    }
+  })
+
+  describe('update', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromRandom({
+            seed: `seed.update.${nodeCount}`,
+            nodeCount,
+            secure: false,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly update node #${i}`, async () => {
+            const test = await generator.makeNodeUpdateTest(
+              i,
+              '0x1234123412341234'
+            )
+
+            expect(
+              await Lib_MerkleTrie.update(
+                test.key,
+                test.val,
+                test.proof,
+                test.root
+              )
+            ).to.equal(test.newRoot)
+          })
+        }
+      })
+    }
+  })
+
+  describe('get', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromRandom({
+            seed: `seed.get.${nodeCount}`,
+            nodeCount,
+            secure: false,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly get the value of node #${i}`, async () => {
+            const test = await generator.makeInclusionProofTest(i)
+
+            expect(
+              await Lib_MerkleTrie.get(test.key, test.proof, test.root)
+            ).to.deep.equal([true, test.val])
+          })
+        }
+      })
+    }
   })
 })

packages/contracts/test/contracts/libraries/trie/Lib_SecureMerkleTrie.spec.ts

+import { expect } from '../../../setup'
+
+/* External Imports */
+import { ethers } from '@nomiclabs/buidler'
+import { Contract } from 'ethers'
+
+/* Internal Imports */
+import { TrieTestGenerator } from '../../../helpers'
+
+const NODE_COUNTS = [1, 2, 128, 256, 512, 1024, 2048, 4096]
+
+describe('Lib_SecureMerkleTrie', () => {
+  let Lib_SecureMerkleTrie: Contract
+  before(async () => {
+    Lib_SecureMerkleTrie = await (
+      await ethers.getContractFactory('TestLib_SecureMerkleTrie')
+    ).deploy()
+  })
+
+  describe('verifyInclusionProof', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromRandom({
+            seed: `seed.incluson.${nodeCount}`,
+            nodeCount,
+            secure: true,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly prove inclusion for node #${i}`, async () => {
+            const test = await generator.makeInclusionProofTest(i)
+
+            expect(
+              await Lib_SecureMerkleTrie.verifyInclusionProof(
+                test.key,
+                test.val,
+                test.proof,
+                test.root
+              )
+            ).to.equal(true)
+          })
+        }
+      })
+    }
+  })
+
+  describe('update', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromRandom({
+            seed: `seed.update.${nodeCount}`,
+            nodeCount,
+            secure: true,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly update node #${i}`, async () => {
+            const test = await generator.makeNodeUpdateTest(
+              i,
+              '0x1234123412341234'
+            )
+
+            expect(
+              await Lib_SecureMerkleTrie.update(
+                test.key,
+                test.val,
+                test.proof,
+                test.root
+              )
+            ).to.equal(test.newRoot)
+          })
+        }
+      })
+    }
+  })
+
+  describe('get', () => {
+    for (const nodeCount of NODE_COUNTS) {
+      describe(`inside a trie with ${nodeCount} nodes`, () => {
+        let generator: TrieTestGenerator
+        before(async () => {
+          generator = await TrieTestGenerator.fromRandom({
+            seed: `seed.get.${nodeCount}`,
+            nodeCount,
+            secure: true,
+          })
+        })
+
+        for (
+          let i = 0;
+          i < nodeCount;
+          i += nodeCount / (nodeCount > 8 ? 8 : 1)
+        ) {
+          it(`should correctly get the value of node #${i}`, async () => {
+            const test = await generator.makeInclusionProofTest(i)
+
+            expect(
+              await Lib_SecureMerkleTrie.get(test.key, test.proof, test.root)
+            ).to.deep.equal([true, test.val])
+          })
+        }
+      })
+    }
+  })
+})

packages/contracts/test/helpers/index.ts

@@ -5,3 +5,4 @@ export * from './resolver'
 export * from './utils'
 export * from './codec'
 export * from './test-runner'
+export * from './trie'

packages/contracts/test/helpers/trie/index.ts

+export * from './trie-test-generator'

packages/contracts/test/helpers/trie/trie-test-generator.ts

+/* External Imports */
+import * as rlp from 'rlp'
+import { default as seedbytes } from 'random-bytes-seed'
+import { SecureTrie, BaseTrie } from 'merkle-patricia-tree'
+
+/* Internal Imports */
+import { fromHexString, toHexString } from '../utils'
+import { NULL_BYTES32 } from '../constants'
+
+export interface TrieNode {
+  key: string
+  val: string
+}
+
+export interface InclusionProofTest {
+  key: string
+  val: string
+  proof: string
+  root: string
+}
+
+export interface NodeUpdateTest extends InclusionProofTest {
+  newRoot: string
+}
+
+export interface EthereumAccount {
+  address?: string
+  nonce: number
+  balance: number
+  codeHash: string
+  storageRoot?: string
+  storage?: TrieNode[]
+}
+
+export interface AccountProofTest {
+  address: string
+  account: EthereumAccount
+  accountTrieWitness: string
+  accountTrieRoot: string
+}
+
+export interface AccountUpdateTest extends AccountProofTest {
+  newAccountTrieRoot: string
+}
+
+const rlpEncodeAccount = (account: EthereumAccount): string => {
+  return toHexString(
+    rlp.encode([
+      account.nonce,
+      account.balance,
+      account.codeHash || NULL_BYTES32,
+      account.storageRoot || NULL_BYTES32,
+    ])
+  )
+}
+
+const rlpDecodeAccount = (encoded: string): EthereumAccount => {
+  const decoded = rlp.decode(fromHexString(encoded)) as any
+  return {
+    nonce: decoded[0].length ? parseInt(decoded[0], 16) : 0,
+    balance: decoded[1].length ? parseInt(decoded[1], 16) : 0,
+    storageRoot: decoded[2].length ? toHexString(decoded[2]) : NULL_BYTES32,
+    codeHash: decoded[3].length ? toHexString(decoded[3]) : NULL_BYTES32,
+  }
+}
+
+const makeTrie = async (
+  nodes: TrieNode[],
+  secure?: boolean
+): Promise<{
+  trie: SecureTrie | BaseTrie
+  TrieClass: any
+}> => {
+  const TrieClass = secure ? SecureTrie : BaseTrie
+  const trie = new TrieClass()
+
+  for (const node of nodes) {
+    await trie.put(fromHexString(node.key), fromHexString(node.val))
+  }
+
+  return {
+    trie,
+    TrieClass,
+  }
+}
+
+export class TrieTestGenerator {
+  constructor(
+    public _TrieClass: any,
+    public _trie: SecureTrie | BaseTrie,
+    public _nodes: TrieNode[],
+    public _subGenerators?: TrieTestGenerator[]
+  ) {}
+
+  static async fromNodes(opts: {
+    nodes: TrieNode[]
+    secure?: boolean
+  }): Promise<TrieTestGenerator> {
+    const { trie, TrieClass } = await makeTrie(opts.nodes, opts.secure)
+
+    return new TrieTestGenerator(TrieClass, trie, opts.nodes)
+  }
+
+  static async fromRandom(opts: {
+    seed: string
+    nodeCount: number
+    secure?: boolean
+    keySize?: number
+    valSize?: number
+  }): Promise<TrieTestGenerator> {
+    const getRandomBytes = seedbytes(opts.seed)
+    const nodes: TrieNode[] = [...Array(opts.nodeCount)].map(() => {
+      return {
+        key: toHexString(getRandomBytes(opts.keySize || 32)),
+        val: toHexString(getRandomBytes(opts.valSize || 32)),
+      }
+    })
+
+    return TrieTestGenerator.fromNodes({
+      nodes,
+      secure: opts.secure,
+    })
+  }
+
+  static async fromAccounts(opts: {
+    accounts: EthereumAccount[]
+    secure?: boolean
+  }): Promise<TrieTestGenerator> {
+    const subGenerators: TrieTestGenerator[] = []
+
+    for (const account of opts.accounts) {
+      if (account.storage) {
+        const subGenerator = await TrieTestGenerator.fromNodes({
+          nodes: account.storage,
+          secure: opts.secure,
+        })
+
+        account.storageRoot = toHexString(subGenerator._trie.root)
+        subGenerators.push(subGenerator)
+      }
+    }
+
+    const nodes = opts.accounts.map((account) => {
+      return {
+        key: account.address,
+        val: rlpEncodeAccount(account),
+      }
+    })
+
+    const { trie, TrieClass } = await makeTrie(nodes, opts.secure)
+
+    return new TrieTestGenerator(TrieClass, trie, nodes, subGenerators)
+  }
+
+  public async makeInclusionProofTest(
+    key: string | number
+  ): Promise<InclusionProofTest> {
+    if (typeof key === 'number') {
+      key = this._nodes[key].key
+    }
+
+    const trie = this._trie.copy()
+
+    const proof = await this.prove(key)
+    const val = await trie.get(fromHexString(key))
+
+    return {
+      proof: toHexString(rlp.encode(proof)),
+      key: toHexString(key),
+      val: toHexString(val),
+      root: toHexString(trie.root),
+    }
+  }
+
+  public async makeAllInclusionProofTests(): Promise<InclusionProofTest[]> {
+    return Promise.all(
+      this._nodes.map(async (node) => {
+        return this.makeInclusionProofTest(node.key)
+      })
+    )
+  }
+
+  public async makeNodeUpdateTest(
+    key: string | number,
+    val: string
+  ): Promise<NodeUpdateTest> {
+    if (typeof key === 'number') {
+      key = this._nodes[key].key
+    }
+
+    const trie = this._trie.copy()
+
+    const proof = await this.prove(key)
+    const oldRoot = trie.root
+
+    await trie.put(fromHexString(key), fromHexString(val))
+    const newRoot = trie.root
+
+    return {
+      proof: toHexString(rlp.encode(proof)),
+      key: toHexString(key),
+      val: toHexString(val),
+      root: toHexString(oldRoot),
+      newRoot: toHexString(newRoot),
+    }
+  }
+
+  public async makeAccountProofTest(
+    address: string | number
+  ): Promise<AccountProofTest> {
+    if (typeof address === 'number') {
+      address = this._nodes[address].key
+    }
+
+    const trie = this._trie.copy()
+
+    const proof = await this.prove(address)
+    const account = await trie.get(fromHexString(address))
+
+    return {
+      address,
+      account: rlpDecodeAccount(toHexString(account)),
+      accountTrieWitness: toHexString(rlp.encode(proof)),
+      accountTrieRoot: toHexString(trie.root),
+    }
+  }
+
+  public async makeAccountUpdateTest(
+    address: string | number,
+    account: EthereumAccount
+  ): Promise<AccountUpdateTest> {
+    if (typeof address === 'number') {
+      address = this._nodes[address].key
+    }
+
+    const trie = this._trie.copy()
+
+    const proof = await this.prove(address)
+    const oldRoot = trie.root
+
+    await trie.put(
+      fromHexString(address),
+      fromHexString(rlpEncodeAccount(account))
+    )
+    const newRoot = trie.root
+
+    return {
+      address,
+      account,
+      accountTrieWitness: toHexString(rlp.encode(proof)),
+      accountTrieRoot: toHexString(oldRoot),
+      newAccountTrieRoot: toHexString(newRoot),
+    }
+  }
+
+  private async prove(key: string): Promise<any> {
+    return this._TrieClass.prove(this._trie, fromHexString(key))
+  }
+}