From 125c453a1308dfc0fe85a82b33bcddd5847390dd Mon Sep 17 00:00:00 2001
From: Miguel Mota <hello@miguelmota.com>
Date: Thu, 2 Dec 2021 13:47:33 -0800
Subject: [PATCH] Add count methods

---
 package.json            |  1 +
 src/MerkleTree.ts       | 29 +++++++++++++++++++++++++++++
 test/MerkleTree.test.js | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+)

diff --git a/package.json b/package.json
index 79eee4e..1060f5b 100644
--- a/package.json
+++ b/package.json
@@ -11,6 +11,7 @@
   "scripts": {
     "test": "tape test/*.js",
     "clean": "rimraf dist",
+    "dev": "tsc -w",
     "build": "npm run clean && tsc && npm run build:browser",
     "build:browser": "browserify -t [ babelify --presets [ @babel/preset-env ] ] dist/index.js | uglifyjs > merkletree.js",
     "lint": "standardx --fix src/*.ts test/*.js",
diff --git a/src/MerkleTree.ts b/src/MerkleTree.ts
index 6daa5d5..40455cf 100644
--- a/src/MerkleTree.ts
+++ b/src/MerkleTree.ts
@@ -199,6 +199,20 @@ export class MerkleTree extends Base {
     return this.leaves
   }
 
+  // TODO: docs
+  getLeaf (index: number):Buffer {
+    if (index < 0 || index > this.leaves.length - 1) {
+      return Buffer.from([])
+    }
+
+    return this.leaves[index]
+  }
+
+  // TODO: docs
+  getLeafCount (): number {
+    return this.leaves.length
+  }
+
   /**
    * getHexLeaves
    * @desc Returns array of leaves of Merkle Tree as hex strings.
@@ -329,6 +343,11 @@ export class MerkleTree extends Base {
     return this.getLayersFlat().map(x => this.bufferToHex(x))
   }
 
+  // TODO: docs
+  getLayerCount ():number {
+    return this.getLayers().length
+  }
+
   /**
    * getRoot
    * @desc Returns the Merkle root hash as a Buffer.
@@ -339,6 +358,10 @@ export class MerkleTree extends Base {
    *```
    */
   getRoot ():Buffer {
+    if (this.layers.length === 0) {
+      return Buffer.from([])
+    }
+
     return this.layers[this.layers.length - 1][0] || Buffer.from([])
   }
 
@@ -917,6 +940,12 @@ export class MerkleTree extends Base {
     return t.getMultiProof(tree, indices)
   }
 
+  // TODO: docs
+  resetTree ():void {
+    this.leaves = []
+    this.layers = []
+  }
+
   /**
    * getPairNode
    * @desc Returns the node at the index for given layer.
diff --git a/test/MerkleTree.test.js b/test/MerkleTree.test.js
index 682ed3d..f91ea88 100644
--- a/test/MerkleTree.test.js
+++ b/test/MerkleTree.test.js
@@ -1007,3 +1007,38 @@ test('fillDefaultHashes', t => {
 
   t.equal(tree.getHexRoot(), '0x11f470d712bb3a84f0b01cb7c73493ec7d06eda480f567c99b9a6dc773679a72')
 })
+
+test('getleafCount', t => {
+  t.plan(1)
+
+  const leaves = ['a', 'b', 'c'].map(x => keccak256(Buffer.from(x)))
+  const tree = new MerkleTree(leaves, sha256)
+
+  t.equal(tree.getLeafCount(), 3)
+})
+
+test('getleaf', t => {
+  t.plan(5)
+
+  const leaves = ['a', 'b', 'c'].map(x => keccak256(Buffer.from(x)))
+  const tree = new MerkleTree(leaves, sha256)
+
+  t.deepEqual(tree.getLeaf(-1), Buffer.from([]))
+  t.deepEqual(tree.getLeaf(0), leaves[0])
+  t.deepEqual(tree.getLeaf(1), leaves[1])
+  t.deepEqual(tree.getLeaf(2), leaves[2])
+  t.deepEqual(tree.getLeaf(3), Buffer.from([]))
+})
+
+test('resetTree', t => {
+  t.plan(2)
+
+  const leaves = ['a', 'b', 'c'].map(x => keccak256(Buffer.from(x)))
+  const tree = new MerkleTree(leaves, sha256)
+
+  t.equal(tree.getLeafCount(), 3)
+
+  tree.resetTree()
+
+  t.equal(tree.getLeafCount(), 0)
+})