feat: track balances and withdraw

contracts/Membership.sol

@@ -15,15 +15,18 @@ error ExceedMaxRateLimitPerEpoch();
 error NotInGracePeriod(uint256 membershipMapIdx);
 error NotHolder(uint256 membershipMapIdx);
 
+error InsufficientBalance();
+error FailedTransfer();
+
 contract Membership {
     using SafeERC20 for IERC20;
 
     // TODO START - add owned setters to all these variables
 
     IPriceCalculator public priceCalculator;
-    uint256 public maxTotalRateLimitPerEpoch;
-    uint16 public maxRateLimitPerMembership;
-    uint16 public minRateLimitPerMembership;
+    uint32 public maxTotalRateLimitPerEpoch; // TO-ASK: what's the theoretical maximum rate limit per epoch we could set? uint32 accepts a max of 4292967295
+    uint16 public maxRateLimitPerMembership; // TO-ASK: what's the theoretical maximum rate limit per epoch a single membership can have? this accepts 65535
+    uint16 public minRateLimitPerMembership; // TO-ASK: what's the theoretical largest minimum rate limit per epoch a single membership can have? this accepts a minimum from 0 to 65535
 
     // TO-ASK: what happens with existing memberships if
     // the expiration term and grace period are updated?
@@ -32,6 +35,15 @@ contract Membership {
 
     // TODO END
 
+    // TO-ASK: is it possible that in the future we change the
+    // token used by the contract? if yes, then it makes sense to
+    // have this balance as a mapping. If not, we can simplify this
+    // mapping and also remove the token setter and the `token`
+    // attribute from the MembershipDetails
+
+    // holder ->  token -> balance
+    mapping(address => mapping(address => uint)) public expiredBalances;
+
     enum MembershipStatus {
         // TODO use in getter to determine state of membership?
         NonExistent,
@@ -49,15 +61,18 @@ contract Membership {
     uint256 public tail = 0;
     uint256 private nextID = 0;
 
+    // TODO: associate membership details with commitment
+
     struct MembershipDetails {
-        address holder;
-        // TODO: should we store the commitment?
+        // Double linked list pointers
+        uint256 prev; // index of the previous membership
+        uint256 next; // index of the next membership
+        // Membership data
         uint256 expirationDate;
-        uint256 rateLimit;
-        address token;
         uint256 amount;
-        uint prev; // index of the previous membership
-        uint next; // index of the next membership
+        uint16 rateLimit;
+        address holder;
+        address token;
     }
 
     // TODO: should it contain the commitment?
@@ -67,11 +82,11 @@ contract Membership {
 
     function __Membership_init(
         address _priceCalculator,
-        uint _maxTotalRateLimitPerEpoch,
+        uint32 _maxTotalRateLimitPerEpoch,
         uint16 _maxRateLimitPerMembership,
         uint16 _minRateLimitPerMembership,
-        uint _expirationTerm,
-        uint _gracePeriod
+        uint256 _expirationTerm,
+        uint256 _gracePeriod
     ) internal {
         priceCalculator = IPriceCalculator(_priceCalculator);
         maxTotalRateLimitPerEpoch = _maxTotalRateLimitPerEpoch;
@@ -84,15 +99,15 @@ contract Membership {
     function registerMembership(
         address _sender,
         uint256[] memory commitments,
-        uint _rateLimit
+        uint16 _rateLimit
     ) internal {
         // TODO: for each commitment
-        (address token, uint amount) = priceCalculator.calculate(_rateLimit);
+        (address token, uint256 amount) = priceCalculator.calculate(_rateLimit);
         acquireRateLimit(_sender, commitments, _rateLimit, token, amount);
-        transferMembershipFees(_sender, token, amount * _rateLimit * commitments.length);
+        transferFees(_sender, token, amount * _rateLimit * commitments.length);
     }
 
-    function transferMembershipFees(address _from, address _token, uint _amount) internal {
+    function transferFees(address _from, address _token, uint256 _amount) internal {
         if (_token == address(0)) {
             if (msg.value != _amount) revert IncorrectAmount();
         } else {
@@ -104,7 +119,7 @@ contract Membership {
     function acquireRateLimit(
         address _sender,
         uint256[] memory _commitments,
-        uint256 _rateLimit,
+        uint16 _rateLimit,
         address _token,
         uint256 _amount
     ) internal {
@@ -131,15 +146,18 @@ contract Membership {
                 uint256 nextOld = oldestMembershipDetails.next;
                 if (nextOld != 0) memberships[nextOld].prev = 0;
 
-                delete memberships[head];
-
                 if (tail == head) {
                     // TODO: test this
                     tail = 0;
                 }
                 head = nextOld;
 
-                // TODO: move balance from expired to the current holder
+                // Move balance from expired membership to holder balance
+                expiredBalances[oldestMembershipDetails.holder][
+                    oldestMembershipDetails.token
+                ] += oldestMembershipDetails.amount;
+
+                delete memberships[head];
             } else {
                 revert ExceedMaxRateLimitPerEpoch();
             }
@@ -228,19 +246,30 @@ contract Membership {
 
     function freeExpiredMemberships(uint256[] calldata expiredMemberships) public {
         // TODO: user can pass a list of expired memberships and free them
-        // Might be useful because then offchain the user can determine which 
+        // Might be useful because then offchain the user can determine which
         // expired memberships slots are available, and proceed to free them.
         // This might be cheaper than the `while` loop used when registering
         // memberships, although easily solved by having a function that receives
-        // the list of memberships to free, and the information for the new 
+        // the list of memberships to free, and the information for the new
         // membership to register
     }
 
-    function withdraw() public {
-        // TODO:
-    }
+    // TODO: expire owned memberships?
 
-    // TODO - keep track of balances, use msg.sender
+    function withdraw(address token) public {
+        // TODO: getSender()  
+        uint256 amount = expiredBalances[msg.sender][token];
+        require(amount > 0, "Insufficient balance");
+
+        expiredBalances[msg.sender][token] = 0;
+        if (token == address(0)) {
+            // ETH
+            (bool success, ) = msg.sender.call{value: amount}("");
+            require(success, "eth transfer failed");
+        } else {
+            IERC20(token).safeTransfer(msg.sender, amount);
+        }
+    }
 
     function getOldestMembership() public view returns (MembershipDetails memory) {
         return memberships[head];