## 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)