Merge branch 'petscheit-refactors_proofFlags'

src/MerkleTree.ts

@@ -638,7 +638,7 @@ export class MerkleTree {
    * getProofFlags
    * @desc Returns list of booleans where proofs should be used instead of hashing.
    * Proof flags are used in the Solidity multiproof verifiers.
-   * @param {Buffer[]} leaves
+   * @param {Number[]|Buffer[]} leaves
    * @param {Buffer[]} proofs
    * @return {Boolean[]} - Boolean flags
    * @example
@@ -648,8 +648,18 @@ export class MerkleTree {
    *const proofFlags = tree.getProofFlags(leaves, proof)
    *```
    */
-  getProofFlags (leaves: Buffer[], proofs: Buffer[]):boolean[] {
-    let ids = leaves.map((el) => this._bufferIndexOf(this.leaves, el)).sort((a, b) => a === b ? 0 : a > b ? 1 : -1)
+  getProofFlags (leaves: any[], proofs: Buffer[]):boolean[] {
+    if (!Array.isArray(leaves) || leaves.length <= 0) {
+      throw new Error('Invalid Inputs!')
+    }
+
+    let ids
+    if (!isNaN(leaves[0])) {
+      ids = leaves.sort((a, b) => a === b ? 0 : a > b ? 1 : -1) // Indices where passed
+    } else {
+      ids = leaves.map((el) => this._bufferIndexOf(this.leaves, el)).sort((a, b) => a === b ? 0 : a > b ? 1 : -1)
+    }
+
     if (!ids.every((idx) => idx !== -1)) {
       throw new Error('Element does not exist in Merkle tree')
     }

test/merkletree.test.js

@@ -850,6 +850,29 @@ test('sha256 getMultiProof', t => {
   ])
 })
 
+test('sha256 getProofFlag with indices', t => {
+  t.plan(2)
+
+  const leaves = ['a', 'b', 'c', 'd'].map(sha256)
+
+  const tree = new MerkleTree(leaves, sha256, {sortPairs: true})
+
+  const root = tree.getHexRoot()
+  t.equal(root, '0x4c6aae040ffada3d02598207b8485fcbe161c03f4cb3f660e4d341e7496ff3b2')
+
+  const treeFlat = tree.getLayersFlat()
+
+  const indices = [1, 2]
+  const proof = tree.getMultiProof(treeFlat, indices)
+
+  const proofFlags = tree.getProofFlags(indices, proof)
+  t.deepEqual(proofFlags, [
+    false,
+    false,
+    true
+  ])
+})
+
 test('sha256 getMultiProof - statusim', t => {
   t.plan(5)