From 798de46ca746c7486ae890d5b28a9ebbcd505bb7 Mon Sep 17 00:00:00 2001
From: andri lim <jangko128@gmail.com>
Date: Thu, 7 May 2020 22:12:34 +0700
Subject: [PATCH] fix json_from_tree

---
 stateless/json_from_tree.nim   | 97 ++++++++++++++++++----------------
 stateless/json_witness_gen.nim |  2 +-
 2 files changed, 51 insertions(+), 48 deletions(-)

diff --git a/stateless/json_from_tree.nim b/stateless/json_from_tree.nim
index 89d6ac073..5ca41d98c 100644
--- a/stateless/json_from_tree.nim
+++ b/stateless/json_from_tree.nim
@@ -215,71 +215,74 @@ proc getBranchRecurse(wb: var WitnessBuilder, z: var StackElem) =
   case nodeRlp.listLen
   of 2:
     let (isLeaf, k) = nodeRlp.extensionNodeKey
-    var match = false
-    # only zero or one group can match the path
-    # but if there is a match, it can be in any position
-    # 1st, 2nd, or max 3rd position
-    # recursion will go deeper depend on the common-prefix length nibbles
-    for mg in groups(z.keys, z.depth, k, z.parentGroup):
-      if not mg.match: continue
-      doAssert(match == false) # should be only one match
-      match = true
-      let value = nodeRlp.listElem(1)
-      if not isLeaf:
-        # ExtensionNodeType
-        writeExtensionNode(wb, k, z.depth, z.node)
-        var zz = StackElem(
-          node: value.getNode,
-          parentGroup: mg.group,
-          keys: z.keys,
-          depth: z.depth + k.len,
-          storageMode: z.storageMode
-        )
-        getBranchRecurse(wb, zz)
-      else:
-        # this should be only one match
-        # if there is more than one match
-        # it means we encounter an invalid address
-        for kd in keyDatas(z.keys, mg.group):
-          if not match(kd, k, z.depth): continue # skip the invalid address
-          kd.visited = true
-          if z.storageMode:
-            doAssert(kd.storageMode)
-            writeAccountStorageLeafNode(wb, kd.storageSlot, value.toBytes.decode(UInt256), k, z.node, z.depth)
-          else:
-            doAssert(not kd.storageMode)
-            writeAccountNode(wb, kd, value.toBytes.decode(Account), k, z.node, z.depth)
-    if not match:
+    let mg = groups(z.keys, z.depth, k, z.parentGroup)
+
+    if not mg.match:
+      # return immediately if there is no match
       writeHashNode(wb, keccak(z.node).data)
+      return
+
+    let value = nodeRlp.listElem(1)
+    if not isLeaf:
+      # recursion will go deeper depend on the common-prefix length nibbles
+      writeExtensionNode(wb, k, z.depth, z.node)
+      var zz = StackElem(
+        node: value.getNode,
+        parentGroup: mg.group,
+        keys: z.keys,
+        depth: z.depth + k.len, # increase the depth by k.len
+        storageMode: z.storageMode
+      )
+      getBranchRecurse(wb, zz)
+      return
+
+    # there should be only one match
+    let kd = z.keys.visitMatch(mg, z.depth, k)
+    if z.storageMode:
+      doAssert(kd.storageMode)
+      writeAccountStorageLeafNode(wb, kd.storageSlot, value.toBytes.decode(UInt256), k, z.node, z.depth)
+    else:
+      doAssert(not kd.storageMode)
+      writeAccountNode(wb, kd, value.toBytes.decode(Account), k, z.node, z.depth)
+
   of 17:
     let branchMask = rlpListToBitmask(nodeRlp)
     writeBranchNode(wb, branchMask, z.depth, z.node)
-    let path = groups(z.keys, z.parentGroup, z.depth)
 
     # if there is a match in any branch elem
     # 1st to 16th, the recursion will go deeper
     # by one nibble
-    let notLeaf = z.depth != 63 # path.len == 0
-    for i in 0..<16:
-      if not branchMask.branchMaskBitIsSet(i): continue
-      var branch = nodeRlp.listElem(i)
-      if notLeaf and branchMaskBitIsSet(path.mask, i):
+    doAssert(z.depth != 64) # notLeaf or path.len == 0
+
+    let path = groups(z.keys, z.parentGroup, z.depth)
+    for i in nonEmpty(branchMask):
+      let branch = nodeRlp.listElem(i)
+      if branchMaskBitIsSet(path.mask, i):
+        # it is a match between multikeys and Branch Node elem
         var zz = StackElem(
           node: branch.getNode,
           parentGroup: path.groups[i],
           keys: z.keys,
-          depth: z.depth + 1,
+          depth: z.depth + 1, # increase the depth by one
           storageMode: z.storageMode
         )
         getBranchRecurse(wb, zz)
+        continue
+
+      if branch.isList:
+        # short node appear in yellow paper
+        # but never in the actual ethereum state trie
+        # an rlp encoded ethereum account will have length > 32 bytes
+        # block witness spec silent about this
+        doAssert(false, "Short node should not exist in block witness")
       else:
-        if branch.isList:
-          doAssert(false, "Short node should not exist in block witness")
-        else:
-          # this is a potential branch for multiproof
-          writeHashNode(wb, branch.expectHash)
+        # if branch elem not empty and not a match, emit hash
+        writeHashNode(wb, branch.expectHash)
 
     # 17th elem should always empty
+    # 17th elem appear in yellow paper but never in
+    # the actual ethereum state trie
+    # the 17th elem also not included in block witness spec
     doAssert branchMask.branchMaskBitIsSet(16) == false
   else:
     raise newException(CorruptedTrieDatabase,
diff --git a/stateless/json_witness_gen.nim b/stateless/json_witness_gen.nim
index 0449ae61d..82411d3fd 100644
--- a/stateless/json_witness_gen.nim
+++ b/stateless/json_witness_gen.nim
@@ -82,7 +82,7 @@ proc runGenerator(numPairs, numSlots: int): string =
   result = wb.buildWitness(mkeys)
 
 proc writeHelp() =
-  echo "jsoin_witness_gen output --pairs:val --slots:val -s:val -p:val"
+  echo "json_witness_gen output --pairs:val --slots:val -s:val -p:val"
 
 proc main() =
   var filename: string