switch to a lib, and generalise the api

contracts/ActiveSlots.sol

@@ -1,28 +0,0 @@
-// SPDX-License-Identifier: MIT
-pragma solidity ^0.8.8;
-
-import "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-
-contract ActiveSlots {
-  mapping(bytes32 =>
-          mapping(address =>
-                  mapping(uint8 =>
-                          EnumerableSet.Bytes32Set))) private activeSlots;
-  mapping(bytes32 => uint8) private activeSlotsIdx;
-
-  function _activeSlotsForHost(address host, bytes32 requestId)
-    internal
-    view
-    returns (EnumerableSet.Bytes32Set storage)
-  {
-    uint8 id = activeSlotsIdx[requestId];
-    return activeSlots[requestId][host][id];
-  }
-
-  /// @notice Clears active slots for a request
-  /// @dev Because there are no efficient ways to clear an EnumerableSet, an index is updated that points to a new instance.
-  /// @param requestId request for which to clear the active slots
-  function _clearActiveSlots(bytes32 requestId) internal {
-    activeSlotsIdx[requestId]++;
-  }
-}

contracts/Marketplace.sol

@@ -6,10 +6,11 @@ import "@openzeppelin/contracts/utils/math/Math.sol";
 import "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
 import "./Collateral.sol";
 import "./Proofs.sol";
-import "./ActiveSlots.sol";
+import "./libs/SetMap.sol";
 
-contract Marketplace is Collateral, Proofs, ActiveSlots {
+contract Marketplace is Collateral, Proofs {
   using EnumerableSet for EnumerableSet.Bytes32Set;
+  using SetMap for SetMap.AddressSetMap;
 
   type RequestId is bytes32;
   type SlotId is bytes32;
@@ -20,6 +21,7 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
   mapping(RequestId => RequestContext) private requestContexts;
   mapping(SlotId => Slot) private slots;
   mapping(address => EnumerableSet.Bytes32Set) private activeRequests;
+  SetMap.AddressSetMap private activeSlots;
 
   constructor(
     IERC20 _token,
@@ -30,7 +32,6 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
   )
     Collateral(_token)
     Proofs(_proofPeriod, _proofTimeout, _proofDowntime)
-    ActiveSlots()
     marketplaceInvariant
   {
     collateral = _collateral;
@@ -45,10 +46,9 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
     view
     returns (SlotId[] memory)
   {
-    EnumerableSet.Bytes32Set
-      storage
-      slotIds = _activeSlotsForHost(msg.sender, RequestId.unwrap(requestId));
-    return _toSlotIds(slotIds.values());
+    bytes32[] memory slotIds = activeSlots.values(_toSetMapKey(requestId),
+                                                 msg.sender);
+    return _toSlotIds(slotIds);
   }
 
   function _equals(RequestId a, RequestId b) internal pure returns (bool) {
@@ -106,8 +106,9 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
     slot.requestId = requestId;
     RequestContext storage context = _context(requestId);
     context.slotsFilled += 1;
-    _activeSlotsForHost(slot.host, RequestId.unwrap(requestId))
-                        .add(SlotId.unwrap(slotId));
+    activeSlots.add(_toSetMapKey(requestId),
+                    slot.host,
+                    SlotId.unwrap(slotId));
     emit SlotFilled(requestId, slotIndex, slotId);
     if (context.slotsFilled == request.ask.slots) {
       context.state = RequestState.Started;
@@ -134,8 +135,9 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
 
     _unexpectProofs(_toProofId(slotId));
 
-    _activeSlotsForHost(slot.host, RequestId.unwrap(requestId))
-                        .remove(SlotId.unwrap(slotId));
+    activeSlots.remove(_toSetMapKey(requestId),
+                       slot.host,
+                       SlotId.unwrap(slotId));
     slot.host = address(0);
     slot.requestId = RequestId.wrap(0);
     context.slotsFilled -= 1;
@@ -151,7 +153,7 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
       _setProofEnd(_toEndId(requestId), block.timestamp - 1);
       context.endsAt = block.timestamp - 1;
       activeRequests[request.client].remove(RequestId.unwrap(requestId));
-      _clearActiveSlots(RequestId.unwrap(requestId));
+      activeSlots.clear(_toSetMapKey(requestId));
       emit RequestFailed(requestId);
 
       // TODO: burn all remaining slot collateral (note: slot collateral not
@@ -172,8 +174,9 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
     SlotId slotId = _toSlotId(requestId, slotIndex);
     Slot storage slot = _slot(slotId);
     require(!slot.hostPaid, "Already paid");
-    _activeSlotsForHost(slot.host, RequestId.unwrap(requestId))
-                        .remove(SlotId.unwrap(slotId));
+    activeSlots.remove(_toSetMapKey(requestId),
+                       slot.host,
+                       SlotId.unwrap(slotId));
     uint256 amount = pricePerSlot(requests[requestId]);
     funds.sent += amount;
     funds.balance -= amount;
@@ -403,6 +406,14 @@ contract Marketplace is Collateral, Proofs, ActiveSlots {
     return EndId.wrap(RequestId.unwrap(requestId));
   }
 
+  function _toSetMapKey(RequestId requestId)
+    internal
+    pure
+    returns (SetMap.Key)
+  {
+    return SetMap.Key.wrap(RequestId.unwrap(requestId));
+  }
+
   function _notEqual(RequestId a, uint256 b) internal pure returns (bool) {
     return RequestId.unwrap(a) != bytes32(b);
   }

contracts/libs/SetMap.sol

@@ -0,0 +1,98 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.8;
+
+import "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+
+library SetMap {
+  using EnumerableSet for EnumerableSet.Bytes32Set;
+
+  type Key is bytes32;
+
+  struct AddressSetMap {
+    mapping(Key =>
+          mapping(address =>
+                  mapping(uint8 =>
+                          EnumerableSet.Bytes32Set))) _values;
+    mapping(Key => uint8) _index;
+  }
+
+  /// @notice Returns the EnumerableSet.Bytes32 containing the values for a key
+  ///         and address in an AddressSetMap
+  /// @dev This is used internally to the library only. `.values()` should only
+  ///      be called on its return value in a view/pure function.
+  /// @param map AddressSetMap to list values
+  /// @param key key of the values to be listed
+  /// @param addr address of the values to be listed
+  /// @return bytes32[] array of bytes32 values
+  function _set(AddressSetMap storage map,
+                Key key,
+                address addr)
+    private
+    view
+    returns (EnumerableSet.Bytes32Set storage)
+  {
+    uint8 id = map._index[key];
+    return map._values[key][addr][id];
+  }
+
+  /// @notice Lists all values for a key and address in an AddressSetMap
+  /// @param map AddressSetMap to list values
+  /// @param key key of the values to be listed
+  /// @param addr address of the values to be listed
+  /// @return bytes32[] array of bytes32 values
+  function values(AddressSetMap storage map,
+                  Key key,
+                  address addr)
+    internal
+    view
+    returns (bytes32[] memory)
+  {
+    return _set(map, key, addr).values();
+  }
+
+  /// @notice Adds a single value to an AddressSetMap
+  /// @param map AddressSetMap to add the value to
+  /// @param key key of the value to be added
+  /// @param addr address of the value to be added
+  /// @param value the value to be added
+  /// @return true if the value was added to the set, that is if it was not
+  ///         already present.
+  function add(AddressSetMap storage map,
+               Key key,
+               address addr,
+               bytes32 value)
+    internal
+    returns (bool)
+  {
+    return _set(map, key, addr).add(value);
+  }
+
+  /// @notice Removes a single value from an AddressSetMap
+  /// @param map AddressSetMap to remove the value from
+  /// @param key key of the value to be removed
+  /// @param addr address of the value to be removed
+  /// @param value the value to be removed
+  /// @return true if the value was removed from the set, that is if it was
+  ///         present.
+  function remove(AddressSetMap storage map,
+                  Key key,
+                  address addr,
+                  bytes32 value)
+    internal
+    returns (bool)
+  {
+    return _set(map, key, addr).remove(value);
+  }
+
+  /// @notice Clears values for a key.
+  /// @dev Does not clear the addresses for the key, simply updates an index
+  ///      such that the next time values for that key and address are
+  ///      retrieved, it will return an empty array.
+  /// @param map AddressSetMap for which to clear values
+  /// @param key key for which to clear values
+  function clear(AddressSetMap storage map, Key key)
+    internal
+  {
+    map._index[key]++;
+  }
+}