190 lines
5.1 KiB
Nim
190 lines
5.1 KiB
Nim
## Nim-Codex
|
|
## Copyright (c) 2022 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 std/sequtils
|
|
import std/math
|
|
import std/bitops
|
|
import std/sugar
|
|
|
|
import pkg/libp2p
|
|
import pkg/stew/byteutils
|
|
import pkg/questionable
|
|
import pkg/questionable/results
|
|
|
|
type
|
|
MerkleHash* = MultiHash
|
|
MerkleTree* = object
|
|
leavesCount: int
|
|
nodes: seq[MerkleHash]
|
|
MerkleProof* = object
|
|
index: int
|
|
path: seq[MerkleHash]
|
|
|
|
# Tree constructed from leaves H0..H2 is
|
|
# H5=H(H3 & H4)
|
|
# / \
|
|
# H3=H(H0 & H1) H4=H(H2 & H2)
|
|
# / \ /
|
|
# H0=H(A) H1=H(B) H2=H(C)
|
|
# | | |
|
|
# A B C
|
|
#
|
|
# Memory layout is [H0, H1, H2, H3, H4, H5]
|
|
#
|
|
# Proofs of inclusion are
|
|
# - [H1, H4] for A
|
|
# - [H0, H4] for B
|
|
# - [H2, H3] for C
|
|
|
|
|
|
func computeTreeHeight(leavesCount: int): int =
|
|
if isPowerOfTwo(leavesCount):
|
|
fastLog2(leavesCount) + 1
|
|
else:
|
|
fastLog2(leavesCount) + 2
|
|
|
|
func getLowHigh(leavesCount, level: int): (int, int) =
|
|
var width = leavesCount
|
|
var low = 0
|
|
for _ in 0..<level:
|
|
low += width
|
|
width = (width + 1) div 2
|
|
|
|
(low, low + width - 1)
|
|
|
|
func getLowHigh(self: MerkleTree, level: int): (int, int) =
|
|
getLowHigh(self.leavesCount, level)
|
|
|
|
func computeTotalSize(leavesCount: int): int =
|
|
let height = computeTreeHeight(leavesCount)
|
|
getLowHigh(leavesCount, height - 1)[1] + 1
|
|
|
|
proc getWidth(self: MerkleTree, level: int): int =
|
|
let (low, high) = self.getLowHigh(level)
|
|
high - low + 1
|
|
|
|
func getChildren(self: MerkleTree, level, i: int): (MerkleHash, MerkleHash) =
|
|
let (low, high) = self.getLowHigh(level - 1)
|
|
let leftIdx = low + 2 * i
|
|
let rightIdx = min(leftIdx + 1, high)
|
|
|
|
(self.nodes[leftIdx], self.nodes[rightIdx])
|
|
|
|
func getSibling(self: MerkleTree, level, i: int): MerkleHash =
|
|
let (low, high) = self.getLowHigh(level)
|
|
if i mod 2 == 0:
|
|
self.nodes[min(low + i + 1, high)]
|
|
else:
|
|
self.nodes[low + i - 1]
|
|
|
|
proc setNode(self: var MerkleTree, level, i: int, value: MerkleHash): void =
|
|
let (low, _) = self.getLowHigh(level)
|
|
self.nodes[low + i] = value
|
|
|
|
proc root*(self: MerkleTree): MerkleHash =
|
|
self.nodes[^1]
|
|
|
|
proc len*(self: MerkleTree): int =
|
|
self.nodes.len
|
|
|
|
proc leaves*(self: MerkleTree): seq[MerkleHash] =
|
|
self.nodes[0..<self.leavesCount]
|
|
|
|
proc nodes*(self: MerkleTree): seq[MerkleHash] =
|
|
self.nodes
|
|
|
|
proc height*(self: MerkleTree): int =
|
|
computeTreeHeight(self.leavesCount)
|
|
|
|
proc `$`*(self: MerkleTree): string =
|
|
result &= "leavesCount: " & $self.leavesCount
|
|
result &= "\nnodes: " & $self.nodes
|
|
|
|
proc getProof*(self: MerkleTree, index: int): ?!MerkleProof =
|
|
if index >= self.leavesCount or index < 0:
|
|
return failure("Index " & $index & " out of range [0.." & $self.leaves.high & "]" )
|
|
|
|
var path = newSeq[MerkleHash](self.height - 1)
|
|
for level in 0..<path.len:
|
|
let i = index div (1 shl level)
|
|
path[level] = self.getSibling(level, i)
|
|
|
|
success(MerkleProof(index: index, path: path))
|
|
|
|
proc initTreeFromLeaves(leaves: openArray[MerkleHash]): ?!MerkleTree =
|
|
without mcodec =? leaves.?[0].?mcodec and
|
|
digestSize =? leaves.?[0].?size:
|
|
return failure("At least one leaf is required")
|
|
|
|
if not leaves.allIt(it.mcodec == mcodec):
|
|
return failure("All leaves must use the same codec")
|
|
|
|
let totalSize = computeTotalSize(leaves.len)
|
|
var tree = MerkleTree(leavesCount: leaves.len, nodes: newSeq[MerkleHash](totalSize))
|
|
|
|
var buf = newSeq[byte](digestSize * 2)
|
|
proc combine(l, r: MerkleHash): ?!MerkleHash =
|
|
copyMem(addr buf[0], unsafeAddr l.data.buffer[0], digestSize)
|
|
copyMem(addr buf[digestSize], unsafeAddr r.data.buffer[0], digestSize)
|
|
|
|
MultiHash.digest($mcodec, buf).mapErr(
|
|
c => newException(CatchableError, "Error calculating hash using codec " & $mcodec & ": " & $c)
|
|
)
|
|
|
|
# copy leaves
|
|
for i in 0..<tree.getWidth(0):
|
|
tree.setNode(0, i, leaves[i])
|
|
|
|
# calculate intermediate nodes
|
|
for level in 1..<tree.height:
|
|
for i in 0..<tree.getWidth(level):
|
|
let (left, right) = tree.getChildren(level, i)
|
|
|
|
without mhash =? combine(left, right), error:
|
|
return failure(error)
|
|
tree.setNode(level, i, mhash)
|
|
|
|
success(tree)
|
|
|
|
func init*(
|
|
T: type MerkleTree,
|
|
root: MerkleHash,
|
|
leavesCount: int
|
|
): MerkleTree =
|
|
let totalSize = computeTotalSize(leavesCount)
|
|
var nodes = newSeq[MerkleHash](totalSize)
|
|
nodes[^1] = root
|
|
MerkleTree(nodes: nodes, leavesCount: leavesCount)
|
|
|
|
proc init*(
|
|
T: type MerkleTree,
|
|
leaves: openArray[MerkleHash]
|
|
): ?!MerkleTree =
|
|
initTreeFromLeaves(leaves)
|
|
|
|
proc index*(self: MerkleProof): int =
|
|
self.index
|
|
|
|
proc path*(self: MerkleProof): seq[MerkleHash] =
|
|
self.path
|
|
|
|
proc `$`*(self: MerkleProof): string =
|
|
result &= "index: " & $self.index
|
|
result &= "\npath: " & $self.path
|
|
|
|
func `==`*(a, b: MerkleProof): bool =
|
|
(a.index == b.index) and (a.path == b.path)
|
|
|
|
proc init*(
|
|
T: type MerkleProof,
|
|
index: int,
|
|
path: seq[MerkleHash]
|
|
): MerkleProof =
|
|
MerkleProof(index: index, path: path)
|