refactor: extracted the sortition tree to a library

Author: jaybuidl

contracts/src/arbitration/KlerosCoreBase.sol

@@ -223,7 +223,7 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         // FORKING_COURT
         // TODO: Fill the properties for the Forking court, emit CourtCreated.
         courts.push();
-        sortitionModule.createTree(bytes32(uint256(FORKING_COURT)), _sortitionExtraData);
+        sortitionModule.createTree(FORKING_COURT, _sortitionExtraData);
 
         // GENERAL_COURT
         Court storage court = courts.push();
@@ -236,7 +236,7 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         court.jurorsForCourtJump = _courtParameters[3];
         court.timesPerPeriod = _timesPerPeriod;
 
-        sortitionModule.createTree(bytes32(uint256(GENERAL_COURT)), _sortitionExtraData);
+        sortitionModule.createTree(GENERAL_COURT, _sortitionExtraData);
 
         uint256[] memory supportedDisputeKits = new uint256[](1);
         supportedDisputeKits[0] = DISPUTE_KIT_CLASSIC;
@@ -343,7 +343,7 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         if (_supportedDisputeKits.length == 0) revert UnsupportedDisputeKit();
         if (_parent == FORKING_COURT) revert InvalidForkingCourtAsParent();
 
-        uint256 courtID = courts.length;
+        uint96 courtID = uint96(courts.length);
         Court storage court = courts.push();
 
         for (uint256 i = 0; i < _supportedDisputeKits.length; i++) {
@@ -364,7 +364,7 @@ abstract contract KlerosCoreBase is IArbitratorV2, Initializable, UUPSProxiable
         court.jurorsForCourtJump = _jurorsForCourtJump;
         court.timesPerPeriod = _timesPerPeriod;
 
-        sortitionModule.createTree(bytes32(courtID), _sortitionExtraData);
+        sortitionModule.createTree(courtID, _sortitionExtraData);
 
         // Update the parent.
         courts[_parent].children.push(courtID);

contracts/src/arbitration/SortitionModuleBase.sol

@@ -7,25 +7,20 @@ import {ISortitionModule} from "./interfaces/ISortitionModule.sol";
 import {IDisputeKit} from "./interfaces/IDisputeKit.sol";
 import {Initializable} from "../proxy/Initializable.sol";
 import {UUPSProxiable} from "../proxy/UUPSProxiable.sol";
+import {SortitionTrees, TreeKey, CourtID} from "../libraries/SortitionTrees.sol";
 import {IRNG} from "../rng/IRNG.sol";
 import "../libraries/Constants.sol";
 
 /// @title SortitionModuleBase
 /// @dev A factory of trees that keeps track of staked values for sortition.
 abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSProxiable {
+    using SortitionTrees for SortitionTrees.Tree;
+    using SortitionTrees for mapping(TreeKey key => SortitionTrees.Tree);
+
     // ************************************* //
     // *         Enums / Structs           * //
     // ************************************* //
 
-    struct SortitionSumTree {
-        uint256 K; // The maximum number of children per node.
-        uint256[] stack; // We use this to keep track of vacant positions in the tree after removing a leaf. This is for keeping the tree as balanced as possible without spending gas on moving nodes around.
-        uint256[] nodes; // The tree nodes.
-        // Two-way mapping of IDs to node indexes. Note that node index 0 is reserved for the root node, and means the ID does not have a node.
-        mapping(bytes32 stakePathID => uint256 nodeIndex) IDsToNodeIndexes;
-        mapping(uint256 nodeIndex => bytes32 stakePathID) nodeIndexesToIDs;
-    }
-
     struct DelayedStake {
         address account; // The address of the juror.
         uint96 courtID; // The ID of the court.
@@ -56,7 +51,7 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     uint256 public rngLookahead; // DEPRECATED: to be removed in the next redeploy
     uint256 public delayedStakeWriteIndex; // The index of the last `delayedStake` item that was written to the array. 0 index is skipped.
     uint256 public delayedStakeReadIndex; // The index of the next `delayedStake` item that should be processed. Starts at 1 because 0 index is skipped.
-    mapping(bytes32 treeHash => SortitionSumTree) sortitionSumTrees; // The mapping trees by keys.
+    mapping(TreeKey key => SortitionTrees.Tree) sortitionSumTrees; // The mapping of sortition trees by keys.
     mapping(address account => Juror) public jurors; // The jurors.
     mapping(uint256 => DelayedStake) public delayedStakes; // Stores the stakes that were changed during Drawing phase, to update them when the phase is switched to Staking.
     mapping(address jurorAccount => mapping(uint96 courtId => uint256)) public latestDelayedStakeIndex; // DEPRECATED. Maps the juror to its latest delayed stake. If there is already a delayed stake for this juror then it'll be replaced. latestDelayedStakeIndex[juror][courtID].
@@ -185,15 +180,12 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     }
 
     /// @dev Create a sortition sum tree at the specified key.
-    /// @param _key The key of the new tree.
+    /// @param _courtID The ID of the court.
     /// @param _extraData Extra data that contains the number of children each node in the tree should have.
-    function createTree(bytes32 _key, bytes memory _extraData) external override onlyByCore {
-        SortitionSumTree storage tree = sortitionSumTrees[_key];
+    function createTree(uint96 _courtID, bytes memory _extraData) external override onlyByCore {
+        TreeKey key = CourtID.wrap(_courtID).toTreeKey();
         uint256 K = _extraDataToTreeK(_extraData);
-        if (tree.K != 0) revert TreeAlreadyExists();
-        if (K <= 1) revert KMustBeGreaterThanOne();
-        tree.K = K;
-        tree.nodes.push(0);
+        sortitionSumTrees.createTree(key, K);
     }
 
     /// @dev Executes the next delayed stakes.
@@ -398,12 +390,13 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         }
 
         // Update the sortition sum tree.
-        bytes32 stakePathID = _accountAndCourtIDToStakePathID(_account, _courtID);
+        bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID);
         bool finished = false;
         uint96 currentCourtID = _courtID;
         while (!finished) {
             // Tokens are also implicitly staked in parent courts through sortition module to increase the chance of being drawn.
-            _set(bytes32(uint256(currentCourtID)), _newStake, stakePathID);
+            TreeKey key = CourtID.wrap(currentCourtID).toTreeKey();
+            sortitionSumTrees[key].set(_newStake, stakePathID);
             if (currentCourtID == GENERAL_COURT) {
                 finished = true;
             } else {
@@ -477,71 +470,32 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
 
     /// @dev Draw an ID from a tree using a number.
     /// Note that this function reverts if the sum of all values in the tree is 0.
-    /// @param _key The key of the tree.
+    /// @param _courtID The ID of the court.
     /// @param _coreDisputeID Index of the dispute in Kleros Core.
     /// @param _nonce Nonce to hash with random number.
     /// @return drawnAddress The drawn address.
     /// `O(k * log_k(n))` where
     /// `k` is the maximum number of children per node in the tree,
     ///  and `n` is the maximum number of nodes ever appended.
     function draw(
-        bytes32 _key,
+        uint96 _courtID,
         uint256 _coreDisputeID,
         uint256 _nonce
     ) public view override returns (address drawnAddress, uint96 fromSubcourtID) {
         if (phase != Phase.drawing) revert NotDrawingPhase();
-        SortitionSumTree storage tree = sortitionSumTrees[_key];
-
-        if (tree.nodes[0] == 0) {
-            return (address(0), 0); // No jurors staked.
-        }
-
-        uint256 currentDrawnNumber = uint256(keccak256(abi.encodePacked(randomNumber, _coreDisputeID, _nonce))) %
-            tree.nodes[0];
-
-        // While it still has children
-        uint256 treeIndex = 0;
-        while ((tree.K * treeIndex) + 1 < tree.nodes.length) {
-            for (uint256 i = 1; i <= tree.K; i++) {
-                // Loop over children.
-                uint256 nodeIndex = (tree.K * treeIndex) + i;
-                uint256 nodeValue = tree.nodes[nodeIndex];
-
-                if (currentDrawnNumber >= nodeValue) {
-                    // Go to the next child.
-                    currentDrawnNumber -= nodeValue;
-                } else {
-                    // Pick this child.
-                    treeIndex = nodeIndex;
-                    break;
-                }
-            }
-        }
 
-        bytes32 stakePathID = tree.nodeIndexesToIDs[treeIndex];
-        (drawnAddress, fromSubcourtID) = _stakePathIDToAccountAndCourtID(stakePathID);
+        TreeKey key = CourtID.wrap(_courtID).toTreeKey();
+        (drawnAddress, fromSubcourtID) = sortitionSumTrees[key].draw(_coreDisputeID, _nonce, randomNumber);
     }
 
     /// @dev Get the stake of a juror in a court.
     /// @param _juror The address of the juror.
     /// @param _courtID The ID of the court.
     /// @return value The stake of the juror in the court.
     function stakeOf(address _juror, uint96 _courtID) public view returns (uint256) {
-        bytes32 stakePathID = _accountAndCourtIDToStakePathID(_juror, _courtID);
-        return stakeOf(bytes32(uint256(_courtID)), stakePathID);
-    }
-
-    /// @dev Get the stake of a juror in a court.
-    /// @param _key The key of the tree, corresponding to a court.
-    /// @param _stakePathID The stake path ID, corresponding to a juror.
-    /// @return The stake of the juror in the court.
-    function stakeOf(bytes32 _key, bytes32 _stakePathID) public view returns (uint256) {
-        SortitionSumTree storage tree = sortitionSumTrees[_key];
-        uint treeIndex = tree.IDsToNodeIndexes[_stakePathID];
-        if (treeIndex == 0) {
-            return 0;
-        }
-        return tree.nodes[treeIndex];
+        bytes32 stakePathID = SortitionTrees.toStakePathID(_juror, _courtID);
+        TreeKey key = CourtID.wrap(_courtID).toTreeKey();
+        return sortitionSumTrees[key].stakeOf(stakePathID);
     }
 
     /// @dev Gets the balance of a juror in a court.
@@ -590,26 +544,6 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     // *            Internal               * //
     // ************************************* //
 
-    /// @dev Update all the parents of a node.
-    /// @param _key The key of the tree to update.
-    /// @param _treeIndex The index of the node to start from.
-    /// @param _plusOrMinus Whether to add (true) or substract (false).
-    /// @param _value The value to add or substract.
-    /// `O(log_k(n))` where
-    /// `k` is the maximum number of children per node in the tree,
-    ///  and `n` is the maximum number of nodes ever appended.
-    function _updateParents(bytes32 _key, uint256 _treeIndex, bool _plusOrMinus, uint256 _value) private {
-        SortitionSumTree storage tree = sortitionSumTrees[_key];
-
-        uint256 parentIndex = _treeIndex;
-        while (parentIndex != 0) {
-            parentIndex = (parentIndex - 1) / tree.K;
-            tree.nodes[parentIndex] = _plusOrMinus
-                ? tree.nodes[parentIndex] + _value
-                : tree.nodes[parentIndex] - _value;
-        }
-    }
-
     function _extraDataToTreeK(bytes memory _extraData) internal pure returns (uint256 K) {
         if (_extraData.length >= 32) {
             assembly {
@@ -621,151 +555,6 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
         }
     }
 
-    /// @dev Set a value in a tree.
-    /// @param _key The key of the tree.
-    /// @param _value The new value.
-    /// @param _stakePathID The ID of the value.
-    /// `O(log_k(n))` where
-    /// `k` is the maximum number of children per node in the tree,
-    ///  and `n` is the maximum number of nodes ever appended.
-    function _set(bytes32 _key, uint256 _value, bytes32 _stakePathID) internal {
-        SortitionSumTree storage tree = sortitionSumTrees[_key];
-        uint256 treeIndex = tree.IDsToNodeIndexes[_stakePathID];
-
-        if (treeIndex == 0) {
-            // No existing node.
-            if (_value != 0) {
-                // Non zero value.
-                // Append.
-                // Add node.
-                if (tree.stack.length == 0) {
-                    // No vacant spots.
-                    // Get the index and append the value.
-                    treeIndex = tree.nodes.length;
-                    tree.nodes.push(_value);
-
-                    // Potentially append a new node and make the parent a sum node.
-                    if (treeIndex != 1 && (treeIndex - 1) % tree.K == 0) {
-                        // Is first child.
-                        uint256 parentIndex = treeIndex / tree.K;
-                        bytes32 parentID = tree.nodeIndexesToIDs[parentIndex];
-                        uint256 newIndex = treeIndex + 1;
-                        tree.nodes.push(tree.nodes[parentIndex]);
-                        delete tree.nodeIndexesToIDs[parentIndex];
-                        tree.IDsToNodeIndexes[parentID] = newIndex;
-                        tree.nodeIndexesToIDs[newIndex] = parentID;
-                    }
-                } else {
-                    // Some vacant spot.
-                    // Pop the stack and append the value.
-                    treeIndex = tree.stack[tree.stack.length - 1];
-                    tree.stack.pop();
-                    tree.nodes[treeIndex] = _value;
-                }
-
-                // Add label.
-                tree.IDsToNodeIndexes[_stakePathID] = treeIndex;
-                tree.nodeIndexesToIDs[treeIndex] = _stakePathID;
-
-                _updateParents(_key, treeIndex, true, _value);
-            }
-        } else {
-            // Existing node.
-            if (_value == 0) {
-                // Zero value.
-                // Remove.
-                // Remember value and set to 0.
-                uint256 value = tree.nodes[treeIndex];
-                tree.nodes[treeIndex] = 0;
-
-                // Push to stack.
-                tree.stack.push(treeIndex);
-
-                // Clear label.
-                delete tree.IDsToNodeIndexes[_stakePathID];
-                delete tree.nodeIndexesToIDs[treeIndex];
-
-                _updateParents(_key, treeIndex, false, value);
-            } else if (_value != tree.nodes[treeIndex]) {
-                // New, non zero value.
-                // Set.
-                bool plusOrMinus = tree.nodes[treeIndex] <= _value;
-                uint256 plusOrMinusValue = plusOrMinus
-                    ? _value - tree.nodes[treeIndex]
-                    : tree.nodes[treeIndex] - _value;
-                tree.nodes[treeIndex] = _value;
-
-                _updateParents(_key, treeIndex, plusOrMinus, plusOrMinusValue);
-            }
-        }
-    }
-
-    /// @dev Packs an account and a court ID into a stake path ID: [20 bytes of address][12 bytes of courtID] = 32 bytes total.
-    /// @param _account The address of the juror to pack.
-    /// @param _courtID The court ID to pack.
-    /// @return stakePathID The stake path ID.
-    function _accountAndCourtIDToStakePathID(
-        address _account,
-        uint96 _courtID
-    ) internal pure returns (bytes32 stakePathID) {
-        assembly {
-            // solium-disable-line security/no-inline-assembly
-            let ptr := mload(0x40)
-
-            // Write account address (first 20 bytes)
-            for {
-                let i := 0x00
-            } lt(i, 0x14) {
-                i := add(i, 0x01)
-            } {
-                mstore8(add(ptr, i), byte(add(0x0c, i), _account))
-            }
-
-            // Write court ID (last 12 bytes)
-            for {
-                let i := 0x14
-            } lt(i, 0x20) {
-                i := add(i, 0x01)
-            } {
-                mstore8(add(ptr, i), byte(i, _courtID))
-            }
-            stakePathID := mload(ptr)
-        }
-    }
-
-    /// @dev Retrieves both juror's address and court ID from the stake path ID.
-    /// @param _stakePathID The stake path ID to unpack.
-    /// @return account The account.
-    /// @return courtID The court ID.
-    function _stakePathIDToAccountAndCourtID(
-        bytes32 _stakePathID
-    ) internal pure returns (address account, uint96 courtID) {
-        assembly {
-            // solium-disable-line security/no-inline-assembly
-            let ptr := mload(0x40)
-
-            // Read account address (first 20 bytes)
-            for {
-                let i := 0x00
-            } lt(i, 0x14) {
-                i := add(i, 0x01)
-            } {
-                mstore8(add(add(ptr, 0x0c), i), byte(i, _stakePathID))
-            }
-            account := mload(ptr)
-
-            // Read court ID (last 12 bytes)
-            for {
-                let i := 0x00
-            } lt(i, 0x0c) {
-                i := add(i, 0x01)
-            } {
-                mstore8(add(add(ptr, 0x14), i), byte(add(i, 0x14), _stakePathID))
-            }
-            courtID := mload(ptr)
-        }
-    }
-
     // ************************************* //
     // *              Errors               * //
     // ************************************* //
@@ -776,8 +565,6 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
     error NoDisputesThatNeedJurors();
     error RandomNumberNotReady();
     error DisputesWithoutJurorsAndMaxDrawingTimeNotPassed();
-    error TreeAlreadyExists();
-    error KMustBeGreaterThanOne();
     error NotStakingPhase();
     error NoDelayedStakeToExecute();
     error NotEligibleForWithdrawal();

contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol

@@ -232,9 +232,7 @@ abstract contract DisputeKitClassicBase is IDisputeKit, Initializable, UUPSProxi
 
         ISortitionModule sortitionModule = core.sortitionModule();
         (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
-        bytes32 key = bytes32(uint256(courtID)); // Get the ID of the tree.
-
-        (drawnAddress, fromSubcourtID) = sortitionModule.draw(key, _coreDisputeID, _nonce);
+        (drawnAddress, fromSubcourtID) = sortitionModule.draw(courtID, _coreDisputeID, _nonce);
         if (drawnAddress == address(0)) {
             // Sortition can return 0 address if no one has staked yet.
             return (drawnAddress, fromSubcourtID);