Merge branch 'main' into annotate-rewards-tests

Author: pdyraga

contracts/SortitionTree.sol

@@ -19,7 +19,18 @@ contract SortitionTree {
   // level6 256k
   // level7 2M
   uint256 internal root;
+
+  // A 2-index mapping from layer => (index (0-index) => branch). For example,
+  // to access the 6th branch in the 2nd layer (right below the root node; the
+  // first branch layer), call branches[2][5]. Mappings are used in place of
+  // arrays for efficiency. The root is the first layer, the branches occupy
+  // layers 2 through 7, and layer 8 is for the leaves. Following this
+  // convention, the first index in `branches` is `2`, and the last index is
+  // `7`.
   mapping(uint256 => mapping(uint256 => uint256)) internal branches;
+
+  // A 0-index mapping from index => leaf, acting as an array. For example, to
+  // access the 42nd leaf, call leaves[41].
   mapping(uint256 => uint256) internal leaves;
 
   // the flagged (see setFlag() and unsetFlag() in Position.sol) positions
@@ -28,6 +39,7 @@ contract SortitionTree {
 
   // the leaf after the rightmost occupied leaf of each stack
   uint256 internal rightmostLeaf;
+
   // the empty leaves in each stack
   // between 0 and the rightmost occupied leaf
   uint256[] internal emptyLeaves;
@@ -36,6 +48,7 @@ contract SortitionTree {
   // which is allocated when they first join the pool
   // and remains unchanged even if they leave and rejoin the pool.
   mapping(address => uint32) internal operatorID;
+
   // The idAddress array records the address corresponding to each ID number.
   // The ID number 0 is initialized with a zero address and is not used.
   address[] internal idAddress;
@@ -46,22 +59,28 @@ contract SortitionTree {
     idAddress.push();
   }
 
-  // Return the ID number of the given operator address.
-  // An ID number of 0 means the operator has not been allocated an ID number yet.
+  /// @notice Return the ID number of the given operator address. An ID number
+  /// of 0 means the operator has not been allocated an ID number yet.
+  /// @param operator Address of the operator.
+  /// @return the ID number of the given operator address
   function getOperatorID(address operator) public view returns (uint32) {
     return operatorID[operator];
   }
 
-  // Get the operator address corresponding to the given ID number.
-  // An empty address means the ID number has not been allocated yet.
+  /// @notice Get the operator address corresponding to the given ID number. A
+  /// zero address means the ID number has not been allocated yet.
+  /// @param id ID of the operator
+  /// @return the address of the operator
   function getIDOperator(uint32 id) public view returns (address) {
     return idAddress.length > id ? idAddress[id] : address(0);
   }
 
-  // Gets the operator addresses corresponding to the given ID numbers.
-  // An empty address means the ID number has not been allocated yet.
-  // This function works just like getIDOperator except that it allows to fetch
-  // operator addresses for multiple IDs in one call.
+  /// @notice Gets the operator addresses corresponding to the given ID
+  /// numbers. A zero address means the ID number has not been allocated yet.
+  /// This function works just like getIDOperator except that it allows to fetch
+  /// operator addresses for multiple IDs in one call.
+  /// @param ids the array of the operator ids
+  /// @return an array of the associated operator addresses
   function getIDOperators(uint32[] calldata ids)
     public
     view
@@ -77,12 +96,15 @@ contract SortitionTree {
     return operators;
   }
 
-  // checks if operator is already registered in the pool
+  /// @notice Checks if operator is already registered in the pool.
+  /// @param operator the address of the operator
+  /// @return whether or not the operator is already registered in the pool
   function isOperatorRegistered(address operator) public view returns (bool) {
     return getFlaggedLeafPosition(operator) != 0;
   }
 
-  // Sum the number of operators in each trunk
+  /// @notice Sum the number of operators in each trunk.
+  /// @return the number of operators in the pool
   function operatorsInPool() public view returns (uint256) {
     // Get the number of leaves that might be occupied;
     // if `rightmostLeaf` equals `firstLeaf()` the tree must be empty,
@@ -96,12 +118,16 @@ contract SortitionTree {
     return (nPossiblyUsedLeaves - nEmptyLeaves);
   }
 
+  /// @notice Convenience method to return the total weight of the pool
+  /// @return the total weight of the pool
   function totalWeight() public view returns (uint256) {
     return root.sumWeight();
   }
 
-  // Give the operator a new ID number
-  // Does not check if the operator already has an ID number
+  /// @notice Give the operator a new ID number.
+  /// Does not check if the operator already has an ID number.
+  /// @param operator the address of the operator
+  /// @return a new ID for that operator
   function allocateOperatorID(address operator) internal returns (uint256) {
     uint256 id = idAddress.length;
 
@@ -112,36 +138,51 @@ contract SortitionTree {
     return id;
   }
 
+  /// @notice Inserts an operator into the sortition pool
+  /// @param operator the address of an operator to insert
+  /// @param weight how much weight that operator has in the pool
   function _insertOperator(address operator, uint256 weight) internal {
     require(
       !isOperatorRegistered(operator),
       "Operator is already registered in the pool"
     );
 
+    // Fetch the operator's ID, and if they don't have one, allocate them one.
     uint256 id = getOperatorID(operator);
     if (id == 0) {
       id = allocateOperatorID(operator);
     }
 
+    // Determine which leaf to insert them into
     uint256 position = getEmptyLeafPosition();
     // Record the block the operator was inserted in
     uint256 theLeaf = Leaf.make(operator, block.number, id);
 
+    // Update the leaf, and propagate the weight changes all the way up to the
+    // root.
     root = setLeaf(position, theLeaf, weight, root);
 
     // Without position flags,
     // the position 0x000000 would be treated as empty
     flaggedLeafPosition[operator] = position.setFlag();
   }
 
+  /// @notice Remove an operator (and their weight) from the pool.
+  /// @param operator the address of the operator to remove
   function _removeOperator(address operator) internal {
     uint256 flaggedPosition = getFlaggedLeafPosition(operator);
     require(flaggedPosition != 0, "Operator is not registered in the pool");
     uint256 unflaggedPosition = flaggedPosition.unsetFlag();
+
+    // Update the leaf, and propagate the weight changes all the way up to the
+    // root.
     root = removeLeaf(unflaggedPosition, root);
     removeLeafPositionRecord(operator);
   }
 
+  /// @notice Update an operator's weight in the pool.
+  /// @param operator the address of the operator to update
+  /// @param weight the new weight
   function updateOperator(address operator, uint256 weight) internal {
     require(
       isOperatorRegistered(operator),
@@ -153,19 +194,28 @@ contract SortitionTree {
     root = updateLeaf(unflaggedPosition, weight, root);
   }
 
+  /// @notice Helper method to remove a leaf position record for an operator.
+  /// @param operator the address of the operator to remove the record for
   function removeLeafPositionRecord(address operator) internal {
     flaggedLeafPosition[operator] = 0;
   }
 
+  /// @notice Removes the data and weight from a particular leaf.
+  /// @param position the leaf index to remove
+  /// @param _root the root node containing the leaf
+  /// @return the updated root node
   function removeLeaf(uint256 position, uint256 _root)
     internal
     returns (uint256)
   {
     uint256 rightmostSubOne = rightmostLeaf - 1;
     bool isRightmost = position == rightmostSubOne;
 
+    // Clears out the data in the leaf node, and then propagates the weight
+    // changes all the way up to the root.
     uint256 newRoot = setLeaf(position, 0, 0, _root);
 
+    // Infer if need to fall back on emptyLeaves yet
     if (isRightmost) {
       rightmostLeaf = rightmostSubOne;
     } else {
@@ -174,6 +224,11 @@ contract SortitionTree {
     return newRoot;
   }
 
+  /// @notice Updates the tree to give a particular leaf a new weight.
+  /// @param position the index of the leaf to update
+  /// @param weight the new weight
+  /// @param _root the root node containing the leaf
+  /// @return the updated root node
   function updateLeaf(
     uint256 position,
     uint256 weight,
@@ -186,6 +241,13 @@ contract SortitionTree {
     }
   }
 
+  /// @notice Places a leaf into a particular position, with a given weight and
+  /// propagates that change.
+  /// @param position the index to place the leaf in
+  /// @param theLeaf the new leaf to place in the position
+  /// @param leafWeight the weight of the leaf
+  /// @param _root the root containing the new leaf
+  /// @return the updated root node
   function setLeaf(
     uint256 position,
     uint256 theLeaf,
@@ -198,6 +260,12 @@ contract SortitionTree {
     return (updateTree(position, leafWeight, _root));
   }
 
+  /// @notice Propagates a weight change at a position through the tree,
+  /// eventually returning the updated root.
+  /// @param position the index of leaf to update
+  /// @param weight the new weight of the leaf
+  /// @param _root the root node containing the leaf
+  /// @return the updated root node
   function updateTree(
     uint256 position,
     uint256 weight,
@@ -224,6 +292,10 @@ contract SortitionTree {
     return _root.setSlot(childSlot, nodeWeight);
   }
 
+  /// @notice Retrieves the next available empty leaf position. Tries to fill
+  /// left to right first, ignoring leaf removals, and then fills
+  /// most-recent-removals first.
+  /// @return the position of the empty leaf
   function getEmptyLeafPosition() internal returns (uint256) {
     uint256 rLeaf = rightmostLeaf;
     bool spaceOnRight = (rLeaf + 1) < Constants.POOL_CAPACITY;
@@ -239,6 +311,9 @@ contract SortitionTree {
     }
   }
 
+  /// @notice Gets the flagged leaf position for an operator.
+  /// @param operator the address of the operator
+  /// @return the leaf position of that operator
   function getFlaggedLeafPosition(address operator)
     internal
     view
@@ -247,13 +322,24 @@ contract SortitionTree {
     return flaggedLeafPosition[operator];
   }
 
+  /// @notice Gets the weight of a leaf at a particular position.
+  /// @param position the index of the leaf
+  /// @return the weight of the leaf at that position
   function getLeafWeight(uint256 position) internal view returns (uint256) {
     uint256 slot = position.slot();
     uint256 parent = position.parent();
+
+    // A leaf's weight information is stored a 32-bit slot in the branch layer
+    // directly above the leaf layer. To access it, we calculate that slot and
+    // parent position, and always know the hard-coded layer index.
     uint256 node = branches[Constants.LEVELS][parent];
     return node.getSlot(slot);
   }
 
+  /// @notice Picks a leaf given a random index.
+  /// @param index a number in `[0, _root.totalWeight())` used to decide
+  /// between leaves
+  /// @param _root the root of the tree
   function pickWeightedLeaf(uint256 index, uint256 _root)
     internal
     view

contracts/test/LeafStub.sol

@@ -0,0 +1,25 @@
+pragma solidity 0.8.9;
+
+import "../Leaf.sol";
+
+contract LeafStub {
+  function make(
+    address operator,
+    uint256 creationBlock,
+    uint256 id
+  ) public pure returns (uint256) {
+    return Leaf.make(operator, creationBlock, id);
+  }
+
+  function operator(uint256 leaf) public pure returns (address) {
+    return Leaf.operator(leaf);
+  }
+
+  function creationBlock(uint256 leaf) public pure returns (uint256) {
+    return Leaf.creationBlock(leaf);
+  }
+
+  function id(uint256 leaf) public pure returns (uint32) {
+    return Leaf.id(leaf);
+  }
+}