Skip to content

Sortition trees extracted to a library #2113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: fix/penalties-reduce-drawing-odds
Choose a base branch
from
Open

Sortition trees extracted to a library #2113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9d15604
refactor: extracted the sortition tree to a library
jaybuidl Aug 27, 2025
e43b335
chore: changelog, natspec decorators
jaybuidl Aug 28, 2025
a12e966
test: coverage of SortitionTrees and optimization of stakePathID pack…
jaybuidl Aug 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions contracts/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ The format is based on [Common Changelog](https://common-changelog.org/).
- Make the primary VRF-based RNG fall back to `BlockhashRNG` if the VRF request is not fulfilled within a timeout ([#2054](https://github.com/kleros/kleros-v2/issues/2054))
- Authenticate the calls to the RNGs to prevent 3rd parties from depleting the Chainlink VRF subscription funds ([#2054](https://github.com/kleros/kleros-v2/issues/2054))
- Use `block.timestamp` rather than `block.number` for `BlockhashRNG` for better reliability on Arbitrum as block production is sporadic depending on network conditions. ([#2054](https://github.com/kleros/kleros-v2/issues/2054))
- Replace the `bytes32 _key` parameter in `SortitionTrees.createTree()` and `SortitionTrees.draw()` by `uint96 courtID` ([#2113](https://github.com/kleros/kleros-v2/issues/2113))
- Extract the sortition sum trees logic into a library `SortitionTrees` ([#2113](https://github.com/kleros/kleros-v2/issues/2113))
- Set the Hardhat Solidity version to v0.8.30 and enable the IR pipeline ([#2069](https://github.com/kleros/kleros-v2/issues/2069))
- Set the Foundry Solidity version to v0.8.30 and enable the IR pipeline ([#2073](https://github.com/kleros/kleros-v2/issues/2073))
- Widen the allowed solc version to any v0.8.x for the interfaces only ([#2083](https://github.com/kleros/kleros-v2/issues/2083))
Expand Down
8 changes: 4 additions & 4 deletions contracts/src/arbitration/KlerosCoreBase.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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;
Expand Down Expand Up @@ -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++) {
Expand All @@ -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);
Expand Down
251 changes: 19 additions & 232 deletions contracts/src/arbitration/SortitionModuleBase.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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].
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -395,12 +387,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 {
Expand Down Expand Up @@ -462,71 +455,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.
Expand Down Expand Up @@ -575,26 +529,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 {
Expand All @@ -606,151 +540,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 * //
// ************************************* //
Expand All @@ -761,8 +550,6 @@ abstract contract SortitionModuleBase is ISortitionModule, Initializable, UUPSPr
error NoDisputesThatNeedJurors();
error RandomNumberNotReady();
error DisputesWithoutJurorsAndMaxDrawingTimeNotPassed();
error TreeAlreadyExists();
error KMustBeGreaterThanOne();
error NotStakingPhase();
error NoDelayedStakeToExecute();
error NotEligibleForWithdrawal();
Expand Down
4 changes: 1 addition & 3 deletions contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Loading
Loading