Use hash writer

execution_chain/db/aristo/aristo_compute.nim

@@ -89,24 +89,70 @@ proc putKeyAtLevel(
 
   ok()
 
-template encodeLeaf(w: var RlpWriter, pfx: NibblesBuf, leafData: untyped): HashKey =
+template appendLeaf(w: var RlpWriter, pfx: NibblesBuf, leafData: untyped) =
   w.startList(2)
   w.append(pfx.toHexPrefix(isLeaf = true).data())
+  w.wrapEncoding(1)
   w.append(leafData)
-  w.finish().digestTo(HashKey)
 
-template encodeBranch(w: var RlpWriter, vtx: VertexRef, subKeyForN: untyped): HashKey =
+template encodeLeaf(pfx: NibblesBuf, leafData: untyped): HashKey =
+  var tracker = DynamicRlpLengthTracker()
+  tracker.initLengthTracker()
+  tracker.appendLeaf(pfx, leafData)
+
+  if tracker.totalLength < 32:
+    var writer = initTwoPassWriter(tracker)
+    writer.appendLeaf(pfx, leafData)
+    let buf = HashKey.fromBytes(writer.finish)
+    buf.value
+  else:
+    var writer = initHashWriter(tracker)
+    writer.appendLeaf(pfx, leafData)
+    let buf = writer.finish()
+    buf.to(HashKey)
+
+template appendBranch(w: var RlpWriter, vtx: VertexRef, subKeyForN: untyped) =
   w.startList(17)
   for (n {.inject.}, subvid {.inject.}) in vtx.allPairs():
     w.append(subKeyForN)
   w.append EmptyBlob
-  w.finish().digestTo(HashKey)
 
-template encodeExt(w: var RlpWriter, pfx: NibblesBuf, branchKey: HashKey): HashKey =
+template encodeBranch(vtx: VertexRef, subKeyForN: untyped): HashKey =
+  var tracker: DynamicRlpLengthTracker
+  tracker.initLengthTracker()
+  tracker.appendBranch(vtx, subKeyForN)
+
+  if tracker.totalLength < 32:
+    var writer = initTwoPassWriter(tracker)
+    writer.appendBranch(vtx, subKeyForN)
+    let buf = HashKey.fromBytes(writer.finish)
+    buf.value
+  else:
+    var writer = initHashWriter(tracker)
+    writer.appendBranch(vtx, subKeyForN)
+    let buf = writer.finish()
+    buf.to(HashKey)
+
+template appendExt(w: var RlpWriter, pfx: NibblesBuf, branchKey: HashKey) =
   w.startList(2)
   w.append(pfx.toHexPrefix(isLeaf = false).data())
   w.append(branchKey)
-  w.finish().digestTo(HashKey)
+
+template encodeExt(pfx: NibblesBuf, branchKey: untyped): HashKey =
+  var tracker: DynamicRlpLengthTracker
+  tracker.initLengthTracker()
+  tracker.appendExt(pfx, branchKey)
+
+  if tracker.totalLength < 32:
+    var writer = initTwoPassWriter(tracker)
+    writer.appendExt(pfx, branchKey)
+    let buf = HashKey.fromBytes(writer.finish)
+    buf.value
+  else:
+    var writer = initHashWriter(tracker)
+    writer.appendExt(pfx, branchKey)
+    let buf = writer.finish()
+    buf.to(HashKey)
 
 proc getKey(
     db: AristoTxRef, rvid: RootedVertexID, skipLayers: static bool
@@ -146,14 +192,11 @@ proc computeKeyImpl(
   # Top-most level of all the verticies this hash computation depends on
   var level = level
 
-  # TODO this is the same code as when serializing NodeRef, without the NodeRef
-  var writer = initRlpWriter()
-
   let key =
     case vtx.vType
     of AccLeaf:
       let vtx = AccLeafRef(vtx)
-      writer.encodeLeaf(vtx.pfx):
+      encodeLeaf(vtx.pfx):
         let
           stoID = vtx.stoID
           skey =
@@ -176,17 +219,16 @@ proc computeKeyImpl(
             else:
               VOID_HASH_KEY
 
-        rlp.encode Account(
+        Account(
           nonce: vtx.account.nonce,
           balance: vtx.account.balance,
           storageRoot: skey.to(Hash32),
           codeHash: vtx.account.codeHash,
         )
     of StoLeaf:
       let vtx = StoLeafRef(vtx)
-      writer.encodeLeaf(vtx.pfx):
-        # TODO avoid memory allocation when encoding storage data
-        rlp.encode(vtx.stoData)
+      encodeLeaf(vtx.pfx):
+        vtx.stoData
     of Branches:
       # For branches, we need to load the vertices before recursing into them
       # to exploit their on-disk order
@@ -244,22 +286,24 @@ proc computeKeyImpl(
               skipLayers = skipLayers,
             )
           batch.leave(n)
-
-      template writeBranch(w: var RlpWriter, vtx: BranchRef): HashKey =
-        w.encodeBranch(vtx):
+ 
+      template writeBranch(vtx: BranchRef): HashKey =
+        encodeBranch(vtx):
           if subvid.isValid:
             level = max(level, keyvtxs[n][1])
             keyvtxs[n][0][0]
           else:
             VOID_HASH_KEY
 
+      var encoding: HashKey
       if vtx.vType == ExtBranch:
         let vtx = ExtBranchRef(vtx)
-        writer.encodeExt(vtx.pfx):
-          var bwriter = initRlpWriter()
-          bwriter.writeBranch(vtx)
+        encoding = encodeExt(vtx.pfx):
+          writeBranch(vtx)
       else:
-        writer.writeBranch(vtx)
+        encoding = writeBranch(vtx)
+
+      encoding
 
   # Cache the hash into the same storage layer as the the top-most value that it
   # depends on (recursively) - this could be an ephemeral in-memory layer or the