wip

package.json

@@ -38,6 +38,7 @@
     "nyc": "^15.1.0",
     "mocha": "^9.2.1",
     "ts-mocha": "^9.0.2",
-    "typescript": "^4.5.5"
+    "typescript": "^4.5.5",
+    "circomlib": "git+https://github.com/tornadocash/circomlib.git#5beb6aee94923052faeecea40135d45b6ce6172c"
   }
 }

src/FixedMerkleTree.ts

@@ -21,13 +21,15 @@ export default class MerkleTree {
     this.zeroElement = zeroElement
 
     this._layers = []
-    this._layers[0] = elements.slice()
+    const leaves = elements.slice()
+    this._layers = [leaves]
     this._buildZeros()
     this._buildHashes()
+    // this._buildHashes2(leaves)
   }
 
   get capacity() {
-    return this.levels ** 2
+    return 2 ** this.levels
   }
 
   get layers(): Array<Element[]> {
@@ -63,6 +65,20 @@ export default class MerkleTree {
     }
   }
 
+  // _buildHashes2(nodes: Element[]) {
+  //   let layerIndex = 0
+  //   while (layerIndex < this.levels) {
+  //     layerIndex = this._layers.length
+  //     this._layers[layerIndex] = []
+  //     for (let i = 0; i < nodes.length; i += 2) {
+  //       const left = nodes[i]
+  //       const right = (i + 1 === nodes.length && nodes.length % 2 === 1) ? this._zeros[layerIndex - 1] : nodes[i + 1]
+  //       this._layers[layerIndex].push(this._hashFn(left, right))
+  //     }
+  //     nodes = this._layers[layerIndex]
+  //   }
+  // }
+
   /**
    * Get tree root
    */
@@ -124,12 +140,11 @@ export default class MerkleTree {
     this._layers[0][index] = element
     for (let level = 1; level <= this.levels; level++) {
       index >>= 1
-      this._layers[level][index] = this._hashFn(
+      const left = this._layers[level - 1][index * 2]
-        this._layers[level - 1][index * 2],
+      const right = index * 2 + 1 < this._layers[level - 1].length
-        index * 2 + 1 < this._layers[level - 1].length
+        ? this._layers[level - 1][index * 2 + 1]
-          ? this._layers[level - 1][index * 2 + 1]
+        : this._zeros[level - 1]
-          : this._zeros[level - 1],
+      this._layers[level][index] = this._hashFn(left, right)
-      )
     }
   }
 
@@ -162,6 +177,7 @@ export default class MerkleTree {
       pathElements,
       pathIndices,
       pathPositions,
+      pathRoot: this.root,
     }
   }
 
@@ -184,9 +200,9 @@ export default class MerkleTree {
     return this.path(index)
   }
 
-  getTreeEdge(edgeElement: Element): TreeEdge {
+  getTreeEdge(edgeIndex: number): TreeEdge {
-    const edgeIndex = this.indexOf(edgeElement)
+    const edgeElement = this._layers[0][edgeIndex]
-    if (edgeIndex <= -1) {
+    if (!edgeElement) {
       throw new Error('Element not found')
     }
     const edgePath = this.path(edgeIndex)

src/PartialMerkleTree.ts

@@ -12,6 +12,10 @@ import {
 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[]
@@ -22,26 +26,28 @@ export class PartialMerkleTree {
   private _initialRoot: Element
   private _hashFn: HashFunction<Element>
   private _edgeLeafProof: ProofPath
+  private _proofMap: Map<number, [i: number, el: Element]>
 
   constructor(levels: number, {
     edgePath,
     edgeElement,
     edgeIndex,
-  }: TreeEdge, leaves: Element[], root: Element, { hashFunction, zeroElement }: MerkleTreeOptions = {}) {
+  }: TreeEdge, leaves: Element[], { hashFunction, zeroElement }: MerkleTreeOptions = {}) {
     hashFunction = hashFunction || defaultHash
-    const hashFn = (left, right) => (left !== null && right !== null) ? hashFunction(left, right) : null
+    const hashFn = (left, right) => (left !== undefined && right !== undefined) ? hashFunction(left, right) : undefined
     this._edgeLeafProof = edgePath
+    this._initialRoot = edgePath.pathRoot
     this.zeroElement = zeroElement ?? 0
     this._edgeLeaf = { data: edgeElement, index: edgeIndex }
     this._leavesAfterEdge = leaves
-    this._initialRoot = root
     this.levels = levels
     this._hashFn = hashFn
+    this._createProofMap()
     this._buildTree()
   }
 
   get capacity() {
-    return this.levels ** 2
+    return 2 ** this.levels
   }
 
   get layers(): Array<Element[]> {
@@ -68,15 +74,25 @@ export class PartialMerkleTree {
     return this._edgeLeaf.data
   }
 
+  private _createProofMap() {
+    this._proofMap = this.edgeLeafProof.pathPositions.reduce((p, c, i) => {
+      p.set(i, [c, this.edgeLeafProof.pathElements[i]])
+      return p
+    }, new Map())
+  }
+
   private _buildTree(): void {
     const edgeLeafIndex = this._edgeLeaf.index
-    this._leaves = [...Array.from({ length: edgeLeafIndex }, () => null), ...this._leavesAfterEdge]
+    this._leaves = Array(edgeLeafIndex).concat(this._leavesAfterEdge)
     if (this._edgeLeafProof.pathIndices[0] === 1) {
       this._leaves[this._edgeLeafProof.pathPositions[0]] = this._edgeLeafProof.pathElements[0]
     }
     this._layers = [this._leaves]
+
     this._buildZeros()
-    this._buildHashes()
+    // this._buildHashes()
+    this._buildHashes2()
+
   }
 
   private _buildZeros() {
@@ -102,6 +118,30 @@ export class PartialMerkleTree {
     }
   }
 
+  _buildHashes2() {
+    let index = this.edgeIndex
+    let nodes: Element[]
+    for (let layerIndex = 1; layerIndex <= this.levels; layerIndex++) {
+      nodes = this._layers[layerIndex - 1]
+      // console.log({ layerIndex, nodes })
+      this._layers[layerIndex] = []
+      index = layerIndex > 1 ? Math.ceil(index / 2) : index
+      for (let i = 0; i < nodes.length; i += 2) {
+        const left = nodes[i]
+        const right = (i + 1 < nodes.length) ? nodes[i + 1] : this._zeros[layerIndex - 1]
+        let hash: Element = this._hashFn(left, right)
+        if (layerIndex === this.levels) hash = hash || this._edgeLeafProof.pathRoot
+        // console.log({ layerIndex, i, left, right, hash })
+        this._layers[layerIndex].push(hash)
+      }
+      if (this._proofMap.has(layerIndex)) {
+        const [proofPos, proofEl] = this._proofMap.get(layerIndex)
+        this._layers[layerIndex][proofPos] = this._layers[layerIndex][proofPos] || proofEl
+      }
+
+    }
+  }
+
   /**
    * Insert new element into the tree
    * @param element Element to insert
@@ -113,7 +153,7 @@ export class PartialMerkleTree {
     this.update(this._layers[0].length, element)
   }
 
-  /**
+  /*
    * Insert multiple elements into the tree.
    * @param {Array} elements Elements to insert
    */
@@ -135,10 +175,11 @@ export class PartialMerkleTree {
       while (index % 2 === 1) {
         level++
         index >>= 1
-        this._layers[level][index] = this._hashFn(
+        const left = this._layers[level - 1][index * 2]
-          this._layers[level - 1][index * 2],
+        const right = this._layers[level - 1][index * 2 + 1]
-          this._layers[level - 1][index * 2 + 1],
+        let hash: Element = this._hashFn(left, right)
-        )
+        if (!hash && this._edgeLeafProof.pathPositions[level] === i) hash = this._edgeLeafProof.pathElements[level]
+        this._layers[level][index] = hash
       }
     }
     this.insert(elements[elements.length - 1])
@@ -157,6 +198,8 @@ 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]
@@ -164,7 +207,7 @@ export class PartialMerkleTree {
         ? this._layers[level - 1][index * 2 + 1]
         : this._zeros[level - 1]
       let hash: Element = this._hashFn(left, right)
-      if (!hash && this._edgeLeafProof.pathPositions[level] === index) {
+      if (!hash && this._edgeLeafProof.pathPositions[level] === index * 2) {
         hash = this._edgeLeafProof.pathElements[level]
       }
       if (level === this.levels) {
@@ -189,7 +232,8 @@ export class PartialMerkleTree {
       pathIndices[level] = elIndex % 2
       const leafIndex = elIndex ^ 1
       if (leafIndex < this._layers[level].length) {
-        pathElements[level] = this._layers[level][leafIndex]
+        const [proofPos, proofEl] = this._proofMap.get(level)
+        pathElements[level] = proofPos === leafIndex ? proofEl : this._layers[level][leafIndex]
         pathPositions[level] = leafIndex
       } else {
         pathElements[level] = this._zeros[level]
@@ -201,6 +245,7 @@ export class PartialMerkleTree {
       pathElements,
       pathIndices,
       pathPositions,
+      pathRoot: this.root,
     }
   }
 
@@ -228,11 +273,12 @@ export class PartialMerkleTree {
       throw new Error(`New edgeIndex should be smaller then ${this._edgeLeaf.index}`)
     }
     if (elements.length !== (this._edgeLeaf.index - edge.edgeIndex)) {
-      throw new Error(`Elements length should be ${elements.length}`)
+      throw new Error(`Elements length should be ${this._edgeLeaf.index - edge.edgeIndex}`)
     }
     this._edgeLeafProof = edge.edgePath
     this._edgeLeaf = { index: edge.edgeIndex, data: edge.edgeElement }
     this._leavesAfterEdge = [...elements, ...this._leavesAfterEdge]
+    this._createProofMap()
     this._buildTree()
   }
 
@@ -254,7 +300,7 @@ export class PartialMerkleTree {
       edgeElement: data._edgeLeaf.data,
       edgeIndex: data._edgeLeaf.index,
     }
-    return new PartialMerkleTree(data.levels, edge, data.leaves, data._initialRoot, {
+    return new PartialMerkleTree(data.levels, edge, data.leaves, {
       hashFunction,
       zeroElement: data._zeros[0],
     })

src/index.ts

@@ -32,11 +32,12 @@ export type ProofPath = {
   pathElements: Element[],
   pathIndices: number[],
   pathPositions: number[],
+  pathRoot: Element
 }
 export type TreeEdge = {
   edgeElement: Element;
   edgePath: ProofPath;
-  edgeIndex: number
+  edgeIndex: number;
 }
 export type LeafWithIndex = { index: number, data: Element }
 

test/fixedMerkleTree.spec.ts

@@ -1,9 +1,12 @@
 import { MerkleTree, TreeEdge } from '../src'
 import { assert, should } from 'chai'
+import { mimcsponge } from 'circomlib'
 import { createHash } from 'crypto'
 import { it } from 'mocha'
 
 const sha256Hash = (left, right) => createHash('sha256').update(`${left}${right}`).digest('hex')
+const mimcHash = (left, right) => mimcsponge.multiHash([BigInt(left), BigInt(right)]).toString()
+const ZERO_ELEMENT = '21663839004416932945382355908790599225266501822907911457504978515578255421292'
 
 describe('MerkleTree', () => {
 
@@ -42,6 +45,11 @@ describe('MerkleTree', () => {
       const tree = new MerkleTree(10, [1, 2, 3, 4, 5, 6], { hashFunction: sha256Hash, zeroElement: 'zero' })
       should().equal(tree.root, 'a377b9fa0ed41add83e56f7e1d0e2ebdb46550b9d8b26b77dece60cb67283f19')
     })
+
+    it('should work with mimc hash function and zero element', () => {
+      const tree = new MerkleTree(10, [1, 2, 3], { hashFunction: mimcHash, zeroElement: ZERO_ELEMENT })
+      should().equal(tree.root, '13605252518346649016266481317890801910232739395710162921320863289825142055129')
+    })
   })
 
   describe('#insert', () => {
@@ -232,11 +240,13 @@ describe('MerkleTree', () => {
       const tree = new MerkleTree(10, [1, 2, 3, 4])
       should().throw((() => tree.path(-1)), 'Index out of bounds: -1')
       should().throw((() => tree.path(5)), 'Index out of bounds: 5')
+      // @ts-ignore
       should().throw((() => tree.path('qwe')), 'Index out of bounds: qwe')
     })
 
     it('should work for correct string index', () => {
       const tree = new MerkleTree(10, [1, 2, 3, 4, 5])
+      // @ts-ignore
       const path = tree.path('2')
       assert.deepEqual(path.pathIndices, [0, 1, 0, 0, 0, 0, 0, 0, 0, 0])
       assert.deepEqual(path.pathElements, [
@@ -273,6 +283,7 @@ describe('MerkleTree', () => {
           ],
           pathIndices: [0, 0, 1, 0],
           pathPositions: [5, 0, 0, 0],
+          pathRoot: '3283298202329284319899364273680487022592',
         },
         edgeElement: 4,
         edgeIndex: 4,

test/partialMerkleTree.spec.ts

@@ -7,15 +7,15 @@ import { createHash } from 'crypto'
 const sha256Hash = (left, right) => createHash('sha256').update(`${left}${right}`).digest('hex')
 
 describe('PartialMerkleTree', () => {
-  const getTestTrees = (levels: number, elements: Element[], edgeElement: Element, treeOptions: MerkleTreeOptions = {}) => {
+  const getTestTrees = (levels: number, elements: Element[], edgeIndex: number, treeOptions: MerkleTreeOptions = {}) => {
     const fullTree = new MerkleTree(levels, elements, treeOptions)
-    const edge = fullTree.getTreeEdge(edgeElement)
+    const edge = fullTree.getTreeEdge(edgeIndex)
     const leavesAfterEdge = elements.slice(edge.edgeIndex)
-    const partialTree = new PartialMerkleTree(levels, edge, leavesAfterEdge, fullTree.root, treeOptions)
+    const partialTree = new PartialMerkleTree(levels, edge, leavesAfterEdge, treeOptions)
     return { fullTree, partialTree }
   }
   describe('#constructor', () => {
-    const { fullTree, partialTree } = getTestTrees(20, ['0', '1', '2', '3', '4', '5'], '2')
+    const { fullTree, partialTree } = getTestTrees(20, ['0', '1', '2', '3', '4', '5'], 2)
     it('should initialize merkle tree with same root', () => {
       should().equal(fullTree.root, partialTree.root)
     })
@@ -25,7 +25,7 @@ describe('PartialMerkleTree', () => {
     })
 
     it('should work with optional hash function and zero element', () => {
-      const { partialTree, fullTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6], 4, {
+      const { partialTree, fullTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6], 3, {
         hashFunction: sha256Hash,
         zeroElement: 'zero',
       })
@@ -36,15 +36,15 @@ describe('PartialMerkleTree', () => {
   describe('#insert', () => {
 
     it('should have equal root to full tree after insertion ', () => {
-      const { fullTree, partialTree } = getTestTrees(10, ['0', '1', '2', '3', '4', '5', '6', '7'], '5')
+      const { fullTree, partialTree } = getTestTrees(10, ['0', '1', '2', '3', '4', '5', '6', '7'], 5)
       fullTree.insert('9')
       partialTree.insert('9')
       should().equal(fullTree.root, partialTree.root)
     })
 
     it('should fail to insert when tree is full', () => {
-      const { partialTree } = getTestTrees(3, ['0', '1', '2', '3', '4', '5', '6', '7', '8'], '5')
+      const { partialTree } = getTestTrees(3, ['0', '1', '2', '3', '4', '5', '6', '7'], 5)
-      const call = () => partialTree.insert('9')
+      const call = () => partialTree.insert('8')
       should().throw(call, 'Tree is full')
     })
   })
@@ -52,7 +52,7 @@ describe('PartialMerkleTree', () => {
   describe('#bulkInsert', () => {
 
     it('should work like full tree', () => {
-      const { fullTree, partialTree } = getTestTrees(20, [1, 2, 3, 4, 5], 3)
+      const { fullTree, partialTree } = getTestTrees(20, [1, 2, 3, 4, 5], 2)
       partialTree.bulkInsert([6, 7, 8])
       fullTree.bulkInsert([6, 7, 8])
       should().equal(fullTree.root, partialTree.root)
@@ -73,8 +73,8 @@ describe('PartialMerkleTree', () => {
       ]
       for (const initial of initialArray) {
         for (const inserted of insertedArray) {
-          const { partialTree: tree1 } = getTestTrees(10, initial, initial.length > 1 ? initial.length - 1 : initial.length)
+          const { partialTree: tree1 } = getTestTrees(10, initial, initial.length - 1)
-          const { partialTree: tree2 } = getTestTrees(10, initial, initial.length > 1 ? initial.length - 1 : initial.length)
+          const { partialTree: tree2 } = getTestTrees(10, initial, initial.length - 1)
           tree1.bulkInsert(inserted)
           for (const item of inserted) {
             tree2.insert(item)
@@ -85,20 +85,20 @@ describe('PartialMerkleTree', () => {
     }).timeout(10000)
 
     it('should fail to insert too many elements', () => {
-      const { partialTree } = getTestTrees(2, [1, 2, 3, 4], 3)
+      const { partialTree } = getTestTrees(2, [1, 2, 3, 4], 2)
       const call = () => partialTree.bulkInsert([5, 6, 7])
       should().throw(call, 'Tree is full')
     })
     it('should bypass empty elements', () => {
       const elements = [1, 2, 3, 4]
-      const { partialTree } = getTestTrees(2, elements, 3)
+      const { partialTree } = getTestTrees(2, elements, 2)
       partialTree.bulkInsert([])
       should().equal(partialTree.elements.length, elements.length, 'No elements inserted')
     })
   })
   describe('#update', () => {
     it('should update last element', () => {
-      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       partialTree.update(4, 42)
       fullTree.update(4, 42)
       should().equal(partialTree.root, fullTree.root)
@@ -106,28 +106,28 @@ describe('PartialMerkleTree', () => {
 
 
     it('should update odd element', () => {
-      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8], 3)
+      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8], 2)
       partialTree.update(4, 42)
       fullTree.update(4, 42)
       should().equal(partialTree.root, fullTree.root)
     })
 
     it('should update even element', () => {
-      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8], 3)
+      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8], 2)
       partialTree.update(3, 42)
       fullTree.update(3, 42)
       should().equal(partialTree.root, fullTree.root)
     })
 
     it('should update extra element', () => {
-      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       partialTree.update(5, 6)
       fullTree.update(5, 6)
       should().equal(fullTree.root, partialTree.root)
     })
 
     it('should fail to update incorrect index', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 4)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
       should().throw((() => partialTree.update(-1, 42)), 'Insert index out of bounds: -1')
       should().throw((() => partialTree.update(6, 42)), 'Insert index out of bounds: 6')
       should().throw((() => partialTree.update(2, 42)), 'Index 2 is below the edge: 3')
@@ -136,36 +136,36 @@ describe('PartialMerkleTree', () => {
     })
 
     it('should fail to update over capacity', () => {
-      const { partialTree } = getTestTrees(2, [1, 2, 3, 4], 2)
+      const { partialTree } = getTestTrees(2, [1, 2, 3, 4], 1)
       const call = () => partialTree.update(4, 42)
       should().throw(call, 'Insert index out of bounds: 4')
     })
   })
   describe('#indexOf', () => {
     it('should return same result as full tree', () => {
-      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8], 4)
+      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8], 3)
       should().equal(partialTree.indexOf(5), fullTree.indexOf(5))
     })
 
     it('should find index', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       should().equal(partialTree.indexOf(3), 2)
     })
 
     it('should work with comparator', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       should().equal(partialTree.indexOf(4, (arg0, arg1) => arg0 === arg1), 3)
     })
 
     it('should return -1 for non existent element', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       should().equal(partialTree.indexOf(42), -1)
     })
   })
 
   describe('#proof', () => {
     it('should return proof for known leaf', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       assert.deepEqual(partialTree.proof(4), partialTree.path(3))
     })
   })
@@ -173,8 +173,8 @@ describe('PartialMerkleTree', () => {
   describe('#getters', () => {
     it('should return capacity', () => {
       const levels = 10
-      const capacity = levels ** 2
+      const capacity = 2 ** levels
-      const { fullTree, partialTree } = getTestTrees(levels, [1, 2, 3, 4, 5], 3)
+      const { fullTree, partialTree } = getTestTrees(levels, [1, 2, 3, 4, 5], 2)
       should().equal(fullTree.capacity, capacity)
       should().equal(partialTree.capacity, capacity)
     })
@@ -188,13 +188,13 @@ describe('PartialMerkleTree', () => {
     })
 
     it('should return copy of layers', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       const layers = partialTree.layers
       should().not.equal(layers, partialTree.layers)
     })
 
     it('should return copy of zeros', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 2)
       const zeros = partialTree.zeros
       should().not.equal(zeros, partialTree.zeros)
     })
@@ -204,14 +204,14 @@ describe('PartialMerkleTree', () => {
 
     it('should return path for known nodes', () => {
       const levels = 20
-      const capacity = levels ** 2
+      const capacity = 2 ** levels
-      const elements = Array.from({ length: capacity }, (_, i) => i)
+      const elements = Array.from({ length: capacity / 2 }, (_, i) => i)
       const { fullTree, partialTree } = getTestTrees(levels, elements, 250)
       assert.deepEqual(fullTree.path(250), partialTree.path(250))
-    })
+    }).timeout(10000)
 
     it('should fail on incorrect index', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8, 9], 5)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8, 9], 4)
       should().throw((() => partialTree.path(-1)), 'Index out of bounds: -1')
       should().throw((() => partialTree.path(10)), 'Index out of bounds: 10')
       // @ts-ignore
@@ -219,39 +219,38 @@ describe('PartialMerkleTree', () => {
     })
 
     it('should fail if index is below edge', () => {
-      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8, 9], 5)
+      const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8, 9], 4)
       const call = () => partialTree.path(2)
       should().throw(call, 'Index 2 is below the edge: 4')
     })
   })
   describe('#shiftEdge', () => {
     it('should work', () => {
-      const levels = 20
+      const levels = 10
-      const elements: Element[] = Array.from({ length: levels ** 2 }, (_, i) => i)
+      const elements: Element[] = Array.from({ length: 20 ** 2 }, (_, i) => i)
       const tree = new MerkleTree(levels, elements)
       const edge1 = tree.getTreeEdge(200)
       const edge2 = tree.getTreeEdge(100)
-      const partialTree1 = new PartialMerkleTree(levels, edge1, elements.slice(edge1.edgeIndex), tree.root)
+      const partialTree = new PartialMerkleTree(levels, edge1, elements.slice(edge1.edgeIndex))
-      const partialTree2 = new PartialMerkleTree(levels, edge2, elements.slice(edge2.edgeIndex), tree.root)
+      partialTree.shiftEdge(edge2, elements.slice(edge2.edgeIndex, partialTree.edgeIndex))
-      partialTree1.shiftEdge(edge2, elements.slice(edge2.edgeIndex, partialTree1.edgeIndex))
+      assert.deepEqual(partialTree.path(110), tree.path(110))
-      assert.deepEqual(partialTree1.path(105), partialTree2.path(105))
     })
     it('should fail if new edge index is over current edge', () => {
-      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8, 9], 5)
+      const { fullTree, partialTree } = getTestTrees(10, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 4)
-      const newEdge = fullTree.getTreeEdge(6)
+      const newEdge = fullTree.getTreeEdge(4)
       const call = () => partialTree.shiftEdge(newEdge, [1, 2])
       should().throw(call, 'New edgeIndex should be smaller then 4')
     })
     it('should fail if elements length are incorrect', () => {
-      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8, 9], 5)
+      const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8, 9], 4)
-      const newEdge = fullTree.getTreeEdge(4)
+      const newEdge = fullTree.getTreeEdge(3)
       const call = () => partialTree.shiftEdge(newEdge, [1, 2])
-      should().throw(call, 'Elements length should be 2')
+      should().throw(call, 'Elements length should be 1')
     })
   })
   describe('#serialize', () => {
     it('should work', () => {
-      const { partialTree } = getTestTrees(5, [1, 2, 3, 4, 5, 6, 7, 8, 9], 6)
+      const { partialTree } = getTestTrees(5, [1, 2, 3, 4, 5, 6, 7, 8, 9], 5)
       const data = partialTree.serialize()
       const dst = PartialMerkleTree.deserialize(data)
       should().equal(partialTree.root, dst.root)
@@ -264,7 +263,7 @@ describe('PartialMerkleTree', () => {
   })
   describe('#toString', () => {
     it('should return correct stringified representation', () => {
-      const { partialTree } = getTestTrees(5, [1, 2, 3, 4, 5, 6, 7, 8, 9], 6)
+      const { partialTree } = getTestTrees(5, [1, 2, 3, 4, 5, 6, 7, 8, 9], 5)
       const str = partialTree.toString()
       const dst = PartialMerkleTree.deserialize(JSON.parse(str))
       should().equal(partialTree.root, dst.root)