From 3d6356b59667eb0a7574d694b51c672fd90dd8a0 Mon Sep 17 00:00:00 2001
From: Sergei SMART <sergei.smart@gmail.com>
Date: Tue, 22 Feb 2022 19:42:49 +1000
Subject: [PATCH] add getTreeEdge method

---
 src/index.ts | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index f186f41..2c80a19 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -130,15 +130,23 @@ export default class MerkleTree {
     let elIndex = +index
     const pathElements: Element[] = []
     const pathIndices: number[] = []
+    const layerIndices: number[] = []
     for (let level = 0; level < this.levels; level++) {
       pathIndices[level] = elIndex % 2
-      pathElements[level] =
-        (elIndex ^ 1) < this._layers[level].length ? this._layers[level][elIndex ^ 1] : this._zeros[level]
+      const leafIndex = elIndex ^ 1
+      if (leafIndex < this._layers[level].length) {
+        pathElements[level] = this._layers[level][leafIndex]
+        layerIndices[level] = leafIndex
+      } else {
+        pathElements[level] = this._zeros[level]
+        layerIndices[level] = 0
+      }
       elIndex >>= 1
     }
     return {
       pathElements,
       pathIndices,
+      layerIndices,
     }
   }
 
@@ -156,6 +164,25 @@ export default class MerkleTree {
     }
   }
 
+  getTreeEdge(edgeElement: Element, index?: number) {
+    if (edgeElement === 'undefined') {
+      throw new Error('element is required')
+    }
+    let edgeIndex: number
+    if (!Number.isInteger(index)) {
+      index = -1
+      const leaves = this._layers[0]
+      index = leaves.indexOf(edgeElement)
+      edgeIndex = index
+    }
+
+    if (index <= -1) {
+      return []
+    }
+    const edgePath = this.path(index)
+    return { edgePath, edgeElement, edgeIndex }
+  }
+
   /**
    * Returns a copy of non-zero tree elements.
    */