Coders for merkle trees

codex/merkletree.nim

@@ -0,0 +1,4 @@
+import ./merkletree/merkletree
+import ./merkletree/coders
+
+export merkletree, coders

codex/merkletree/coders.nim

@@ -0,0 +1,46 @@
+## Nim-Codex
+## Copyright (c) 2023 Status Research & Development GmbH
+## Licensed under either of
+##  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
+##  * MIT license ([LICENSE-MIT](LICENSE-MIT))
+## at your option.
+## This file may not be copied, modified, or distributed except according to
+## those terms.
+
+import pkg/libp2p
+import pkg/questionable
+import pkg/questionable/results
+
+import ./merkletree
+import ../units
+import ../errors
+
+const MaxMerkleTreeSize = 100.MiBs.uint
+
+proc encode*(self: MerkleTree): seq[byte] =
+  var pb = initProtoBuffer(maxSize = MaxMerkleTreeSize)
+  pb.write(1, self.mcodec.uint64)
+  pb.write(2, self.digestSize.uint64)
+  pb.write(3, self.leavesCount.uint64)
+  pb.write(4, self.nodesBuffer)
+  pb.finish
+  pb.buffer
+
+proc decode*(_: type MerkleTree, data: seq[byte]): ?!MerkleTree =
+  var pb = initProtoBuffer(data, maxSize = MaxMerkleTreeSize)
+  var mcodecCode: uint64
+  var digestSize: uint64
+  var leavesCount: uint64
+  discard ? pb.getField(1, mcodecCode).mapFailure
+  discard ? pb.getField(2, digestSize).mapFailure
+  discard ? pb.getField(3, leavesCount).mapFailure
+
+  let mcodec = MultiCodec.codec(cast[int](mcodecCode))
+  if mcodec == InvalidMultiCodec:
+    return failure("Invalid MultiCodec code " & $cast[int](mcodec))
+
+  var nodesBuffer = newSeq[byte]()
+  discard ? pb.getField(4, nodesBuffer).mapFailure
+
+  let tree = ? MerkleTree.init(mcodec, digestSize, leavesCount, nodesBuffer)
+  success(tree)

codex/merkletree/merkletree.nim

@@ -15,6 +15,7 @@ import std/sugar
 import pkg/questionable/results
 import pkg/nimcrypto/sha2
 import pkg/libp2p/[multicodec, multihash, vbuffer]
+import pkg/stew/base58
 
 import ../errors
 
@@ -23,7 +24,7 @@ type
     mcodec: MultiCodec
     digestSize: Natural
     leavesCount: Natural
-    nodesBuffer: seq[byte]
+    nodesBuffer*: seq[byte]
   MerkleProof* = object
     mcodec: MultiCodec
     digestSize: Natural
@@ -111,8 +112,8 @@ proc build*(self: MerkleTreeBuilder): ?!MerkleTree =
   copyMem(addr tree.nodesBuffer[0], unsafeAddr self.buffer[0], leavesCount * digestSize)
 
   # calculate intermediate nodes
-  var zero = newSeq[byte](self.digestSize)
+  var zero = newSeq[byte](digestSize)
-  var one = newSeq[byte](self.digestSize)
+  var one = newSeq[byte](digestSize)
   one[^1] = 0x01
 
   var concatBuf = newSeq[byte](2 * digestSize)
@@ -154,6 +155,12 @@ proc len*(self: (MerkleTree | MerkleProof)): Natural =
 proc nodes*(self: (MerkleTree | MerkleProof)): seq[MultiHash] =
   toSeq(0..<self.len).map(i => self.nodeBufferToMultiHash(i))
 
+proc mcodec*(self: (MerkleTree | MerkleProof)): MultiCodec =
+  self.mcodec
+
+proc digestSize*(self: (MerkleTree | MerkleProof)): Natural = 
+  self.digestSize
+
 proc root*(self: MerkleTree): MultiHash =
   let rootIndex = self.len - 1
   self.nodeBufferToMultiHash(rootIndex)
@@ -161,6 +168,9 @@ proc root*(self: MerkleTree): MultiHash =
 proc leaves*(self: MerkleTree): seq[MultiHash] =
   toSeq(0..<self.leavesCount).map(i => self.nodeBufferToMultiHash(i))
 
+proc leavesCount*(self: MerkleTree): Natural =
+  self.leavesCount
+
 proc height*(self: MerkleTree): Natural =
   computeTreeHeight(self.leavesCount)
 
@@ -213,6 +223,33 @@ proc `$`*(self: MerkleTree): string =
     "\nleavesCount: " & $self.leavesCount & 
     "\nnodes: " & $self.nodes
 
+proc `==`*(a, b: MerkleTree): bool =
+  (a.mcodec == b.mcodec) and
+  (a.digestSize == b.digestSize) and
+  (a.leavesCount == b.leavesCount) and
+    (a.nodesBuffer == b.nodesBuffer)
+
+func init*(
+  T: type MerkleTree,
+  mcodec: MultiCodec,
+  digestSize: Natural,
+  leavesCount: Natural,
+  nodesBuffer: seq[byte]
+): ?!MerkleTree =
+  let levels = computeLevels(leavesCount)
+  let totalNodes = levels[^1].offset + 1
+  if totalNodes * digestSize == nodesBuffer.len:
+    success(
+      MerkleTree(
+        mcodec: mcodec, 
+        digestSize: digestSize, 
+        leavesCount: leavesCount, 
+        nodesBuffer: nodesBuffer
+      )
+    )
+  else:
+    failure("Expected nodesBuffer len to be " & $(totalNodes * digestSize) & " but was " & $nodesBuffer.len)
+
 ###########################################################
 # MerkleProof
 ###########################################################
@@ -226,7 +263,10 @@ proc `$`*(self: MerkleProof): string =
     "\nnodes: " & $self.nodes
 
 func `==`*(a, b: MerkleProof): bool =
-  (a.index == b.index) and (a.mcodec == b.mcodec) and (a.digestSize == b.digestSize) == (a.nodesBuffer == b.nodesBuffer)
+  (a.index == b.index) and 
+    (a.mcodec == b.mcodec) and 
+    (a.digestSize == b.digestSize) and
+    (a.nodesBuffer == b.nodesBuffer)
 
 proc init*(
   T: type MerkleProof,

tests/codex/merkletree/testcoders.nim

@@ -0,0 +1,43 @@
+import std/unittest
+
+import pkg/questionable/results
+import pkg/stew/byteutils
+
+import pkg/codex/merkletree
+import ../helpers
+
+checksuite "merkletree - coders":
+  const data =
+    [
+      "0123456789012345678901234567890123456789".toBytes,
+      "1234567890123456789012345678901234567890".toBytes,
+      "2345678901234567890123456789012345678901".toBytes,
+      "3456789012345678901234567890123456789012".toBytes,
+      "4567890123456789012345678901234567890123".toBytes,
+      "5678901234567890123456789012345678901234".toBytes,
+      "6789012345678901234567890123456789012345".toBytes,
+      "7890123456789012345678901234567890123456".toBytes,
+      "8901234567890123456789012345678901234567".toBytes,
+      "9012345678901234567890123456789012345678".toBytes,
+    ]
+
+  test "encoding and decoding a tree yields the same tree":
+    var builder = MerkleTreeBuilder.init(multiCodec("sha2-256")).tryGet()
+    builder.addDataBlock(data[0]).tryGet()
+    builder.addDataBlock(data[1]).tryGet()
+    builder.addDataBlock(data[2]).tryGet()
+    builder.addDataBlock(data[3]).tryGet()
+    builder.addDataBlock(data[4]).tryGet()
+    builder.addDataBlock(data[5]).tryGet()
+    builder.addDataBlock(data[6]).tryGet()
+    builder.addDataBlock(data[7]).tryGet()
+    builder.addDataBlock(data[8]).tryGet()
+    builder.addDataBlock(data[9]).tryGet()
+
+    let tree = builder.build().tryGet()
+    let encodedBytes = tree.encode()
+    echo "encode success, size " & $encodedBytes.len
+    let decodedTree = MerkleTree.decode(encodedBytes).tryGet()
+
+    check:
+      tree == decodedTree

tests/codex/merkletree/testmerkletree.nim

@@ -6,7 +6,7 @@ import pkg/questionable/results
 import pkg/stew/byteutils
 import pkg/nimcrypto/sha2
 
-import codex/merkletree/merkletree
+import pkg/codex/merkletree
 import ../helpers
 
 checksuite "merkletree":

tests/codex/testmerkletree.nim

@@ -1,3 +1,4 @@
 import ./merkletree/testmerkletree
+import ./merkletree/testcoders
 
 {.warning[UnusedImport]: off.}