code duplication reduced

src/BaseTree.ts

@@ -0,0 +1,172 @@
+import { Element, HashFunction, ProofPath } from './index'
+
+export class BaseTree {
+  levels: number
+  protected _hashFn: HashFunction<Element>
+  protected zeroElement: Element
+  protected _zeros: Element[]
+  protected _layers: Array<Element[]>
+
+  get capacity() {
+    return 2 ** this.levels
+  }
+
+  get layers(): Array<Element[]> {
+    return this._layers.slice()
+  }
+
+  get zeros(): Element[] {
+    return this._zeros.slice()
+  }
+
+  get elements(): Element[] {
+    return this._layers[0].slice()
+  }
+
+  get root(): Element {
+    return this._layers[this.levels][0] ?? this._zeros[this.levels]
+  }
+
+  /**
+   * Find an element in the tree
+   * @param element An element to find
+   * @param comparator A function that checks leaf value equality
+   * @returns {number} Index if element is found, otherwise -1
+   */
+  indexOf(element: Element, comparator?: <T> (arg0: T, arg1: T) => boolean): number {
+    if (comparator) {
+      return this._layers[0].findIndex((el) => comparator<Element>(element, el))
+    } else {
+      return this._layers[0].indexOf(element)
+    }
+  }
+
+  /**
+   * Insert new element into the tree
+   * @param element Element to insert
+   */
+  insert(element: Element) {
+    if (this._layers[0].length >= this.capacity) {
+      throw new Error('Tree is full')
+    }
+    this.update(this._layers[0].length, element)
+  }
+
+  /*
+   * Insert multiple elements into the tree.
+   * @param {Array} elements Elements to insert
+   */
+  bulkInsert(elements: Element[]): void {
+    if (!elements.length) {
+      return
+    }
+
+    if (this._layers[0].length + elements.length > this.capacity) {
+      throw new Error('Tree is full')
+    }
+    // First we insert all elements except the last one
+    // updating only full subtree hashes (all layers where inserted element has odd index)
+    // the last element will update the full path to the root making the tree consistent again
+    for (let i = 0; i < elements.length - 1; i++) {
+      this._layers[0].push(elements[i])
+      let level = 0
+      let index = this._layers[0].length - 1
+      while (index % 2 === 1) {
+        level++
+        index >>= 1
+        const left = this._layers[level - 1][index * 2]
+        const right = this._layers[level - 1][index * 2 + 1]
+        this._layers[level][index] = this._hashFn(left, right)
+      }
+    }
+    this.insert(elements[elements.length - 1])
+  }
+
+
+  /**
+   * Change an element in the tree
+   * @param {number} index Index of element to change
+   * @param element Updated element value
+   */
+  update(index: number, element: Element) {
+    if (isNaN(Number(index)) || index < 0 || index > this._layers[0].length || index >= this.capacity) {
+      throw new Error('Insert index out of bounds: ' + index)
+    }
+    this._layers[0][index] = element
+    this._processUpdate(index)
+  }
+
+  proof(element: Element): ProofPath {
+    const index = this.indexOf(element)
+    return this.path(index)
+  }
+
+  /**
+   * Get merkle path to a leaf
+   * @param {number} index Leaf index to generate path for
+   * @returns {{pathElements: Object[], pathIndex: number[]}} An object containing adjacent elements and left-right index
+   */
+  path(index: number): ProofPath {
+    if (isNaN(Number(index)) || index < 0 || index >= this._layers[0].length) {
+      throw new Error('Index out of bounds: ' + index)
+    }
+    let elIndex = +index
+    const pathElements: Element[] = []
+    const pathIndices: number[] = []
+    const pathPositions: number [] = []
+    for (let level = 0; level < this.levels; level++) {
+      pathIndices[level] = elIndex % 2
+      const leafIndex = elIndex ^ 1
+      if (leafIndex < this._layers[level].length) {
+        pathElements[level] = this._layers[level][leafIndex]
+        pathPositions[level] = leafIndex
+      } else {
+        pathElements[level] = this._zeros[level]
+        pathPositions[level] = 0
+      }
+      elIndex >>= 1
+    }
+    return {
+      pathElements,
+      pathIndices,
+      pathPositions,
+      pathRoot: this.root,
+    }
+  }
+
+  protected _buildZeros() {
+    this._zeros = [this.zeroElement]
+    for (let i = 1; i <= this.levels; i++) {
+      this._zeros[i] = this._hashFn(this._zeros[i - 1], this._zeros[i - 1])
+    }
+  }
+
+  protected _processNodes(nodes: Element[], layerIndex: number) {
+    const length = nodes.length
+    let currentLength = Math.ceil(length / 2)
+    const currentLayer = new Array(currentLength)
+    currentLength--
+    const starFrom = length - ((length % 2) ^ 1)
+    let j = 0
+    for (let i = starFrom; i >= 0; i -= 2) {
+      if (nodes[i - 1] === undefined) break
+      const left = nodes[i - 1]
+      const right = (i === starFrom && length % 2 === 1) ? this._zeros[layerIndex - 1] : nodes[i]
+      currentLayer[currentLength - j] = this._hashFn(left, right)
+      j++
+    }
+    return currentLayer
+  }
+
+  protected _processUpdate(index: number) {
+    for (let level = 1; level <= this.levels; level++) {
+      index >>= 1
+      const left = this._layers[level - 1][index * 2]
+      const right = index * 2 + 1 < this._layers[level - 1].length
+        ? this._layers[level - 1][index * 2 + 1]
+        : this._zeros[level - 1]
+      this._layers[level][index] = this._hashFn(left, right)
+    }
+  }
+
+}

src/FixedMerkleTree.ts

@@ -1,96 +1,34 @@
-import {
-  Element,
-  HashFunction,
-  MerkleTreeOptions,
-  ProofPath,
-  SerializedTreeState,
-  simpleHash,
-  TreeEdge,
-  TreeSlice,
-} from './'
+import { Element, HashFunction, MerkleTreeOptions, SerializedTreeState, TreeEdge, TreeSlice } from './'
+import defaultHash from './simpleHash'
+import { BaseTree } from './BaseTree'
 
-const defaultHash = (left: Element, right: Element): string => simpleHash([left, right])
-
-export default class MerkleTree {
-  levels: number
-  private _hashFn: HashFunction<Element>
-  private zeroElement: Element
-  private _zeros: Element[]
-  private _layers: Array<Element[]>
 
+export default class MerkleTree extends BaseTree {
   constructor(levels: number, elements: Element[] = [], {
     hashFunction = defaultHash,
     zeroElement = 0,
   }: MerkleTreeOptions = {}) {
+    super()
     this.levels = levels
     if (elements.length > this.capacity) {
       throw new Error('Tree is full')
     }
     this._hashFn = hashFunction
     this.zeroElement = zeroElement
-
     this._layers = []
     const leaves = elements.slice()
     this._layers = [leaves]
     this._buildZeros()
     this._buildHashes()
-    // this._buildHashes2(leaves)
   }
 
-  get capacity() {
-    return 2 ** this.levels
-  }
-
-  get layers(): Array<Element[]> {
-    return this._layers.slice()
-  }
-
-  get zeros(): Element[] {
-    return this._zeros.slice()
-  }
-
-  get elements(): Element[] {
-    return this._layers[0].slice()
-  }
-
-  private _buildZeros() {
-    this._zeros = [this.zeroElement]
-    for (let i = 1; i <= this.levels; i++) {
-      this._zeros[i] = this._hashFn(this._zeros[i - 1], this._zeros[i - 1])
+  private _buildHashes() {
+    for (let layerIndex = 1; layerIndex <= this.levels; layerIndex++) {
+      const nodes = this._layers[layerIndex - 1]
+      this._layers[layerIndex] = this._processNodes(nodes, layerIndex)
     }
   }
 
-  _buildHashes() {
-    for (let level = 1; level <= this.levels; level++) {
-      this._layers[level] = []
-      for (let i = 0; i < Math.ceil(this._layers[level - 1].length / 2); i++) {
-        this._layers[level][i] = this._hashFn(
-          this._layers[level - 1][i * 2],
-          i * 2 + 1 < this._layers[level - 1].length
-            ? this._layers[level - 1][i * 2 + 1]
-            : this._zeros[level - 1],
-        )
-      }
-    }
-  }
-
-  /**
-   * Get tree root
-   */
-  get root(): Element {
-    return this._layers[this.levels][0] ?? this._zeros[this.levels]
-  }
-
-  /**
-   * Insert new element into the tree
-   * @param element Element to insert
-   */
-  insert(element: Element) {
-    if (this._layers[0].length >= this.capacity) {
-      throw new Error('Tree is full')
-    }
-    this.update(this._layers[0].length, element)
-  }
 
   /**
    * Insert multiple elements into the tree.
@@ -123,78 +61,6 @@ export default class MerkleTree {
     this.insert(elements[elements.length - 1])
   }
 
-  /**
-   * Change an element in the tree
-   * @param {number} index Index of element to change
-   * @param element Updated element value
-   */
-  update(index: number, element: Element) {
-    if (isNaN(Number(index)) || index < 0 || index > this._layers[0].length || index >= this.capacity) {
-      throw new Error('Insert index out of bounds: ' + index)
-    }
-    this._layers[0][index] = element
-    for (let level = 1; level <= this.levels; level++) {
-      index >>= 1
-      const left = this._layers[level - 1][index * 2]
-      const right = index * 2 + 1 < this._layers[level - 1].length
-        ? this._layers[level - 1][index * 2 + 1]
-        : this._zeros[level - 1]
-      this._layers[level][index] = this._hashFn(left, right)
-    }
-  }
-
-  /**
-   * Get merkle path to a leaf
-   * @param {number} index Leaf index to generate path for
-   * @returns {{pathElements: Object[], pathIndex: number[]}} An object containing adjacent elements and left-right index
-   */
-  path(index: number): ProofPath {
-    if (isNaN(Number(index)) || index < 0 || index >= this._layers[0].length) {
-      throw new Error('Index out of bounds: ' + index)
-    }
-    let elIndex = +index
-    const pathElements: Element[] = []
-    const pathIndices: number[] = []
-    const pathPositions: number [] = []
-    for (let level = 0; level < this.levels; level++) {
-      pathIndices[level] = elIndex % 2
-      const leafIndex = elIndex ^ 1
-      if (leafIndex < this._layers[level].length) {
-        pathElements[level] = this._layers[level][leafIndex]
-        pathPositions[level] = leafIndex
-      } else {
-        pathElements[level] = this._zeros[level]
-        pathPositions[level] = 0
-      }
-      elIndex >>= 1
-    }
-    return {
-      pathElements,
-      pathIndices,
-      pathPositions,
-      pathRoot: this.root,
-    }
-  }
-
-  /**
-   * Find an element in the tree
-   * @param element An element to find
-   * @param comparator A function that checks leaf value equality
-   * @returns {number} Index if element is found, otherwise -1
-   */
-  indexOf(element: Element, comparator?: <T> (arg0: T, arg1: T) => boolean): number {
-    if (comparator) {
-      return this._layers[0].findIndex((el) => comparator<Element>(element, el))
-    } else {
-      return this._layers[0].indexOf(element)
-    }
-  }
-
-  proof(element: Element): ProofPath {
-    const index = this.indexOf(element)
-    return this.path(index)
-  }
-
   getTreeEdge(edgeIndex: number): TreeEdge {
     const edgeElement = this._layers[0][edgeIndex]
     if (edgeElement === undefined) {

src/PartialMerkleTree.ts

@@ -5,26 +5,16 @@ import {
   MerkleTreeOptions,
   ProofPath,
   SerializedPartialTreeState,
-  simpleHash,
   TreeEdge,
 } from './'
+import defaultHash from './simpleHash'
+import { BaseTree } from './BaseTree'
 
-export const defaultHash = (left: Element, right: Element): string => simpleHash([left, right])
-
-export class PartialMerkleTree {
-  get edgeLeafProof(): ProofPath {
-    return this._edgeLeafProof
-  }
-
-  levels: number
-  private zeroElement: Element
-  private _zeros: Element[]
-  private _layers: Array<Element[]>
+export class PartialMerkleTree extends BaseTree {
   private _leaves: Element[]
   private _leavesAfterEdge: Element[]
   private _edgeLeaf: LeafWithIndex
   private _initialRoot: Element
-  private _hashFn: HashFunction<Element>
   private _edgeLeafProof: ProofPath
   private _proofMap: Map<number, [i: number, el: Element]>
 
@@ -34,6 +24,7 @@ export class PartialMerkleTree {
     edgeIndex,
     edgeElementsCount,
   }: TreeEdge, leaves: Element[], { hashFunction, zeroElement }: MerkleTreeOptions = {}) {
+    super()
     if (edgeIndex + leaves.length !== edgeElementsCount) throw new Error('Invalid number of elements')
     this._edgeLeafProof = edgePath
     this._initialRoot = edgePath.pathRoot
@@ -46,26 +37,6 @@ export class PartialMerkleTree {
     this._buildTree()
   }
 
-  get capacity() {
-    return 2 ** this.levels
-  }
-
-  get layers(): Array<Element[]> {
-    return this._layers.slice()
-  }
-
-  get zeros(): Element[] {
-    return this._zeros.slice()
-  }
-
-  get elements(): Element[] {
-    return this._layers[0].slice()
-  }
-
-  get root(): Element {
-    return this._layers[this.levels][0] ?? this._zeros[this.levels]
-  }
-
   get edgeIndex(): number {
     return this._edgeLeaf.index
   }
@@ -74,6 +45,10 @@ export class PartialMerkleTree {
     return this._edgeLeaf.data
   }
 
+  get edgeLeafProof(): ProofPath {
+    return this._edgeLeafProof
+  }
+
   private _createProofMap() {
     this._proofMap = this.edgeLeafProof.pathPositions.reduce((p, c, i) => {
       p.set(i, [c, this.edgeLeafProof.pathElements[i]])
@@ -94,29 +69,10 @@ export class PartialMerkleTree {
     this._buildHashes()
   }
 
-  private _buildZeros() {
-    this._zeros = [this.zeroElement]
-    for (let i = 1; i <= this.levels; i++) {
-      this._zeros[i] = this._hashFn(this._zeros[i - 1], this._zeros[i - 1])
-    }
-  }
-
-  _buildHashes() {
+  private _buildHashes() {
     for (let layerIndex = 1; layerIndex <= this.levels; layerIndex++) {
       const nodes = this._layers[layerIndex - 1]
-      const length = nodes.length
-      let currentLength = Math.ceil(length / 2)
-      const currentLayer = new Array(currentLength)
-      currentLength--
-      const starFrom = length - ((length % 2) ^ 1)
-      let j = 0
-      for (let i = starFrom; i >= 0; i -= 2) {
-        if (nodes[i - 1] === undefined) break
-        const left = nodes[i - 1]
-        const right = (i === starFrom && length % 2 === 1) ? this._zeros[layerIndex - 1] : nodes[i]
-        currentLayer[currentLength - j] = this._hashFn(left, right)
-        j++
-      }
+      const currentLayer = this._processNodes(nodes, layerIndex)
       if (this._proofMap.has(layerIndex)) {
         const [proofPos, proofEl] = this._proofMap.get(layerIndex)
         if (!currentLayer[proofPos]) currentLayer[proofPos] = proofEl
@@ -125,46 +81,6 @@ export class PartialMerkleTree {
     }
   }
 
-  /**
-   * Insert new element into the tree
-   * @param element Element to insert
-   */
-  insert(element: Element) {
-    if (this._layers[0].length >= this.capacity) {
-      throw new Error('Tree is full')
-    }
-    this.update(this._layers[0].length, element)
-  }
-
-  /*
-   * Insert multiple elements into the tree.
-   * @param {Array} elements Elements to insert
-   */
-  bulkInsert(elements: Element[]): void {
-    if (!elements.length) {
-      return
-    }
-
-    if (this._layers[0].length + elements.length > this.capacity) {
-      throw new Error('Tree is full')
-    }
-    // First we insert all elements except the last one
-    // updating only full subtree hashes (all layers where inserted element has odd index)
-    // the last element will update the full path to the root making the tree consistent again
-    for (let i = 0; i < elements.length - 1; i++) {
-      this._layers[0].push(elements[i])
-      let level = 0
-      let index = this._layers[0].length - 1
-      while (index % 2 === 1) {
-        level++
-        index >>= 1
-        const left = this._layers[level - 1][index * 2]
-        const right = this._layers[level - 1][index * 2 + 1]
-        this._layers[level][index] = this._hashFn(left, right)
-      }
-    }
-    this.insert(elements[elements.length - 1])
-  }
 
   /**
    * Change an element in the tree
@@ -179,14 +95,7 @@ export class PartialMerkleTree {
       throw new Error(`Index ${index} is below the edge: ${this._edgeLeaf.index}`)
     }
     this._layers[0][index] = element
-    for (let level = 1; level <= this.levels; level++) {
-      index >>= 1
-      const left = this._layers[level - 1][index * 2]
-      const right = index * 2 + 1 < this._layers[level - 1].length
-        ? this._layers[level - 1][index * 2 + 1]
-        : this._zeros[level - 1]
-      this._layers[level][index] = this._hashFn(left, right)
-    }
+    this._processUpdate(index)
   }
 
   path(index: number): ProofPath {
@@ -221,18 +130,6 @@ export class PartialMerkleTree {
     }
   }
 
-  indexOf(element: Element, comparator?: <T> (arg0: T, arg1: T) => boolean): number {
-    if (comparator) {
-      return this._layers[0].findIndex((el) => comparator<Element>(element, el))
-    } else {
-      return this._layers[0].indexOf(element)
-    }
-  }
-
-  proof(element: Element): ProofPath {
-    const index = this.indexOf(element)
-    return this.path(index)
-  }
 
   /**
    * Shifts edge of tree to left

src/simpleHash.ts

@@ -1,3 +1,5 @@
+import { Element } from './index'
+
 /***
  * This is insecure hash function, just for example only
  * @param data
@@ -17,3 +19,4 @@ export function simpleHash<T>(data: T[], seed?: number, hashLength = 40): string
   return BigInt('0x' + hash.padEnd(hashLength - (hash.length - 1), '0')).toString(10)
 }
 
+export default (left: Element, right: Element): string => simpleHash([left, right])