Merge pull request #226 from privacy-scaling-explorations/refactor/in…

…dex-of

Handle `LeanIMT.indexOf` exception when a leaf does not exist

packages/imt.sol/contracts/README.md

@@ -9,8 +9,8 @@
     <a href="https://github.com/privacy-scaling-explorations/zk-kit">
         <img src="https://img.shields.io/badge/project-zk--kit-blue.svg?style=flat-square">
     </a>
-    <a href="https://github.com/privacy-scaling-explorations/zk-kit/tree/main/packages/imt.sol/LICENSE">
-        <img alt="NPM license" src="https://img.shields.io/npm/l/%40zk-kit%2Fsmt.sol?style=flat-square">
+    <a href="https://github.com/privacy-scaling-explorations/zk-kit/tree/main/packages/imt.sol/contracts/LICENSE">
+        <img alt="NPM license" src="https://img.shields.io/npm/l/%40zk-kit%2Fimt.sol?style=flat-square">
     </a>
     <a href="https://www.npmjs.com/package/@zk-kit/imt.sol">
         <img alt="NPM version" src="https://img.shields.io/npm/v/@zk-kit/imt.sol?style=flat-square" />
@@ -37,14 +37,15 @@
 > [!WARNING]  
 > If you are looking for the first version of this package, please visit this [link](https://github.com/privacy-scaling-explorations/zk-kit/tree/imt-v1/packages/incremental-merkle-tree.sol).
 
----
+> [!WARNING]  
+> These libraries have **not** been audited.
 
 ## Libraries
 
-✔️ [BinaryIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/BinaryIMT.sol) (Poseidon)\
-✔️ [QuinaryIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/QuinaryIMT.sol) (Poseidon)\
-✔️ [LazyIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/LazyIMT.sol) (Poseidon)\
-✔️ [LeanIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/LeanIMT.sol) (Poseidon)
+✔️ [BinaryIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/internal/InternalBinaryIMT.sol) (Poseidon)\
+✔️ [QuinaryIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/internal/InternalQuinaryIMT.sol) (Poseidon)\
+✔️ [LazyIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/internal/InternalLazyIMT.sol) (Poseidon)\
+✔️ [LeanIMT](https://github.com/privacy-scaling-explorations/zk-kit/blob/main/packages/imt.sol/contracts/internal/InternalLeanIMT.sol) (Poseidon)
 
 ## 🛠 Install
 

packages/imt.sol/contracts/internal/InternalLeanIMT.sol

@@ -321,8 +321,13 @@ library InternalLeanIMT {
     /// @dev Retrieves the index of a given leaf in the tree.
     /// @param self: A storage reference to the 'LeanIMTData' struct.
     /// @param leaf: The value of the leaf whose index is to be found.
-    /// @return The index of the specified leaf within the tree. If the leaf is not present, the function returns 0.
+    /// @return The index of the specified leaf within the tree. If the leaf is not present, the function
+    /// reverts with a custom error.
     function _indexOf(LeanIMTData storage self, uint256 leaf) internal view returns (uint256) {
+        if (self.leaves[leaf] == 0) {
+            revert LeafDoesNotExist();
+        }
+
         return self.leaves[leaf] - 1;
     }
 

packages/imt.sol/test/LeanIMT.ts

@@ -268,6 +268,7 @@ describe("LeanIMT", () => {
 
             expect(hasLeaf).to.equal(true)
         })
+
         it("Should return false because the node is not the tree", async () => {
             const hasLeaf = await leanIMTTest.has(2)
 
@@ -282,6 +283,7 @@ describe("LeanIMT", () => {
 
             expect(index).to.equal(0)
         })
+
         it("Should return the indices of the leaves", async () => {
             await leanIMTTest.insertMany([1, 2])
 
@@ -291,6 +293,14 @@ describe("LeanIMT", () => {
             expect(index1).to.equal(0)
             expect(index2).to.equal(1)
         })
+
+        it("Should throw a custom error if the leaf does not exist", async () => {
+            await leanIMTTest.insertMany([1, 2])
+
+            const transaction = leanIMTTest.indexOf(3)
+
+            await expect(transaction).to.be.revertedWithCustomError(leanIMT, "LeafDoesNotExist")
+        })
     })
 
     describe("# root", () => {