Update _eraseMemberships to upate roots

src/WakuRlnV2.sol

@@ -200,6 +200,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
 
     /// @notice Returns the list of recent roots, newest first
     function getRecentRoots() external view returns (uint256[HISTORY_SIZE] memory ordered) {
+        // refresh the latest root - needed if memberships were changed/erased but no new root stored yet
         if (!initialized) {
             return ordered; // empty array if no roots yet
         }
@@ -221,6 +222,12 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
         return recentRoots[index];
     }
 
+    /// @notice Returns the root of the Merkle tree that stores rate commitments of memberships
+    /// @return The root of the Merkle tree that stores rate commitments of memberships
+    function root() public view returns (uint256) {
+        return LazyIMT.root(merkleTree, MERKLE_TREE_DEPTH);
+    }
+
     /// @dev Register a membership (internal function)
     /// @param idCommitment The idCommitment of the membership
     /// @param rateLimit The rate limit of the membership
@@ -234,14 +241,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
         }
 
         // ✅ After updating or inserting, capture and store the new root
-        uint256 newRoot = LazyIMT.root(merkleTree, MERKLE_TREE_DEPTH);
+        _storeNewRoot(root());
-        _storeNewRoot(newRoot);
-    }
-
-    /// @notice Returns the root of the Merkle tree that stores rate commitments of memberships
-    /// @return The root of the Merkle tree that stores rate commitments of memberships
-    function root() external view returns (uint256) {
-        return LazyIMT.root(merkleTree, MERKLE_TREE_DEPTH);
     }
 
     /// @notice Returns the Merkle proof that a given membership is in the membership set
@@ -294,6 +294,7 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
         // eraseFromMembershipSet == false means lazy erasure.
         //  Only erase memberships from the memberships array (consume less gas).
         //  Merkle tree data will be overwritten when the correspondind index is reused.
+        bool treeModified = false;
         for (uint256 i = 0; i < idCommitmentsToErase.length; i++) {
             // Erase the membership from the memberships array in contract storage
             uint32 indexToErase = _eraseMembershipLazily(_msgSender(), idCommitmentsToErase[i]);
@@ -301,8 +302,13 @@ contract WakuRlnV2 is Initializable, OwnableUpgradeable, UUPSUpgradeable, Member
             // This does not affect the total rate limit control, or index reusal for new membership registrations.
             if (eraseFromMembershipSet) {
                 LazyIMT.update(merkleTree, 0, indexToErase);
+                treeModified = true;
             }
         }
+        // Record only the final root once per batch full clean-up call to avoid duplicates.
+        if (treeModified) {
+            _storeNewRoot(root());
+        }
     }
 
     /// @notice Withdraw any available deposit balance in tokens after a membership is erased

test/WakuRlnV2.t.sol

@@ -1644,4 +1644,63 @@ contract WakuRlnV2Test is Test {
         assertEq(w.getRootAt(0), root2);
         assertEq(w.getRootAt(1), root1);
     }
+
+    function test__RecentRoots_FullCleanupErasureUpdatesHistory() external {
+        uint32 rate = w.minMembershipRateLimit();
+        (, uint256 price) = w.priceCalculator().calculate(rate);
+
+        // Register three memberships (1,2,3), capturing roots after each
+        token.approve(address(w), price);
+        w.register(1, rate, noIdCommitmentsToErase);
+        uint256 r1 = w.root();
+
+        token.approve(address(w), price);
+        w.register(2, rate, noIdCommitmentsToErase);
+        uint256 r2 = w.root();
+
+        token.approve(address(w), price);
+        w.register(3, rate, noIdCommitmentsToErase);
+        uint256 r3 = w.root();
+
+        // Expire id 1 and erase with full clean-up (tree changes)
+        (,, uint256 gStart1, uint32 gDur1,,,,) = w.memberships(1);
+        vm.warp(gStart1 + gDur1 + 1);
+        uint256[] memory toErase = new uint256[](1);
+        toErase[0] = 1;
+        w.eraseMemberships(toErase, true);
+        uint256 r4 = w.root();
+        assertNotEq(r4, r3);
+
+        // Recent roots should be [r4, r3, r2, r1, 0]
+        {
+            uint256[HISTORY] memory recent = w.getRecentRoots();
+            assertEq(recent[0], r4);
+            assertEq(recent[1], r3);
+            assertEq(recent[2], r2);
+            assertEq(recent[3], r1);
+            assertEq(recent[4], 0);
+        }
+
+        // Now expire ids 2 and 3 and erase both in a single full clean-up call
+        (,, uint256 gStart3, uint32 gDur3,,,,) = w.memberships(3);
+        vm.warp(gStart3 + gDur3 + 1);
+        uint256[] memory batchErase = new uint256[](2);
+        batchErase[0] = 2;
+        batchErase[1] = 3;
+        // Capture the root before the batch to verify only one new history entry is pushed
+        uint256 preBatchRoot = w.root();
+        w.eraseMemberships(batchErase, true);
+        uint256 r5 = w.root();
+        assertNotEq(r5, preBatchRoot);
+
+        // Expect exactly one new entry for the batch (final root), so recent = [r5, r4, r3, r2, r1]
+        {
+            uint256[HISTORY] memory recent = w.getRecentRoots();
+            assertEq(recent[0], r5);
+            assertEq(recent[1], r4);
+            assertEq(recent[2], r3);
+            assertEq(recent[3], r2);
+            assertEq(recent[4], r1);
+        }
+    }
 }