@crpdo/merkle/lib/merkle.js

A dynamic, in-memory merkle tree implementation in js

/**
 * @file merkle.js
 * 
 * @description The Merkle module, `@crpdo/merkle`, is responsible for creating 
 * Merkle trees, which are fundamental data structures in various cryptographic 
 * applications. 
 * 
 * @module Merkle
 */

const { Crypto, _: { _, log }} = require('@crpdo/crypto')
const HASH_LENGTH = undefined

/**
 * Class representing a Merkle tree.
 * A Merkle tree is a tree of hashes and is used for efficient data verification.
 * It's widely used in distributed systems and blockchains for verifying the 
 * integrity of transaction data.
 */
class Merkle {
  // static Merkle = Merkle
  static get Merkle() { return Merkle }
  /**
   * Create a Merkle tree.
   * @param {Array} [leaves=[]] - The leaves of the tree.
   * @param {boolean} [checkpoint=true] - Whether to create a checkpoint.
   */
  constructor(leaves = [], checkpoint = true) {
    this.root = this.hash('')
    this.stacks = {}
    if (checkpoint) {
      const res = this.checkpoint()
      this.levels = [leaves]
      this.derive(res.stack)
    } else {
      this.levels = [leaves]
      this.derive()
    }
  }
  
  /**
   * Generates a hash for the provided input using the static hash function.
   * 
   * @param {string} input - The input to hash.
   * @returns {string} The generated hash.
   */
  hash(input) {
    return this.constructor.hash(input)
  }

  /**
   * Static method to create a hash from a given input.
   * 
   * @static
   * @param {string} input - The input to be hashed.
   * @param {number} [hashLength=HASH_LENGTH] - The length of the hash.
   * @returns {string} - The resulting hash.
   */
  static hash(input, hashLength = HASH_LENGTH) {
    return Crypto.hash(input, hashLength)
  }

  /**
   * The depth of the Merkle tree, calculated as the length of the levels.
   * @returns {number} The depth of the Merkle tree.
   */
  get depth() {
    return this.levels.length
  }

  /**
   * Returns the key from the `this.stacks` object. 
   * If no parameter is provided, it returns the last key.
   * @param {number} back - The index of the key to return.
   * @returns {string} The key at the provided index or the last key.
   */
  peek(back) {
    const keys = _.keys(this.stacks)
    if (_.isNil(back))
      return _.last(keys)
    return keys[back]
  }

  /**
   * Creates a new checkpoint for the Merkle tree.
   * This involves creating a new stack and associating it with the root.
   * @returns {Object} An object containing the root and the newly created stack.
   */
  checkpoint() {
    const root = this.root
    const stack = []
    this.stacks[root] = stack
    return { root, stack }
  }

  /**
   * Commits changes made since the last checkpoint.
   * The root and its associated stack are removed from the `this.stacks` object.
   * If `restart` is true, a new checkpoint is created.
   * @param {boolean} [restart=true] - Whether to create a new checkpoint 
   * after committing.
   * @returns {Object} An object containing the root and the stack that was 
   * just committed.
   * @throws {Error} If there is no checkpoint to commit.
   */
  commit(restart = true) {
    const root = this.peek(0)
    const stack = this.stacks[root]
    if (!stack)
      throw new Error(`No checkpoint to commit`)
    delete this.stacks[root]
    if (restart)
      return this.checkpoint()
    return { root, stack }
  }

  /**
   * Reverts changes made since the last checkpoint.
   * The stack associated with the current root is popped until empty, 
   * undoing the operations stored in it.
   * If the root does not match the original root after reverting, false 
   * is returned.
   * @returns {boolean|string} False if the root after reverting does not 
   * match the original root, the root otherwise.
   * @throws {Error} If there is no checkpoint to revert to.
   */
  revert() {
    const root = this.peek()
    const stack = this.stacks[root]
    if (!stack)
      throw new Error(`No checkpoint to revert to`)
    do {
      const task = stack.pop()
      const { op, depth } = task
      const level = this.levels[depth]
      const height = level.length
      if (op === 'row') {
        this.levels.splice(depth, 1)
      } else if (op === 'put') {
        const { width, value } = task
        level[width] = value
      } else if (op === 'del') {
        const { width } = task
        level.splice(width, 1)
      }
    } while (stack.length)
    if (root !== this.getRoot())
      return false
    return this.genRoot()
  }

  /**
   * Derive new levels for the Merkle tree based on the existing ones and generate
   * a new root.
   * 
   * @param {Array} undos - The undo operations.
   */
  derive(undos) {
    this.constructor._derive(this.levels, undos)
    this.genRoot()
  }

  /**
   * Static method to derive new levels for the Merkle tree based on the existing
   * ones.
   * 
   * @static
   * @param {Array} levels - The existing levels of the Merkle tree.
   * @param {Array} undos - The undo operations.
   * @param {number} hashLength - The length of the hash.
   * @returns {Array} The derived levels.
   */
  static derive(levels, undos, hashLength) {
    return this._derive([levels], undos, hashLength)
  }

  /**
   * Internal static method to derive new levels for the Merkle tree based on
   * the existing ones.
   * 
   * @static
   * @param {Array} levels - The existing levels of the Merkle tree.
   * @param {Array} undos - The undo operations.
   * @param {number} hashLength - The length of the hash.
   * @returns {Array} The derived levels.
   */
  static _derive(levels, undos, hashLength) {
    let nodes = _.last(levels)
    while (nodes.length > 1) {
      if (undos)
        undos.push({ op: 'row', depth: levels.length })
      const level = []
      for (let ii = 0; ii < nodes.length; ii += 2) {
        const left = nodes[ii]
        const right = nodes[ii + 1] ? nodes[ii + 1] : left
        level.push(this.hash(left + right, hashLength))
      }
      nodes = level
      levels.push(level)
    }
    return levels
  }

  /**
   * Static method to append a new value to the Merkle tree and derive the new
   * levels.
   * 
   * @static
   * @param {string} value - The value to append.
   * @param {Array} levels - The existing levels of the Merkle tree.
   * @param {Array} [undos=[]] - The undo operations.
   * @returns {Array} The derived levels after appending the new value.
   */
  static append(value, levels, undos = []) {
    const level = _.first(levels)
    level.push(value)
    const undo = { op: 'del', depth: 0, width: level.length }
    undos.push(undo)
    for (let ii = 1; ii < levels.length; ii++) {
      const prev = levels[ii - 1]
      const curr = levels[ii]
      let index, hash
      if (_.even(prev.length)) {
        const left = _.nth(prev, -2)
        const right = _.last(prev)
        index = curr.length - 1
        hash = this.hash(left + right)
      } else {
        const left = _.last(prev)
        index = Math.floor((prev.length - 1) / 2)
        hash = this.hash(left + left)
      }
      const last = curr[index]
      const undo = { op: 'del', depth: ii, width: index }
      if (last) {
        undo.op = 'put'
        undo.value = last
      }
      undos.push(undo)
      curr[index] = hash
    }
    this._derive(levels)
    return levels
  }

  /**
   * Method to append a new value to the Merkle tree and generate a new root.
   * 
   * @param {string} value - The value to append.
   */
  append(value) {
    let stack = []
    if (_.keys(this.stacks).length) {
      const root = this.peek()
      stack = this.stacks[root]
    }
    this.constructor.append(value, this.levels, stack)
    this.genRoot()
  }

  /**
   * Retrieves the root of the Merkle tree.
   * @param {Array} levels - The levels of the tree.
   * @param {boolean} mutate - Whether or not to mutate the levels.
   * @returns {*} - The root of the tree.
   */
  static getRoot(levels, mutate) {
    if (mutate)
      return levels.pop().pop()
    return _.get(_.last(levels), 0, null)
  }

  /**
   * Gets the root of the Merkle tree instance.
   * @returns {*} - The root of the tree.
   */
  getRoot() {
    return this.constructor.getRoot(this.levels)
  }

  /**
   * Generates and sets the root of the Merkle tree instance.
   * @returns {*} - The newly set root of the tree.
   */
  genRoot() {
    this.root = this.getRoot()
    return this.root
  }

  /**
   * Gets the leaf at the specified index.
   * @param {number} index - The index of the leaf to retrieve.
   * @returns {*} - The leaf at the specified index.
   */
  getLeaf(index) {
    return this.levels[0][index]
  }

  /**
   * Gets the count of leaves in the Merkle tree instance.
   * @returns {number} - The count of leaves.
   */
  getLeafCount() {
    return this.levels[0].length
  }

  /**
   * Gets the index of the specified leaf in the Merkle tree instance.
   * @param {*} leaf - The leaf to find.
   * @returns {number} - The index of the leaf if found, -1 otherwise.
   */
  getIndex(leaf) {
    return this.levels[0].indexOf(leaf)
  }

  /**
   * Generates a proof for a given leaf in a Merkle tree.
   * The proof consists of sibling hashes from leaf to root.
   * 
   * @static
   * @param {number} index - Index of the leaf for which the proof is generated.
   * @param {Array} levels - Array containing all levels of the Merkle tree.
   * @returns {Array} - An array containing pairs of left and right node hashes
   * for each level of the tree, starting from the given leaf level.
   */
  static getProof(index, levels) {
    let proof = []
    for (let ii = 0; ii < levels.length; ii++) {
      let level = levels[ii]
      let width = level.length
      if (!(index === width - 1 && width % 2 === 1)) {
        const left = (index % 2) ? level[index - 1] : level[index]
        const right = (index % 2) ? level[index] : level[index + 1]
        proof.push([left, right])
      }
      index = Math.floor(index / 2)
    }
    return proof
  }

  /**
   * Generates a Merkle proof for a given leaf.
   * 
   * @param {number} index - Index of the leaf for which to generate the proof.
   * @returns {Object} - An object containing the Merkle proof, the target leaf,
   * its index, and the root of the tree.
   */
  getProof(index) {
    let target
    if (_.isInteger(index))
      target = this.getLeaf(index)
    else
      [target, index] = [index, this.getIndex(index)]
    const root = this.getRoot()
    const proof = this.constructor.getProof(index, this.levels)
    return { proof, target, index, root }
  }

  /**
   * Verifies a Merkle proof.
   * 
   * @param {Object} proofData - An object containing the Merkle proof, the root
   * of the tree, and the target leaf.
   * @returns {boolean} - Returns true if the proof is valid, false otherwise.
   */
  verifyProof({ proof, root, target }) {
    return this.constructor.verifyProof(proof, root, target)
  }

  /**
   * Verifies a Merkle proof.
   * 
   * @static
   * @param {Array} proof - Array containing pairs of left and right node hashes
   * for each level of the tree.
   * @param {string} root - The root hash of the tree.
   * @param {string} target - The target leaf value.
   * @returns {boolean} - Returns true if the proof hash equals the root,
   * indicating the proof is valid, false otherwise.
   */
  static verifyProof(proof, root, target) {
    if (!proof.length)
      return target === root
    let proofHash
    for (let ii = 0; ii < proof.length; ii++) {
      const [left, right] = proof[ii]
      proofHash = this.hash(left + right)
    }
    return proofHash === root
  }
}

module.exports = Merkle