From 6303b54e2af0d729f1f7597e1494b6638048c982 Mon Sep 17 00:00:00 2001
From: Eric Mastro <eric.mastro@gmail.com>
Date: Tue, 15 Nov 2022 16:53:31 +1100
Subject: [PATCH] Clear active slots by host

Allow for clearing of active slots by host, by incrementing an mapping index. The new index points to a fresh instance of EnumerableSet, effectively wiping it clean.
---
 contracts/Marketplace.sol | 34 ++++++++++++++++++++++++++--------
 test/Marketplace.test.js  |  3 +--
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/contracts/Marketplace.sol b/contracts/Marketplace.sol
index 9e92634..b788aff 100644
--- a/contracts/Marketplace.sol
+++ b/contracts/Marketplace.sol
@@ -19,7 +19,8 @@ contract Marketplace is Collateral, Proofs {
   mapping(RequestId => RequestContext) private requestContexts;
   mapping(SlotId => Slot) private slots;
   mapping(address => EnumerableSet.Bytes32Set) private activeRequests;
-  mapping(address => EnumerableSet.Bytes32Set) private activeSlots;
+  mapping(address => mapping(uint8 => EnumerableSet.Bytes32Set)) private activeSlots;
+  mapping(address => uint8) private activeSlotsIdx;
 
   constructor(
     IERC20 _token,
@@ -40,7 +41,25 @@ contract Marketplace is Collateral, Proofs {
   }
 
   function mySlots() public view returns (SlotId[] memory) {
-    return _toSlotIds(activeSlots[msg.sender].values());
+    EnumerableSet.Bytes32Set storage slotIds = _activeSlotsForHost(msg.sender);
+    return _toSlotIds(slotIds.values());
+  }
+
+  function _activeSlotsForHost(address host)
+    internal
+    view
+    returns (EnumerableSet.Bytes32Set storage)
+  {
+    mapping(uint8 => EnumerableSet.Bytes32Set) storage active = activeSlots[host];
+    uint8 idx = activeSlotsIdx[host];
+    return active[idx];
+  }
+
+  /// @notice Clears active slots for a host
+  /// @dev Because there are no efficient ways to clear an EnumerableSet, an index is updated that points to a new instance.
+  /// @param host address of the host for which to clear the active slots
+  function _clearActiveSlots(address host) internal {
+    activeSlotsIdx[host] = activeSlotsIdx[host] + 1;
   }
 
   function requestStorage(Request calldata request)
@@ -94,7 +113,7 @@ contract Marketplace is Collateral, Proofs {
     slot.requestId = requestId;
     RequestContext storage context = _context(requestId);
     context.slotsFilled += 1;
-    activeSlots[slot.host].add(SlotId.unwrap(slotId));
+    _activeSlotsForHost(slot.host).add(SlotId.unwrap(slotId));
     emit SlotFilled(requestId, slotIndex, slotId);
     if (context.slotsFilled == request.ask.slots) {
       context.state = RequestState.Started;
@@ -121,7 +140,8 @@ contract Marketplace is Collateral, Proofs {
 
     _unexpectProofs(_toProofId(slotId));
 
-    activeSlots[slot.host].remove(SlotId.unwrap(slotId));
+    address slotHost = slot.host;
+    _activeSlotsForHost(slotHost).remove(SlotId.unwrap(slotId));
     slot.host = address(0);
     slot.requestId = RequestId.wrap(0);
     context.slotsFilled -= 1;
@@ -137,9 +157,7 @@ contract Marketplace is Collateral, Proofs {
       _setProofEnd(_toEndId(requestId), block.timestamp - 1);
       context.endsAt = block.timestamp - 1;
       activeRequests[request.client].remove(RequestId.unwrap(requestId));
-      // NOTE: not removing any remaining active slots for the host as listing
-      // values of enumerable set could be too expensive (copies from storage to
-      // memory)
+      _clearActiveSlots(slotHost);
       emit RequestFailed(requestId);
 
       // TODO: burn all remaining slot collateral (note: slot collateral not
@@ -160,7 +178,7 @@ contract Marketplace is Collateral, Proofs {
     SlotId slotId = _toSlotId(requestId, slotIndex);
     Slot storage slot = _slot(slotId);
     require(!slot.hostPaid, "Already paid");
-    activeSlots[slot.host].remove(SlotId.unwrap(slotId));
+    _activeSlotsForHost(slot.host).remove(SlotId.unwrap(slotId));
     uint256 amount = pricePerSlot(requests[requestId]);
     funds.sent += amount;
     funds.balance -= amount;
diff --git a/test/Marketplace.test.js b/test/Marketplace.test.js
index d125b84..738eeb8 100644
--- a/test/Marketplace.test.js
+++ b/test/Marketplace.test.js
@@ -763,8 +763,7 @@ describe("Marketplace", function () {
     it("removes all slots from list when request fails", async function () {
       await waitUntilStarted(marketplace, request, proof)
       await waitUntilFailed(marketplace, request, slot)
-      let expectedLength = request.ask.slots - (request.ask.maxSlotLoss + 1)
-      expect((await marketplace.mySlots()).length).to.equal(expectedLength)
+      expect((await marketplace.mySlots())).to.deep.equal([])
     })
 
     it("removes slots from list when request finishes", async function () {