Comment the sortition tree

Author: Beau Shinkle

contracts/SortitionTree.sol

@@ -35,7 +35,14 @@ contract SortitionTree {
   // level6 256k
   // level7 2M
   uint256 internal root;
+
+  // A mapping from layer => (index => branch). For example, to access the 6th
+  // (0-index) branch in the 2nd layer (right below the root node), call
+  // branches[2][6]. Mappings are used in place of arrays for efficiency.
   mapping(uint256 => mapping(uint256 => uint256)) internal branches;
+
+  // A mapping from index => leaf, acting as an array. For example, to access
+  // the 42nd (0-index) leaf, call leaves[42].
   mapping(uint256 => uint256) internal leaves;
 
   // the flagged (see setFlag() and unsetFlag() in Position.sol) positions
@@ -44,6 +51,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;
@@ -52,6 +60,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;
@@ -62,22 +71,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. An
+  /// empty 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. 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.
+  /// @param ids the array of the operator ids
+  /// @return an array of the associated operator addresses
   function getIDOperators(uint32[] calldata ids)
     public
     view
@@ -93,12 +108,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,
@@ -112,12 +130,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;
 
@@ -128,36 +150,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),
@@ -169,19 +206,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 {
@@ -190,6 +236,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,
@@ -202,6 +253,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,
@@ -214,6 +272,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,
@@ -240,6 +304,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) < POOL_CAPACITY;
@@ -255,6 +323,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
@@ -263,13 +334,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[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