Merge branch 'main' into implementation-details-readme

Author: pdyraga

contracts/Branch.sol

@@ -1,26 +1,11 @@
 pragma solidity 0.8.9;
 
+import "./Constants.sol";
+
 /// @notice The implicit 8-ary trees of the sortition pool
 /// rely on packing 8 "slots" of 32-bit values into each uint256.
 /// The Branch library permits efficient calculations on these slots.
 library Branch {
-  ////////////////////////////////////////////////////////////////////////////
-  // Parameters for configuration
-
-  // How many bits a position uses per level of the tree;
-  // each branch of the tree contains 2**SLOT_BITS slots.
-  uint256 private constant SLOT_BITS = 3;
-  ////////////////////////////////////////////////////////////////////////////
-
-  ////////////////////////////////////////////////////////////////////////////
-  // Derived constants, do not touch
-  uint256 private constant SLOT_COUNT = 2**SLOT_BITS;
-  uint256 private constant SLOT_WIDTH = 256 / SLOT_COUNT;
-  uint256 private constant LAST_SLOT = SLOT_COUNT - 1;
-  uint256 private constant SLOT_MAX = (2**SLOT_WIDTH) - 1;
-
-  ////////////////////////////////////////////////////////////////////////////
-
   /// @notice Calculate the right shift required
   /// to make the 32 least significant bits of an uint256
   /// be the bits of the `position`th slot
@@ -31,7 +16,7 @@ library Branch {
   /// I wish solidity had macros, even C macros.
   function slotShift(uint256 position) internal pure returns (uint256) {
     unchecked {
-      return position * SLOT_WIDTH;
+      return position * Constants.SLOT_WIDTH;
     }
   }
 
@@ -43,12 +28,12 @@ library Branch {
     returns (uint256)
   {
     unchecked {
-      uint256 shiftBits = position * SLOT_WIDTH;
+      uint256 shiftBits = position * Constants.SLOT_WIDTH;
       // Doing a bitwise AND with `SLOT_MAX`
       // clears all but the 32 least significant bits.
       // Because of the right shift by `slotShift(position)` bits,
       // those 32 bits contain the 32 bits in the `position`th slot of `node`.
-      return (node >> shiftBits) & SLOT_MAX;
+      return (node >> shiftBits) & Constants.SLOT_MAX;
     }
   }
 
@@ -59,7 +44,7 @@ library Branch {
     returns (uint256)
   {
     unchecked {
-      uint256 shiftBits = position * SLOT_WIDTH;
+      uint256 shiftBits = position * Constants.SLOT_WIDTH;
       // Shifting `SLOT_MAX` left by `slotShift(position)` bits
       // gives us a number where all bits of the `position`th slot are set,
       // and all other bits are unset.
@@ -71,7 +56,7 @@ library Branch {
       // Bitwise ANDing the original `node` with this number
       // sets the bits of `position`th slot to zero,
       // leaving all other bits unchanged.
-      return node & ~(SLOT_MAX << shiftBits);
+      return node & ~(Constants.SLOT_MAX << shiftBits);
     }
   }
 
@@ -86,17 +71,17 @@ library Branch {
     uint256 weight
   ) internal pure returns (uint256) {
     unchecked {
-      uint256 shiftBits = position * SLOT_WIDTH;
+      uint256 shiftBits = position * Constants.SLOT_WIDTH;
       // Clear the `position`th slot like in `clearSlot()`.
-      uint256 clearedNode = node & ~(SLOT_MAX << shiftBits);
+      uint256 clearedNode = node & ~(Constants.SLOT_MAX << shiftBits);
       // Bitwise AND `weight` with `SLOT_MAX`
       // to clear all but the 32 least significant bits.
       //
       // Shift this left by `slotShift(position)` bits
       // to obtain a uint256 with all bits unset
       // except in the `position`th slot
       // which contains the 32-bit value of `weight`.
-      uint256 shiftedWeight = (weight & SLOT_MAX) << shiftBits;
+      uint256 shiftedWeight = (weight & Constants.SLOT_MAX) << shiftBits;
       // When we bitwise OR these together,
       // all other slots except the `position`th one come from the left argument,
       // and the `position`th gets filled with `weight` from the right argument.
@@ -107,14 +92,14 @@ library Branch {
   /// @notice Calculate the summed weight of all slots in the `node`.
   function sumWeight(uint256 node) internal pure returns (uint256 sum) {
     unchecked {
-      sum = node & SLOT_MAX;
+      sum = node & Constants.SLOT_MAX;
       // Iterate through each slot
       // by shifting `node` right in increments of 32 bits,
       // and adding the 32 least significant bits to the `sum`.
-      uint256 newNode = node >> SLOT_WIDTH;
+      uint256 newNode = node >> Constants.SLOT_WIDTH;
       while (newNode > 0) {
-        sum += (newNode & SLOT_MAX);
-        newNode = newNode >> SLOT_WIDTH;
+        sum += (newNode & Constants.SLOT_MAX);
+        newNode = newNode >> Constants.SLOT_WIDTH;
       }
       return sum;
     }
@@ -143,12 +128,12 @@ library Branch {
     unchecked {
       newIndex = index;
       uint256 newNode = node;
-      uint256 currentSlotWeight = newNode & SLOT_MAX;
+      uint256 currentSlotWeight = newNode & Constants.SLOT_MAX;
       while (newIndex >= currentSlotWeight) {
         newIndex -= currentSlotWeight;
         slot++;
-        newNode = newNode >> SLOT_WIDTH;
-        currentSlotWeight = newNode & SLOT_MAX;
+        newNode = newNode >> Constants.SLOT_WIDTH;
+        currentSlotWeight = newNode & Constants.SLOT_MAX;
       }
       return (slot, newIndex);
     }

contracts/Constants.sol

@@ -0,0 +1,32 @@
+pragma solidity 0.8.9;
+
+library Constants {
+  ////////////////////////////////////////////////////////////////////////////
+  // Parameters for configuration
+
+  // How many bits a position uses per level of the tree;
+  // each branch of the tree contains 2**SLOT_BITS slots.
+  uint256 constant SLOT_BITS = 3;
+  uint256 constant LEVELS = 7;
+  ////////////////////////////////////////////////////////////////////////////
+
+  ////////////////////////////////////////////////////////////////////////////
+  // Derived constants, do not touch
+  uint256 constant SLOT_COUNT = 2**SLOT_BITS;
+  uint256 constant SLOT_WIDTH = 256 / SLOT_COUNT;
+  uint256 constant LAST_SLOT = SLOT_COUNT - 1;
+  uint256 constant SLOT_MAX = (2**SLOT_WIDTH) - 1;
+  uint256 constant POOL_CAPACITY = SLOT_COUNT**LEVELS;
+
+  uint256 constant ID_WIDTH = SLOT_WIDTH;
+  uint256 constant ID_MAX = SLOT_MAX;
+
+  uint256 constant BLOCKHEIGHT_WIDTH = 96 - ID_WIDTH;
+  uint256 constant BLOCKHEIGHT_MAX = (2**BLOCKHEIGHT_WIDTH) - 1;
+
+  uint256 constant SLOT_POINTER_MAX = (2**SLOT_BITS) - 1;
+  uint256 constant LEAF_FLAG = 1 << 255;
+
+  uint256 constant WEIGHT_WIDTH = 256 / SLOT_COUNT;
+  ////////////////////////////////////////////////////////////////////////////
+}

contracts/Leaf.sol

@@ -1,28 +1,8 @@
 pragma solidity 0.8.9;
 
-library Leaf {
-  ////////////////////////////////////////////////////////////////////////////
-  // Parameters for configuration
-
-  // How many bits a position uses per level of the tree;
-  // each branch of the tree contains 2**SLOT_BITS slots.
-  uint256 private constant SLOT_BITS = 3;
-  ////////////////////////////////////////////////////////////////////////////
-
-  ////////////////////////////////////////////////////////////////////////////
-  // Derived constants, do not touch
-  uint256 private constant SLOT_COUNT = 2**SLOT_BITS;
-  uint256 private constant SLOT_WIDTH = 256 / SLOT_COUNT;
-  uint256 private constant SLOT_MAX = (2**SLOT_WIDTH) - 1;
-
-  uint256 private constant ID_WIDTH = SLOT_WIDTH;
-  uint256 private constant ID_MAX = SLOT_MAX;
-
-  uint256 private constant BLOCKHEIGHT_WIDTH = 96 - ID_WIDTH;
-  uint256 private constant BLOCKHEIGHT_MAX = (2**BLOCKHEIGHT_WIDTH) - 1;
-
-  ////////////////////////////////////////////////////////////////////////////
+import "./Constants.sol";
 
+library Leaf {
   function make(
     address _operator,
     uint256 _creationBlock,
@@ -35,10 +15,11 @@ library Leaf {
     uint256 op = uint256(bytes32(bytes20(_operator)));
     // Bitwise AND the id to erase
     // all but the 32 least significant bits
-    uint256 uid = _id & ID_MAX;
+    uint256 uid = _id & Constants.ID_MAX;
     // Erase all but the 64 least significant bits,
     // then shift left by 32 bits to make room for the id
-    uint256 cb = (_creationBlock & BLOCKHEIGHT_MAX) << ID_WIDTH;
+    uint256 cb = (_creationBlock & Constants.BLOCKHEIGHT_MAX) <<
+      Constants.ID_WIDTH;
     // Bitwise OR them all together to get
     // [address operator || uint64 creationBlock || uint32 id]
     return (op | cb | uid);
@@ -52,12 +33,12 @@ library Leaf {
 
   /// @notice Return the block number the leaf was created in.
   function creationBlock(uint256 leaf) internal pure returns (uint256) {
-    return ((leaf >> ID_WIDTH) & BLOCKHEIGHT_MAX);
+    return ((leaf >> Constants.ID_WIDTH) & Constants.BLOCKHEIGHT_MAX);
   }
 
   function id(uint256 leaf) internal pure returns (uint32) {
     // Id is stored in the 32 least significant bits.
     // Bitwise AND ensures that we only get the contents of those bits.
-    return uint32(leaf & ID_MAX);
+    return uint32(leaf & Constants.ID_MAX);
   }
 }

contracts/Position.sol

@@ -1,43 +1,30 @@
 pragma solidity 0.8.9;
 
-library Position {
-  ////////////////////////////////////////////////////////////////////////////
-  // Parameters for configuration
-
-  // How many bits a position uses per level of the tree;
-  // each branch of the tree contains 2**SLOT_BITS slots.
-  uint256 private constant SLOT_BITS = 3;
-  ////////////////////////////////////////////////////////////////////////////
-
-  ////////////////////////////////////////////////////////////////////////////
-  // Derived constants, do not touch
-  uint256 private constant SLOT_POINTER_MAX = (2**SLOT_BITS) - 1;
-  uint256 private constant LEAF_FLAG = 1 << 255;
-
-  ////////////////////////////////////////////////////////////////////////////
+import "./Constants.sol";
 
+library Position {
   // Return the last 3 bits of a position number,
   // corresponding to its slot in its parent
   function slot(uint256 a) internal pure returns (uint256) {
-    return a & SLOT_POINTER_MAX;
+    return a & Constants.SLOT_POINTER_MAX;
   }
 
   // Return the parent of a position number
   function parent(uint256 a) internal pure returns (uint256) {
-    return a >> SLOT_BITS;
+    return a >> Constants.SLOT_BITS;
   }
 
   // Return the location of the child of a at the given slot
   function child(uint256 a, uint256 s) internal pure returns (uint256) {
-    return (a << SLOT_BITS) | (s & SLOT_POINTER_MAX); // slot(s)
+    return (a << Constants.SLOT_BITS) | (s & Constants.SLOT_POINTER_MAX); // slot(s)
   }
 
   // Return the uint p as a flagged position uint:
   // the least significant 21 bits contain the position
   // and the 22nd bit is set as a flag
   // to distinguish the position 0x000000 from an empty field.
   function setFlag(uint256 p) internal pure returns (uint256) {
-    return p | LEAF_FLAG;
+    return p | Constants.LEAF_FLAG;
   }
 
   // Turn a flagged position into an unflagged position
@@ -47,6 +34,6 @@ library Position {
   // as all position-manipulating code should ignore non-position bits anyway
   // but it's cheap to call so might as well do it.
   function unsetFlag(uint256 p) internal pure returns (uint256) {
-    return p & (~LEAF_FLAG);
+    return p & (~Constants.LEAF_FLAG);
   }
 }

contracts/RNG.sol

@@ -1,23 +1,9 @@
 pragma solidity 0.8.9;
 
 import "./Leaf.sol";
+import "./Constants.sol";
 
 library RNG {
-  ////////////////////////////////////////////////////////////////////////////
-  // Parameters for configuration
-
-  // How many bits a position uses per level of the tree;
-  // each branch of the tree contains 2**SLOT_BITS slots.
-  uint256 private constant SLOT_BITS = 3;
-  ////////////////////////////////////////////////////////////////////////////
-
-  ////////////////////////////////////////////////////////////////////////////
-  // Derived constants, do not touch
-  uint256 private constant SLOT_COUNT = 2**SLOT_BITS;
-  uint256 private constant WEIGHT_WIDTH = 256 / SLOT_COUNT;
-
-  ////////////////////////////////////////////////////////////////////////////
-
   /// @notice Get an index in the range `[0 .. range-1]`
   /// and the new state of the RNG,
   /// using the provided `state` of the RNG.
@@ -76,7 +62,7 @@ library RNG {
         return 0;
       }
 
-      uint256 bits = WEIGHT_WIDTH - 1;
+      uint256 bits = Constants.WEIGHT_WIDTH - 1;
 
       // Left shift by `bits`,
       // so we have a 1 in the (bits + 1)th least significant bit

contracts/Rewards.sol

@@ -34,6 +34,8 @@ contract Rewards {
     // The amount of rewards collected by the operator after the latest update.
     // The amount the operator could withdraw may equal `available`
     // or it may be greater, if more rewards have been paid in since then.
+    // To evaulate the most recent amount including rewards potentially paid
+    // since the last update, use `availableRewards` function.
     uint96 available;
     // If nonzero, the operator is ineligible for rewards
     // and may only re-enable rewards after the specified timestamp.
@@ -196,4 +198,21 @@ contract Rewards {
     o.ineligibleUntil = 0;
     operatorRewards[operator] = o;
   }
+
+  /// @notice Returns the amount of rewards currently available for withdrawal
+  ///         for the given operator.
+  function availableRewards(uint32 operator) internal view returns (uint96) {
+    uint96 acc = globalRewardAccumulator;
+    OperatorRewards memory o = operatorRewards[operator];
+    if (o.ineligibleUntil == 0) {
+      // If operator is not ineligible, calculate newly accrued rewards and add
+      // them to the available ones, calculated during the last update.
+      uint96 accruedRewards = (acc - o.accumulated) * uint96(o.weight);
+      return o.available + accruedRewards;
+    } else {
+      // If ineligible, return only the rewards calculated during the last
+      // update.
+      return o.available;
+    }
+  }
 }

contracts/SortitionPool.sol

@@ -56,16 +56,27 @@ contract SortitionPool is SortitionTree, Rewards, Ownable, IReceiveApproval {
     Rewards.addRewards(uint96(amount), uint32(root.sumWeight()));
   }
 
+  /// @notice Withdraws all available rewards for the given operator to the
+  ///         given beneficiary.
+  /// @dev Can be called only be the owner. Does not validate if the provided
+  ///      beneficiary is associated with the provided operator - this needs to
+  ///      be done by the owner calling this function.
+  /// @return The amount of rewards withdrawn in this call.
   function withdrawRewards(address operator, address beneficiary)
     public
     onlyOwner
+    returns (uint96)
   {
     uint32 id = getOperatorID(operator);
     Rewards.updateOperatorRewards(id, uint32(getPoolWeight(operator)));
     uint96 earned = Rewards.withdrawOperatorRewards(id);
     rewardToken.transfer(beneficiary, uint256(earned));
+    return earned;
   }
 
+  /// @notice Withdraws rewards not allocated to operators marked as ineligible
+  ///         to the given recipient address.
+  /// @dev Can be called only by the owner.
   function withdrawIneligible(address recipient) public onlyOwner {
     uint96 earned = Rewards.withdrawIneligibleRewards();
     rewardToken.transfer(recipient, uint256(earned));
@@ -140,6 +151,12 @@ contract SortitionPool is SortitionTree, Rewards, Ownable, IReceiveApproval {
     emit RewardEligibilityRestored(operator, id);
   }
 
+  /// @notice Returns the amount of rewards withdrawable for the given operator.
+  function getAvailableRewards(address operator) public view returns (uint96) {
+    uint32 id = getOperatorID(operator);
+    return availableRewards(id);
+  }
+
   /// @notice Return whether the operator is present in the pool.
   function isOperatorInPool(address operator) public view returns (bool) {
     return getFlaggedLeafPosition(operator) != 0;