From 9251400051a2774ffb07dd3545fca78cc3d9031d Mon Sep 17 00:00:00 2001 From: Richard Ramos Date: Fri, 6 Sep 2024 16:12:19 -0400 Subject: [PATCH] chore: some fixes and gas optimizations --- src/Membership.sol | 88 ++++++++++++++++++++++++++------------------ test/WakuRlnV2.t.sol | 88 ++++++++++++++++++++++++++++++++++++-------- 2 files changed, 125 insertions(+), 51 deletions(-) diff --git a/src/Membership.sol b/src/Membership.sol index 7fa48b4..e7d5744 100644 --- a/src/Membership.sol +++ b/src/Membership.sol @@ -202,33 +202,39 @@ contract Membership { revert InvalidRateLimit(); } + // Storing in local variable to not access the storage frequently + // And we're using/modifying these variables in each iteration + uint256 _head = head; + uint256 _tail = tail; + uint256 _totalRateLimitPerEpoch = totalRateLimitPerEpoch; + uint32 _maxTotalRateLimitPerEpoch = maxTotalRateLimitPerEpoch; + // Determine if we exceed the total rate limit - if (totalRateLimitPerEpoch + _rateLimit > maxTotalRateLimitPerEpoch) { - if (head == 0) revert ExceedAvailableMaxRateLimitPerEpoch(); // List is empty + if (_totalRateLimitPerEpoch + _rateLimit > _maxTotalRateLimitPerEpoch) { + if (_head == 0) revert ExceedAvailableMaxRateLimitPerEpoch(); // List is empty // Attempt to free expired membership slots - while (totalRateLimitPerEpoch + _rateLimit > maxTotalRateLimitPerEpoch) { + while (_totalRateLimitPerEpoch + _rateLimit > _maxTotalRateLimitPerEpoch && _head != 0) { // Determine if there are any available spot in the membership map // by looking at the oldest membership. If it's expired, we can free it - MembershipInfo memory oldestMembership = members[head]; + MembershipInfo memory oldestMembership = members[_head]; if ( - oldestMembership.holder != address(0) // membership has a holder - && _isExpired( - oldestMembership.gracePeriodStartDate, - oldestMembership.gracePeriod, - oldestMembership.numberOfPeriods - ) + _isExpired( + oldestMembership.gracePeriodStartDate, + oldestMembership.gracePeriod, + oldestMembership.numberOfPeriods + ) ) { - emit MemberExpired(head, oldestMembership.userMessageLimit, oldestMembership.index); + emit MemberExpired(_head, oldestMembership.userMessageLimit, oldestMembership.index); // Deduct the expired membership rate limit - totalRateLimitPerEpoch -= oldestMembership.userMessageLimit; + _totalRateLimitPerEpoch -= oldestMembership.userMessageLimit; // Remove the element from the list - delete members[head]; + delete members[_head]; // Promote the next oldest membership to oldest - head = oldestMembership.next; + _head = oldestMembership.next; // Move balance from expired membership to holder balance balancesToWithdraw[oldestMembership.holder][oldestMembership.token] += oldestMembership.amount; @@ -240,24 +246,22 @@ contract Membership { } // Ensure new head and tail are pointing to the correct memberships - if (head != 0) { - members[head].prev = 0; + if (_head != 0) { + members[_head].prev = 0; } else { - tail = 0; + _tail = 0; } } - uint256 prev = 0; - if (tail != 0) { - MembershipInfo storage latestMembership = members[tail]; - latestMembership.next = _idCommitment; - prev = tail; + if (_tail != 0) { + members[_tail].next = _idCommitment; } else { // First item - head = _idCommitment; + _head = _idCommitment; } - totalRateLimitPerEpoch += _rateLimit; + // Adding the rate limit of the new registration + _totalRateLimitPerEpoch += _rateLimit; // Reuse available slots from previously removed expired memberships uint256 arrLen = availableExpiredIndices.length; @@ -269,6 +273,7 @@ contract Membership { index = commitmentIndex; } + totalRateLimitPerEpoch = _totalRateLimitPerEpoch; members[_idCommitment] = MembershipInfo({ holder: _sender, gracePeriodStartDate: block.timestamp + (uint256(billingPeriod) * uint256(_numberOfPeriods)), @@ -278,10 +283,10 @@ contract Membership { amount: _amount, userMessageLimit: _rateLimit, next: 0, // It's the newest value, so point to nowhere - prev: prev, + prev: _tail, index: index }); - + head = _head; tail = _idCommitment; } @@ -300,29 +305,40 @@ contract Membership { uint256 newGracePeriodStartDate = block.timestamp + (uint256(billingPeriod) * uint256(mdetails.numberOfPeriods)); - uint256 mdetailsNext = mdetails.next; - uint256 mdetailsPrev = mdetails.prev; + uint256 next = mdetails.next; + uint256 prev = mdetails.prev; + uint256 _tail = tail; + uint256 _head = head; // Remove current membership references - if (mdetailsPrev != 0) { - members[mdetailsPrev].next = mdetailsNext; + if (prev != 0) { + members[prev].next = next; } else { - head = mdetailsNext; + _head = next; } - if (mdetailsNext != 0) { - members[mdetailsNext].prev = mdetailsPrev; + if (next != 0) { + members[next].prev = prev; } else { - tail = mdetailsPrev; + _tail = prev; } // Move membership to the end (since it will be the newest) mdetails.next = 0; - mdetails.prev = tail; + mdetails.prev = _tail; mdetails.gracePeriodStartDate = newGracePeriodStartDate; mdetails.gracePeriod = gracePeriod; - members[tail].next = _idCommitment; + // Link previous tail with membership that was just extended + if (_tail != 0) { + members[_tail].next = _idCommitment; + } else { + // There are no other items in the list. + // The head will become the extended commitment + _head = _idCommitment; + } + + head = _head; tail = _idCommitment; emit MemberExtended(_idCommitment, mdetails.userMessageLimit, mdetails.index, newGracePeriodStartDate); diff --git a/test/WakuRlnV2.t.sol b/test/WakuRlnV2.t.sol index f2274a0..24c9b92 100644 --- a/test/WakuRlnV2.t.sol +++ b/test/WakuRlnV2.t.sol @@ -328,6 +328,38 @@ contract WakuRlnV2Test is Test { w.extend(commitmentsToExtend); } + function test__ValidRegistrationExtendSingleMembership(uint32 userMessageLimit, uint32 numberOfPeriods) external { + vm.pauseGasMetering(); + uint256 idCommitment = 2; + vm.assume(numberOfPeriods > 0 && numberOfPeriods < 100); + (, uint256 price) = w.priceCalculator().calculate(userMessageLimit, numberOfPeriods); + vm.assume( + userMessageLimit >= w.minRateLimitPerMembership() && userMessageLimit <= w.maxRateLimitPerMembership() + ); + vm.assume(w.isValidUserMessageLimit(userMessageLimit)); + vm.resumeGasMetering(); + + w.register{ value: price }(idCommitment, userMessageLimit, numberOfPeriods); + (,,,, uint256 gracePeriodStartDate,,,,,) = w.members(idCommitment); + + vm.warp(gracePeriodStartDate); + + uint256[] memory commitmentsToExtend = new uint256[](1); + commitmentsToExtend[0] = idCommitment; + + // Extend the membership + vm.expectEmit(true, false, false, false); // only check the first parameter of the event (the idCommitment) + emit Membership.MemberExtended(idCommitment, 0, 0, 0); + w.extend(commitmentsToExtend); + + // Verify list order is correct + assertEq(w.tail(), idCommitment); + assertEq(w.head(), idCommitment); + (uint256 prev, uint256 next,,,,,,,,) = w.members(idCommitment); + assertEq(next, 0); + assertEq(prev, 0); + } + function test__ValidRegistrationExpiry(uint32 userMessageLimit, uint32 numberOfPeriods) external { vm.pauseGasMetering(); uint256 idCommitment = 2; @@ -473,39 +505,63 @@ contract WakuRlnV2Test is Test { vm.stopPrank(); vm.resumeGasMetering(); - (, uint256 price) = w.priceCalculator().calculate(1, 10); - w.register{ value: price }(1, 1, 10); + (, uint256 priceA) = w.priceCalculator().calculate(1, 10); + w.register{ value: priceA }(1, 1, 10); vm.warp(block.timestamp + 100); - w.register{ value: price }(2, 1, 10); + w.register{ value: priceA }(2, 1, 10); vm.warp(block.timestamp + 100); uint256 expirationDate = w.expirationDate(2); vm.warp(expirationDate); - w.register{ value: price }(3, 1, 10); + w.register{ value: priceA }(3, 1, 10); // Make sure only the first 2 memberships are expired assertTrue(w.isExpired(1)); assertTrue(w.isExpired(2)); assertFalse(w.isExpired(3) || w.isGracePeriod(3)); + (,,,,,,, uint32 index1,,) = w.members(1); + (,,,,,,, uint32 index2,,) = w.members(2); + // Attempt to register a membership that will require to expire 2 memberships // Currently there is 2 available, and we want to register 4 // If we remove first membership, we'll have 3 available // If we also remove the second, we'll have 4 available vm.expectEmit(true, false, false, false); - emit Membership.MemberExpired(1, 0, 0); + emit Membership.MemberExpired(1, 0, 0); vm.expectEmit(true, false, false, false); emit Membership.MemberExpired(2, 0, 0); - (, price) = w.priceCalculator().calculate(4, 10); - w.register{ value: price }(4, 4, 10); + (, uint256 priceB) = w.priceCalculator().calculate(4, 10); + w.register{ value: priceB }(4, 4, 10); - // TODO: validate reuse of index - // TODO: validate balance - // TODO: check that the expired memberships are gone and membership 3 is still there - // TODO: check the usermessage limits (total) - // TODO: check that it reused the index of the one gone - // TODO: check head and tail are correct and next and prev - // TODO: there should be at least a single index available - // TODO: validate balance + // idCommitment4 will use the last removed index available (since we push to an array) + (,,,,,,, uint32 index4,,) = w.members(4); + assertEq(index4, index2); + + // the index of the first removed membership is still available for further registrations + assertEq(index1, w.availableExpiredIndices(0)); + + // The previous expired memberships should have been erased + (,,,,,,,, address holder,) = w.members(1); + assertEq(holder, address(0)); + (,,,,,,,, holder,) = w.members(2); + assertEq(holder, address(0)); + + // The total rate limit used should be those from idCommitment 3 and 4 + assertEq(5, w.totalRateLimitPerEpoch()); + + // There should only be 2 memberships, the non expired and the new one + assertEq(w.head(), 3); + assertEq(w.tail(), 4); + (uint256 prev, uint256 next,,,,,,,,) = w.members(3); + assertEq(prev, 0); + assertEq(next, 4); + (prev, next,,,,,,,,) = w.members(4); + assertEq(prev, 3); + assertEq(next, 0); + + // The balance available for withdrawal should match the amount of the expired membership + uint256 availableBalance = w.balancesToWithdraw(address(this), address(0)); + assertEq(availableBalance, priceA * 2); } function test__RegistrationWhenMaxRateLimitIsReachedAndMultipleExpiredMembersAvailableButTheyDontHaveEnoughRateLimit( @@ -513,6 +569,8 @@ contract WakuRlnV2Test is Test { external { } + function test__indexReuse() external { } + function test__RemoveExpiredMemberships(uint32 userMessageLimit, uint32 numberOfPeriods) external { vm.pauseGasMetering(); uint256 idCommitment = 2;