From 4de8eaa81d19620bc4dc195365ae5c40b51fd32b Mon Sep 17 00:00:00 2001
From: andri lim <jangko128@gmail.com>
Date: Tue, 28 Apr 2020 16:58:16 +0700
Subject: [PATCH] implement account code encoding and decoding

---
 stateless/test_witness_keys.nim | 58 ++++++++++++++++-----------------
 stateless/tree_from_witness.nim | 12 +++++--
 stateless/witness_from_tree.nim |  6 ++--
 3 files changed, 43 insertions(+), 33 deletions(-)

diff --git a/stateless/test_witness_keys.nim b/stateless/test_witness_keys.nim
index e3c201c3f..904194fcb 100644
--- a/stateless/test_witness_keys.nim
+++ b/stateless/test_witness_keys.nim
@@ -1,9 +1,12 @@
 import
   randutils, stew/byteutils, random,
-  eth/[common, rlp], eth/trie/[hexary, db],
+  eth/[common, rlp], eth/trie/[hexary, db, trie_defs],
   faststreams/input_stream, nimcrypto/[utils, sysrand],
   ../stateless/[witness_from_tree, tree_from_witness]
 
+type
+   DB = TrieDatabaseRef
+
 proc randU256(): UInt256 =
   var bytes: array[32, byte]
   discard randomBytes(bytes[0].addr, sizeof(result))
@@ -12,56 +15,53 @@ proc randU256(): UInt256 =
 proc randNonce(): AccountNonce =
   discard randomBytes(result.addr, sizeof(result))
 
+proc randCode(db: DB): Hash256 =
+  if rand(0..1) == 0:
+    result = blankStringHash
+  else:
+    let codeLen = rand(1..150)
+    let code = randList(byte, rng(0, 255), codeLen, unique = false)
+    result = hexary.keccak(code)
+    db.put(result.data, code)
+
 proc randHash(): Hash256 =
   discard randomBytes(result.data[0].addr, sizeof(result))
 
-proc randAccount(): Account =
+proc randAccount(db: DB): Account =
   result.nonce = randNonce()
   result.balance = randU256()
-  result.codeHash = randHash()
+  result.codeHash = randCode(db)
   result.storageRoot = randHash()
 
-proc runTest(keyBytes: int, valBytes: int, numPairs: int) =
+proc runTest(numPairs: int) =
   var memDB = newMemoryDB()
-  var trie = initHexaryTrie(memDB)
-
-  var
-    keys = newSeq[Bytes](numPairs)
-    vals = newSeq[Bytes](numPairs)
+  var trie = initSecureHexaryTrie(memDB)
+  var addrs = newSeq[Bytes](numPairs)
 
   for i in 0..<numPairs:
-    keys[i] = randList(byte, rng(0, 255), keyBytes, unique = false)
-    vals[i] = randList(byte, rng(0, 255), valBytes, unique = false)
-    trie.put(keys[i], vals[i])
+    addrs[i] = randList(byte, rng(0, 255), 20, unique = false)
+    let acc = randAccount(memDB)
+    trie.put(addrs[i], rlp.encode(acc))
 
   let rootHash = trie.rootHash
 
   var wb = initWitnessBuilder(memDB, rootHash)
-  var witness = wb.getBranchRecurse(keys[0])
-  var input = memoryInput(witness)
-
+  var witness = wb.getBranchRecurse(addrs[0])
   var db = newMemoryDB()
-  var tb = initTreeBuilder(input, db)
+  when defined(useInputStream):
+    var input = memoryInput(witness)
+    var tb = initTreeBuilder(input, db)
+  else:
+    var tb = initTreeBuilder(witness, db)
   var root = tb.treeNode()
   debugEcho "root: ", root.data.toHex
   debugEcho "rootHash: ", rootHash.data.toHex
   doAssert root.data == rootHash.data
 
 proc main() =
-  runTest(7, 100, 50)
-  runTest(1, 1, 1)
-  runTest(5, 5, 1)
-  runTest(6, 7, 3)
-  runTest(7, 10, 7)
-  runTest(8, 15, 11)
-  runTest(9, 30, 13)
-  runTest(11, 40, 10)
-  runTest(20, 1, 15)
-  runTest(25, 10, 20)
-
   randomize()
-  for i in 0..<30:
-    runTest(rand(1..30), rand(1..50), rand(1..30))
 
+  for i in 0..<30:
+    runTest(rand(1..30))
 
 main()
diff --git a/stateless/tree_from_witness.nim b/stateless/tree_from_witness.nim
index b8345e168..0be37028d 100644
--- a/stateless/tree_from_witness.nim
+++ b/stateless/tree_from_witness.nim
@@ -213,8 +213,10 @@ proc extensionNode(t: var TreeBuilder, depth: int): NodeKey =
 
 proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 65)
-  let len = t.readU32().int
-  result = toNodeKey(t.read(len))
+
+  when defined(debugHash):
+    let len = t.readU32().int
+    let nodeKey = toNodeKey(t.read(len))
 
   when defined(debugDepth):
     let readDepth = t.readByte.int
@@ -246,6 +248,12 @@ proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
     let storageRoot = t.treeNode(0, accountMode = true)
     doAssert(storageRoot.usedBytes == 32)
     acc.storageRoot.data = storageRoot.data
+    
+  r.append rlp.encode(acc)
+  result = toNodeKey(r.finish)    
+  
+  when defined(debugHash):
+    doAssert(result == nodeKey, "account node parsing error")
 
 proc accountStorageLeafNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 65)
diff --git a/stateless/witness_from_tree.nim b/stateless/witness_from_tree.nim
index d441f9f49..3b440e8ea 100644
--- a/stateless/witness_from_tree.nim
+++ b/stateless/witness_from_tree.nim
@@ -95,8 +95,10 @@ proc writeHashNode(wb: var WitnessBuilder, node: openArray[byte]) =
 proc writeAccountNode(wb: var WitnessBuilder, acc: Account, nibbles: NibblesSeq, node: openArray[byte], depth: int) =
   # write type
   wb.output.append(AccountNodeType.byte)
-  wb.writeU32(node.len.uint32)
-  wb.output.append(node)
+
+  when defined(debugHash):
+    wb.writeU32(node.len.uint32)
+    wb.output.append(node)
 
   when defined(debugDepth):
     wb.output.append(depth.byte)