added strictAdd & strictRemove (#93)

dl-solarity · Mar 14, 2024 · f54793c · f54793c
1 parent 06971b1
commit f54793c
Show file tree

Hide file tree

Showing 4 changed files with 161 additions and 2 deletions.
diff --git a/contracts/libs/arrays/SetHelper.sol b/contracts/libs/arrays/SetHelper.sol
@@ -42,6 +42,35 @@ library SetHelper {
         }
     }
 
+    /**
+     * @notice The function for the strict insertion of an array of elements into the address set
+     * @param set the set to insert the elements into
+     * @param array_ the elements to be inserted
+     */
+    function strictAdd(EnumerableSet.AddressSet storage set, address[] memory array_) internal {
+        for (uint256 i = 0; i < array_.length; i++) {
+            require(set.add(array_[i]), "SetHelper: element already exists");
+        }
+    }
+
+    /**
+     * @notice The function for the strict insertion of an array of elements into the uint256 set
+     */
+    function strictAdd(EnumerableSet.UintSet storage set, uint256[] memory array_) internal {
+        for (uint256 i = 0; i < array_.length; i++) {
+            require(set.add(array_[i]), "SetHelper: element already exists");
+        }
+    }
+
+    /**
+     * @notice The function for the strict insertion of an array of elements into the string set
+     */
+    function strictAdd(StringSet.Set storage set, string[] memory array_) internal {
+        for (uint256 i = 0; i < array_.length; i++) {
+            require(set.add(array_[i]), "SetHelper: element already exists");
+        }
+    }
+
     /**
      * @notice The function to remove an array of elements from the address set
      * @param set the set to remove the elements from
@@ -70,4 +99,33 @@ library SetHelper {
             set.remove(array_[i]);
         }
     }
+
+    /**
+     * @notice The function for the strict removal of an array of elements from the address set
+     * @param set the set to remove the elements from
+     * @param array_ the elements to be removed
+     */
+    function strictRemove(EnumerableSet.AddressSet storage set, address[] memory array_) internal {
+        for (uint256 i = 0; i < array_.length; i++) {
+            require(set.remove(array_[i]), "SetHelper: no such element");
+        }
+    }
+
+    /**
+     * @notice The function for the strict removal of an array of elements from the uint256 set
+     */
+    function strictRemove(EnumerableSet.UintSet storage set, uint256[] memory array_) internal {
+        for (uint256 i = 0; i < array_.length; i++) {
+            require(set.remove(array_[i]), "SetHelper: no such element");
+        }
+    }
+
+    /**
+     * @notice The function for the strict removal of an array of elements from the string set
+     */
+    function strictRemove(StringSet.Set storage set, string[] memory array_) internal {
+        for (uint256 i = 0; i < array_.length; i++) {
+            require(set.remove(array_[i]), "SetHelper: no such element");
+        }
+    }
 }
diff --git a/contracts/mock/libs/arrays/SetHelperMock.sol b/contracts/mock/libs/arrays/SetHelperMock.sol
@@ -30,6 +30,18 @@ contract SetHelperMock {
         stringSet.add(arr_);
     }
 
+    function strictAddToAddressSet(address[] memory arr_) external {
+        addressSet.strictAdd(arr_);
+    }
+
+    function strictAddToUintSet(uint256[] memory arr_) external {
+        uintSet.strictAdd(arr_);
+    }
+
+    function strictAddToStringSet(string[] memory arr_) external {
+        stringSet.strictAdd(arr_);
+    }
+
     function removeFromAddressSet(address[] memory arr_) external {
         addressSet.remove(arr_);
     }
@@ -42,6 +54,18 @@ contract SetHelperMock {
         stringSet.remove(arr_);
     }
 
+    function strictRemoveFromAddressSet(address[] memory arr_) external {
+        addressSet.strictRemove(arr_);
+    }
+
+    function strictRemoveFromUintSet(uint256[] memory arr_) external {
+        uintSet.strictRemove(arr_);
+    }
+
+    function strictRemoveFromStringSet(string[] memory arr_) external {
+        stringSet.strictRemove(arr_);
+    }
+
     function getAddressSet() external view returns (address[] memory) {
         return addressSet.values();
     }

diff --git a/package-lock.json b/package-lock.json
diff --git a/test/libs/arrays/SetHelper.test.ts b/test/libs/arrays/SetHelper.test.ts
@@ -45,6 +45,40 @@ describe("SetHelper", () => {
     });
   });
 
+  describe("strictAdd", () => {
+    it("should not strict add to address set if element already exists", async () => {
+      await expect(mock.strictAddToAddressSet([FIRST.address, FIRST.address])).to.be.revertedWith(
+        "SetHelper: element already exists",
+      );
+    });
+
+    it("should not strict add to uint set if element already exists", async () => {
+      await expect(mock.strictAddToUintSet([1, 1])).to.be.revertedWith("SetHelper: element already exists");
+    });
+
+    it("should strict add to string set if element already exists", async () => {
+      await expect(mock.strictAddToStringSet(["1", "1"])).to.be.revertedWith("SetHelper: element already exists");
+    });
+
+    it("should strict add to address set", async () => {
+      await mock.strictAddToAddressSet([FIRST.address, SECOND.address]);
+
+      expect(await mock.getAddressSet()).to.deep.equal([FIRST.address, SECOND.address]);
+    });
+
+    it("should strict add to uint set", async () => {
+      await mock.strictAddToUintSet([1]);
+
+      expect(await mock.getUintSet()).to.deep.equal([1n]);
+    });
+
+    it("should strict add to string set", async () => {
+      await mock.strictAddToStringSet(["1", "2", "3"]);
+
+      expect(await mock.getStringSet()).to.deep.equal(["1", "2", "3"]);
+    });
+  });
+
   describe("remove", () => {
     it("should remove from address set", async () => {
       await mock.addToAddressSet([FIRST.address, SECOND.address]);
@@ -67,4 +101,47 @@ describe("SetHelper", () => {
       expect(await mock.getStringSet()).to.deep.equal(["3", "2"]);
     });
   });
+
+  describe("remove", () => {
+    it("should not strict remove from address set if no such element", async () => {
+      await mock.strictAddToAddressSet([FIRST.address, SECOND.address]);
+
+      await expect(mock.strictRemoveFromAddressSet([SECOND.address, SECOND.address])).to.be.revertedWith(
+        "SetHelper: no such element",
+      );
+    });
+
+    it("should not strict remove from uint set if no such element", async () => {
+      await mock.strictAddToUintSet([1]);
+
+      await expect(mock.strictRemoveFromUintSet([1, 1])).to.be.revertedWith("SetHelper: no such element");
+    });
+
+    it("should not strict remove from string set if no such element", async () => {
+      await mock.strictAddToStringSet(["1", "2", "3"]);
+
+      await expect(mock.strictRemoveFromStringSet(["1", "1"])).to.be.revertedWith("SetHelper: no such element");
+    });
+
+    it("should strict remove from address set", async () => {
+      await mock.strictAddToAddressSet([FIRST.address, SECOND.address]);
+      await mock.strictRemoveFromAddressSet([SECOND.address]);
+
+      expect(await mock.getAddressSet()).to.deep.equal([FIRST.address]);
+    });
+
+    it("should strict remove from uint set", async () => {
+      await mock.strictAddToUintSet([1]);
+      await mock.strictRemoveFromUintSet([1]);
+
+      expect(await mock.getUintSet()).to.deep.equal([]);
+    });
+
+    it("should strict remove from string set", async () => {
+      await mock.strictAddToStringSet(["1", "2", "3"]);
+      await mock.strictRemoveFromStringSet(["1"]);
+
+      expect(await mock.getStringSet()).to.deep.equal(["3", "2"]);
+    });
+  });
 });