chore: some fixes and gas optimizations

This commit is contained in:
Richard Ramos 2024-09-06 16:12:19 -04:00
parent 9a1c1dc33d
commit 9251400051
No known key found for this signature in database
GPG Key ID: 1CE87DB518195760
2 changed files with 125 additions and 51 deletions

View File

@ -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);

View File

@ -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;