nim-dagger/codex/merkletree/merkletree.nim

157 lines
4.1 KiB
Nim

## 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.
{.push raises: [].}
import std/bitops
import pkg/questionable/results
import ../errors
type
CompressFn*[H, K] = proc (x, y: H, key: K): ?!H {.noSideEffect, raises: [].}
MerkleTree*[H, K] = ref object of RootObj
layers* : seq[seq[H]]
compress*: CompressFn[H, K]
zero* : H
MerkleProof*[H, K] = ref object of RootObj
index* : int # linear index of the leaf, starting from 0
path* : seq[H] # order: from the bottom to the top
nleaves* : int # number of leaves in the tree (=size of input)
compress*: CompressFn[H, K] # compress function
zero* : H # zero value
func depth*[H, K](self: MerkleTree[H, K]): int =
return self.layers.len - 1
func leavesCount*[H, K](self: MerkleTree[H, K]): int =
return self.layers[0].len
func levels*[H, K](self: MerkleTree[H, K]): int =
return self.layers.len
func leaves*[H, K](self: MerkleTree[H, K]): seq[H] =
return self.layers[0]
iterator layers*[H, K](self: MerkleTree[H, K]): seq[H] =
for layer in self.layers:
yield layer
iterator nodes*[H, K](self: MerkleTree[H, K]): H =
for layer in self.layers:
for node in layer:
yield node
func root*[H, K](self: MerkleTree[H, K]): ?!H =
let last = self.layers[^1]
if last.len != 1:
return failure "invalid tree"
return success last[0]
func getProof*[H, K](
self: MerkleTree[H, K],
index: int,
proof: MerkleProof[H, K]): ?!void =
let depth = self.depth
let nleaves = self.leavesCount
if not (index >= 0 and index < nleaves):
return failure "index out of bounds"
var path : seq[H] = newSeq[H](depth)
var k = index
var m = nleaves
for i in 0..<depth:
let j = k xor 1
path[i] = if (j < m): self.layers[i][j] else: self.zero
k = k shr 1
m = (m + 1) shr 1
proof.index = index
proof.path = path
proof.nleaves = nleaves
proof.compress = self.compress
success()
func getProof*[H, K](self: MerkleTree[H, K], index: int): ?!MerkleProof[H, K] =
var
proof = MerkleProof[H, K]()
? self.getProof(index, proof)
success proof
func reconstructRoot*[H, K](proof: MerkleProof[H, K], leaf: H): ?!H =
var
m = proof.nleaves
j = proof.index
h = leaf
bottomFlag = K.KeyBottomLayer
for p in proof.path:
let oddIndex : bool = (bitand(j,1) != 0)
if oddIndex:
# the index of the child is odd, so the node itself can't be odd (a bit counterintuitive, yeah :)
h = ? proof.compress( p, h, bottomFlag )
else:
if j == m - 1:
# single child => odd node
h = ? proof.compress( h, p, K(bottomFlag.ord + 2) )
else:
# even node
h = ? proof.compress( h , p, bottomFlag )
bottomFlag = K.KeyNone
j = j shr 1
m = (m+1) shr 1
return success h
func verify*[H, K](proof: MerkleProof[H, K], leaf: H, root: H): ?!void =
return if bool(root == ? proof.reconstructRoot(leaf)):
success()
else:
failure("invalid proof")
func merkleTreeWorker*[H, K](
self: MerkleTree[H, K],
xs: openArray[H],
isBottomLayer: static bool): ?!seq[seq[H]] =
let a = low(xs)
let b = high(xs)
let m = b - a + 1
when not isBottomLayer:
if m == 1:
return success @[ @xs ]
let halfn: int = m div 2
let n : int = 2 * halfn
let isOdd: bool = (n != m)
var ys: seq[H]
if not isOdd:
ys = newSeq[H](halfn)
else:
ys = newSeq[H](halfn + 1)
for i in 0..<halfn:
const key = when isBottomLayer: K.KeyBottomLayer else: K.KeyNone
ys[i] = ? self.compress( xs[a + 2 * i], xs[a + 2 * i + 1], key = key )
if isOdd:
const key = when isBottomLayer: K.KeyOddAndBottomLayer else: K.KeyOdd
ys[halfn] = ? self.compress( xs[n], self.zero, key = key )
success @[ @xs ] & ? self.merkleTreeWorker(ys, isBottomLayer = false)