logos-storage-nim/codex/merkletree/merkletree.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, atomics]

import pkg/questionable/results
import pkg/taskpools
import pkg/chronos/threadsync

import ../errors

type UniqueSeq*[T] = object
  ## A unique pointer to a seq[seq[T]] in shared memory
  ## Can only be moved, not copied
  data: ptr seq[seq[T]]

proc newUniqueSeq*[T](data: sink Isolated[seq[seq[T]]]): UniqueSeq[T] =
  ## Creates a new unique sequence in shared memory
  ## The memory is automatically freed when the object is destroyed
  result.data = cast[ptr seq[seq[T]]](allocShared0(sizeof(seq[seq[T]])))

  result.data[] = extract(data)

proc `=destroy`*[T](p: var UniqueSeq[T]) =
  ## Destructor for UniqueSeq
  if p.data != nil:
    # Clear the sequence to release inner sequences
    p.data[].setLen(0)
    echo "destroying unique seq"
    deallocShared(p.data)
    p.data = nil

proc `=copy`*[T](
  dest: var UniqueSeq[T], src: UniqueSeq[T]
) {.error: "UniqueSeq cannot be copied, only moved".}

proc `=sink`*[T](dest: var UniqueSeq[T], src: UniqueSeq[T]) =
  ## Move constructor for UniqueSeq
  if dest.data != nil:
    `=destroy`(dest)
  dest.data = src.data
  # We need to nil out the source data to prevent double-free
  # This is handled by Nim's destructive move semantics

proc `[]`*[T](p: UniqueSeq[T]): lent seq[seq[T]] =
  ## Access the data (read-only)
  if p.data == nil:
    raise newException(NilAccessDefect, "accessing nil UniqueSeq")
  p.data[]

proc `[]`*[T](p: var UniqueSeq[T]): var seq[seq[T]] =
  ## Access the data (mutable)
  if p.data == nil:
    raise newException(NilAccessDefect, "accessing nil UniqueSeq")
  p.data[]

proc isNil*[T](p: UniqueSeq[T]): bool =
  ## Check if the UniqueSeq is nil
  p.data == nil

proc extractValue*[T](p: var UniqueSeq[T]): seq[seq[T]] =
  ## Extract the value from the UniqueSeq and release the memory
  if p.data == nil:
    raise newException(NilAccessDefect, "extracting from nil UniqueSeq")

  # Move the value out
  var isolated = isolate(p.data[])
  result = extract(isolated)

  # Free the shared memory
  deallocShared(p.data)
  p.data = nil

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

  MerkleTask*[H, K] = object
    tree*: ptr MerkleTree[H, K]
    leaves*: seq[H]
    signal*: ThreadSignalPtr
    layers*: UniqueSeq[H]
    success*: Atomic[bool]

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): ?!bool =
  success bool(root == ?proof.reconstructRoot(leaf))

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)

proc merkleTreeWorker*[H, K](task: ptr MerkleTask[H, K]) {.gcsafe.} =
  defer:
    discard task[].signal.fireSync()

  let res = merkleTreeWorker(task[].tree[], task[].leaves, isBottomLayer = true)

  if res.isErr:
    task[].success.store(false)
    return

  var l = res.get()
  var newOuterSeq = newSeq[seq[H]](l.len)
  for i in 0 ..< l.len:
    var isoInner = isolate(l[i])
    newOuterSeq[i] = extract(isoInner)

  var isolatedLayers = isolate(newOuterSeq)
  task[].layers = newUniqueSeq(isolatedLayers)
  task[].success.store(true)