From febe67e8ea7b7e220af58271be3bb9ef09ff2292 Mon Sep 17 00:00:00 2001
From: Miguel Mota <hello@miguelmota.com>
Date: Mon, 6 Dec 2021 15:07:54 -0800
Subject: [PATCH] Add method to add additional leaves

---
 src/MerkleTree.ts       | 30 +++++++++++++++++++++++++++
 test/MerkleTree.test.js | 45 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/src/MerkleTree.ts b/src/MerkleTree.ts
index 40455cf..f024032 100644
--- a/src/MerkleTree.ts
+++ b/src/MerkleTree.ts
@@ -91,6 +91,10 @@ export class MerkleTree extends Base {
     this.duplicateOdd = !!options.duplicateOdd
 
     this.hashFn = this._bufferifyFn(hashFn)
+    this.processLeaves(leaves)
+  }
+
+  private processLeaves (leaves: TLeaf[]) {
     if (this.hashLeaves) {
       leaves = leaves.map(x => this.hashFn(x))
     }
@@ -208,6 +212,20 @@ export class MerkleTree extends Base {
     return this.leaves[index]
   }
 
+  // TODO: docs
+  getLeafIndex (target: Buffer):number {
+    target = this.bufferify(target)
+    const leaves = this.getLeaves()
+    for (let i = 0; i < leaves.length; i++) {
+      const leaf = leaves[i]
+      if (leaf.equals(target)) {
+        return i
+      }
+    }
+
+    return -1
+  }
+
   // TODO: docs
   getLeafCount (): number {
     return this.leaves.length
@@ -993,6 +1011,18 @@ export class MerkleTree extends Base {
   toString ():string {
     return this._toTreeString()
   }
+
+  // TODO: docs
+  addLeaf (leaf: TLeaf, shouldHash: boolean = false) {
+    const leaves = this.leaves
+    leaves.push(leaf)
+    this.processLeaves(this.leaves)
+  }
+
+  // TODO: docs
+  addLeaves (leaves: TLeaf[], shouldHash: boolean = false) {
+    this.processLeaves(leaves)
+  }
 }
 
 if (typeof window !== 'undefined') {
diff --git a/test/MerkleTree.test.js b/test/MerkleTree.test.js
index f91ea88..09f7b4c 100644
--- a/test/MerkleTree.test.js
+++ b/test/MerkleTree.test.js
@@ -1008,6 +1008,19 @@ test('fillDefaultHashes', t => {
   t.equal(tree.getHexRoot(), '0x11f470d712bb3a84f0b01cb7c73493ec7d06eda480f567c99b9a6dc773679a72')
 })
 
+test('getLeafIndex', t => {
+  t.plan(5)
+
+  const leaves = ['a', 'b', 'c'].map(x => keccak256(Buffer.from(x)))
+  const tree = new MerkleTree(leaves, sha256)
+
+  t.equal(tree.getHexRoot(), '0x311d2e46f49b15fff8b746b74ad57f2cc9e0d9939fda94387141a2d3fdf187ae')
+  t.equal(tree.getLeafIndex(leaves[1]), 1)
+  t.equal(tree.getLeafIndex(keccak256(Buffer.from('b'))), 1)
+  t.equal(tree.getLeafIndex(leaves[2]), 2)
+  t.equal(tree.getLeafIndex(Buffer.from('invalid')), -1)
+})
+
 test('getleafCount', t => {
   t.plan(1)
 
@@ -1030,6 +1043,38 @@ test('getleaf', t => {
   t.deepEqual(tree.getLeaf(3), Buffer.from([]))
 })
 
+test('addLeaf', t => {
+  t.plan(2)
+
+  const leaves = ['a', 'b', 'c'].map(x => keccak256(Buffer.from(x)))
+  const tree = new MerkleTree([], sha256)
+  tree.addLeaf(leaves[0])
+  tree.addLeaf(leaves[1])
+  tree.addLeaf(leaves[2])
+
+  t.equal(tree.getHexRoot(), '0x311d2e46f49b15fff8b746b74ad57f2cc9e0d9939fda94387141a2d3fdf187ae')
+  t.deepEqual(tree.getHexLeaves(), [
+    '0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb',
+    '0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510',
+    '0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2'
+  ])
+})
+
+test('addLeaves', t => {
+  t.plan(2)
+
+  const leaves = ['a', 'b', 'c'].map(x => keccak256(Buffer.from(x)))
+  const tree = new MerkleTree([], sha256)
+  tree.addLeaves(leaves)
+
+  t.equal(tree.getHexRoot(), '0x311d2e46f49b15fff8b746b74ad57f2cc9e0d9939fda94387141a2d3fdf187ae')
+  t.deepEqual(tree.getHexLeaves(), [
+    '0x3ac225168df54212a25c1c01fd35bebfea408fdac2e31ddd6f80a4bbf9a5f1cb',
+    '0xb5553de315e0edf504d9150af82dafa5c4667fa618ed0a6f19c69b41166c5510',
+    '0x0b42b6393c1f53060fe3ddbfcd7aadcca894465a5a438f69c87d790b2299b9b2'
+  ])
+})
+
 test('resetTree', t => {
   t.plan(2)