From 4974ca49dbf2277c62e222097b753c9ab4f28d64 Mon Sep 17 00:00:00 2001
From: andri lim <jangko128@gmail.com>
Date: Wed, 29 Apr 2020 11:12:15 +0700
Subject: [PATCH] fix buildWitness API

---
 stateless/test_block_witness.nim | 2 +-
 stateless/test_witness_keys.nim  | 9 ++++++---
 stateless/witness_from_tree.nim  | 5 +++--
 stateless/witness_types.nim      | 3 ++-
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/stateless/test_block_witness.nim b/stateless/test_block_witness.nim
index 546e6ab59..9830120d3 100644
--- a/stateless/test_block_witness.nim
+++ b/stateless/test_block_witness.nim
@@ -30,7 +30,7 @@ proc testGetBranch(tester: Tester, rootHash: KeccakHash, testStatusIMPL: var Tes
   #try:
   for address in tester.address:
     var wb = initWitnessBuilder(tester.memdb, rootHash)
-    var witness = wb.getBranchRecurse(address)
+    var witness = wb.buildWitness(address)
 
     var db = newMemoryDB()
     when defined(useInputStream):
diff --git a/stateless/test_witness_keys.nim b/stateless/test_witness_keys.nim
index d3b48eb79..5b01e0770 100644
--- a/stateless/test_witness_keys.nim
+++ b/stateless/test_witness_keys.nim
@@ -43,20 +43,23 @@ proc randAccount(db: DB): Account =
   result.codeHash = randCode(db)
   result.storageRoot = randStorage(db)
 
+proc randAddress(): EthAddress =
+  discard randomBytes(result.addr, sizeof(result))
+
 proc runTest(numPairs: int) =
   var memDB = newMemoryDB()
   var trie = initSecureHexaryTrie(memDB)
-  var addrs = newSeq[Bytes](numPairs)
+  var addrs = newSeq[EthAddress](numPairs)
 
   for i in 0..<numPairs:
-    addrs[i] = randList(byte, rng(0, 255), 20, unique = false)
+    addrs[i] = randAddress()
     let acc = randAccount(memDB)
     trie.put(addrs[i], rlp.encode(acc))
 
   let rootHash = trie.rootHash
 
   var wb = initWitnessBuilder(memDB, rootHash)
-  var witness = wb.getBranchRecurse(hexary.keccak(addrs[0]).data)
+  var witness = wb.buildWitness(addrs[0])
   var db = newMemoryDB()
   when defined(useInputStream):
     var input = memoryInput(witness)
diff --git a/stateless/witness_from_tree.nim b/stateless/witness_from_tree.nim
index 88a3b3d05..7cb9555fc 100644
--- a/stateless/witness_from_tree.nim
+++ b/stateless/witness_from_tree.nim
@@ -241,7 +241,8 @@ proc getBranchRecurseAux(wb: var WitnessBuilder, node: openArray[byte], path: Ni
     raise newException(CorruptedTrieDatabase,
                        "HexaryTrie node with an unexpected number of children")
 
-proc getBranchRecurse*(wb: var WitnessBuilder; key: openArray[byte]): seq[byte] =
+proc buildWitness*(wb: var WitnessBuilder; address: EthAddress): seq[byte] =
+  let key = keccak(address)
   var node = wb.db.get(wb.root.data)
-  getBranchRecurseAux(wb, node, initNibbleRange(key), 0, false)
+  getBranchRecurseAux(wb, node, initNibbleRange(key.data), 0, false)
   result = wb.output.getOutput(seq[byte])
diff --git a/stateless/witness_types.nim b/stateless/witness_types.nim
index d6da55108..4f8c824bb 100644
--- a/stateless/witness_types.nim
+++ b/stateless/witness_types.nim
@@ -13,7 +13,8 @@ type
 
 const
   StorageLeafNodeType* = AccountNodeType
-  
+  BlockWitnessVersion* = 0x01
+
 proc setBranchMaskBit*(x: var uint, i: int) {.inline.} =
   assert(i >= 0 and i < 17)
   x = x or (1 shl i).uint