Sortition trees extracted to a library #2113
Conversation
✅ Deploy Preview for kleros-v2-testnet ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
✅ Deploy Preview for kleros-v2-testnet-devtools ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
WalkthroughRefactors sortition into a new SortitionTrees library, replaces bytes32 tree keys with uint96 courtIDs across createTree/draw/setStake APIs, updates draw to return (drawnAddress, fromSubcourtID), and migrates per-court stake/penalty operations to the library-backed implementation. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Core as KlerosCoreBase
participant DK as DisputeKitClassicBase
participant SM as SortitionModule
participant Lib as SortitionTrees (library)
Note over Core,DK: Juror selection
Core->>DK: request draw(coreDisputeID, nonce)
DK->>SM: draw(uint96 courtID, coreDisputeID, nonce)
SM->>Lib: trees[courtID].draw(coreDisputeID, nonce, rand)
Lib-->>SM: (drawnAddress, fromSubcourtID)
SM-->>DK: (drawnAddress, fromSubcourtID)
DK-->>Core: (drawnAddress, fromSubcourtID)
sequenceDiagram
autonumber
actor Core as KlerosCoreBase
participant SM as SortitionModule
participant Lib as SortitionTrees (library)
Note over Core,SM: Apply per-court stake penalty
Core->>SM: setStakePenalty(account, courtID, penalty)
SM->>Lib: trees[courtID].set(...) / stake update
Lib-->>SM: updated stake info
SM-->>Core: (pnkBalance, newCourtStake, availablePenalty)
alt newCourtStake < minStake
Core->>SM: setStake(account, courtID, 0, ...)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
🧪 Generate unit tests
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
jaybuidl
changed the base branch from
dev
to
fix/penalties-reduce-drawing-odds
jaybuidl
added
Type: Enhancement ✨
Compatibility: ABI change 🗯
❌ Deploy Preview for kleros-v2-university failed. Why did it fail? →
|
✅ Deploy Preview for kleros-v2-neo ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
contracts/src/arbitration/SortitionModuleBase.sol (1)
390-403: Potential issue with parent court iteration.In the
_setStakefunction, Line 395 uses_courtIDto create the tree key, but the loop iterates through parent courts viacurrentCourtID. This appears to be incorrect - the key should be created fromcurrentCourtIDfor each iteration.while (!finished) { // Tokens are also implicitly staked in parent courts through sortition module to increase the chance of being drawn. - TreeKey key = CourtID.wrap(_courtID).toTreeKey(); + TreeKey key = CourtID.wrap(currentCourtID).toTreeKey(); sortitionSumTrees[key].set(_newStake, stakePathID); if (currentCourtID == GENERAL_COURT) { finished = true; } else { (currentCourtID, , , , , , ) = core.courts(currentCourtID); // Get the parent court. } }
🧹 Nitpick comments (3)
contracts/src/libraries/SortitionTrees.sol (3)
51-53: Consider validating the tree state more explicitly.While returning
(address(0), 0)when the tree is empty is reasonable, consider whether this should revert with a descriptive error to make the failure mode more explicit to callers.if (_tree.nodes[0] == 0) { - return (address(0), 0); // No jurors staked. + revert NoJurorsStaked(); // More explicit error }Add the error definition:
error NoJurorsStaked();
219-248: Assembly implementation for unpacking looks correct but could be simplified.The
toAccountAndCourtIDfunction correctly unpacks the stake path ID, but like the packing function, it could benefit from a more straightforward implementation.Consider this simpler implementation:
function toAccountAndCourtID(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) - } + account = address(uint160(uint256(_stakePathID) >> 96)); + courtID = uint96(uint256(_stakePathID)); }
189-217: Byte layout verified; optional bitwise refactor recommended
- Verified that the current assembly implementation correctly packs the 20-byte address into bytes 0–19 and the 12-byte courtID into bytes 20–31 as intended.
- The proposed one-liner
yields an identical byte layout and is more concise and gas-efficient.stakePathID = bytes32(bytes20(_account)) | bytes32(uint256(_courtID));Feel free to replace the inline assembly with the above bitwise OR implementation if you’d like to simplify the code.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
contracts/src/arbitration/SortitionModuleBase.sol(5 hunks)contracts/src/libraries/SortitionTrees.sol(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: SonarCloud
🔇 Additional comments (7)
contracts/src/libraries/SortitionTrees.sol (2)
1-13: LGTM! Clean type definitions and global extension.The use of custom types
TreeKeyandCourtIDwith a global extension function provides type safety and clear abstractions for the sortition tree implementation.
88-157: Complex tree update logic handles edge cases well.The
setfunction correctly handles all cases: new nodes, removals, and updates. The logic for maintaining vacant spots via the stack and handling parent sum nodes is sound.contracts/src/arbitration/SortitionModuleBase.sol (5)
10-18: Good integration with the SortitionTrees library.The import and usage declarations properly set up the library integration, enabling clean delegation of tree operations.
465-474: Draw function correctly returns the subcourt information.The updated
drawfunction properly returns both the drawn address and the originating subcourt ID, which aligns with the penalty mechanism requirements.
493-507: New getJurorBalance function provides comprehensive juror information.The function efficiently returns all relevant juror balance information in a single call, which will help reduce external calls when querying juror state.
312-334: Verify penalty edge cases and inter-module assumptionsI updated the search to cover all
.solfiles (removing the unrecognized--type solidityfilter) and did not find any other checks or calls that assume_penaltyis always ≤juror.stakedPnk. However, absence of automated matches doesn’t guarantee there aren’t inter-module dependencies or off-chain tooling that rely on the old behavior. Please manually verify the following before approving:
- Review any contracts or libraries that call
setStakePenalty(or otherwise handle slashing) to ensure they correctly handleavailablePenaltybeing capped atjuror.stakedPnkeven when the input_penaltyis larger.- Add or update unit tests for scenarios where
_penalty>juror.stakedPnkto confirm no underflow occurs and event/state updates remain correct.- Check off-chain integrations (scripts, indexers, UI) that read the penalty values—confirm they don’t break if they encounter a capped penalty.
Once these points have been confirmed, the new penalty mechanism can be approved.
183-189: No additional guards needed forcreateTree
ThesortitionSumTrees.createTreelibrary call already reverts on duplicate trees (TreeAlreadyExists), and all invocations fromKlerosCoreBasehappen exactly once per court ID:
- Initial FORKING_COURT creation in the constructor
- GENERAL_COURT setup immediately afterward
- Each new subcourt uses a fresh
courtID(derived fromcourts.length) and callscreateTreeonly once.Since there are no code paths that attempt to recreate a tree for the same key, the existing revert behavior sufficiently prevents duplicates.
jaybuidl
force-pushed
the
refactor/sortition-trees-library
branch
from
1d86159 to
9d15604
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
contracts/src/arbitration/SortitionModuleBase.sol (1)
313-334: Penalty can push stakedPnk below lockedPnk (invariant breach).setStakePenalty deducts from juror.stakedPnk without respecting juror.lockedPnk, allowing stakedPnk < lockedPnk (e.g., staked=100, locked=90, penalty=50 → staked=50 < 90). This breaks the locked≤staked invariant enforced elsewhere and can cause underflows or inconsistent accounting.
Apply:
function setStakePenalty( address _account, uint96 _courtID, uint256 _penalty ) external override onlyByCore returns (uint256 pnkBalance, uint256 newCourtStake, uint256 availablePenalty) { Juror storage juror = jurors[_account]; - availablePenalty = _penalty; - newCourtStake = stakeOf(_account, _courtID); - if (juror.stakedPnk < _penalty) { - availablePenalty = juror.stakedPnk; - } + // Do not deduct below the locked amount. + uint256 maxPenalty = juror.stakedPnk > juror.lockedPnk ? juror.stakedPnk - juror.lockedPnk : 0; + availablePenalty = _penalty > maxPenalty ? maxPenalty : _penalty; + newCourtStake = stakeOf(_account, _courtID); if (availablePenalty == 0) return (juror.stakedPnk, newCourtStake, 0); // No penalty to apply. - uint256 currentStake = stakeOf(_account, _courtID); - uint256 newStake = 0; - if (currentStake >= availablePenalty) { - newStake = currentStake - availablePenalty; - } + uint256 currentStake = newCourtStake; + uint256 newStake = currentStake > availablePenalty ? currentStake - availablePenalty : 0; _setStake(_account, _courtID, 0, availablePenalty, newStake); pnkBalance = juror.stakedPnk; // updated by _setStake() newCourtStake = stakeOf(_account, _courtID); // updated by _setStake() }
🧹 Nitpick comments (5)
contracts/src/arbitration/SortitionModuleBase.sol (5)
183-189: Guard K locally or confirm library-level validation.If SortitionTrees.createTree reverts for K < 2, fine; otherwise add a local check to fail fast with a clear reason.
Example:
function createTree(uint96 _courtID, bytes memory _extraData) external override onlyByCore { TreeKey key = CourtID.wrap(_courtID).toTreeKey(); uint256 K = _extraDataToTreeK(_extraData); + require(K > 1, "Sortition: K must be > 1"); sortitionSumTrees.createTree(key, K); }
336-345: Fix Natspec: function describes reward but performs penalty.The header says “reward deposit” for setStakePenalty; it should say “penalty deduction.”
-/// @dev Update the state of the stakes with a PNK reward deposit, called by KC during rewards execution. +/// @dev Update the state of the stakes with a PNK penalty deduction, called by KC during penalties execution.
389-397: OK; consider small gas nits in tree updates.Loop is correct. Minor nits:
- Cache toTreeKey for GENERAL_COURT if hot.
- A do/while eliminates the finished flag.
-bool finished = false; -uint96 currentCourtID = _courtID; -while (!finished) { +uint96 currentCourtID = _courtID; +for (;;) { TreeKey key = CourtID.wrap(currentCourtID).toTreeKey(); sortitionSumTrees[key].set(_newStake, stakePathID); - if (currentCourtID == GENERAL_COURT) { - finished = true; - } else { - (currentCourtID, , , , , , ) = core.courts(currentCourtID); // Get the parent court. - } + if (currentCourtID == GENERAL_COURT) break; + (currentCourtID, , , , , , ) = core.courts(currentCourtID); // parent }
63-75: Unify courtID types in events.StakeSet uses uint256 for _courtID while the rest of the API moved to uint96 (StakeDelayed already uses uint96). Standardize to uint96 to avoid decoding friction and ABI surprises.
-event StakeSet(address indexed _address, uint256 _courtID, uint256 _amount, uint256 _amountAllCourts); +event StakeSet(address indexed _address, uint96 _courtID, uint256 _amount, uint256 _amountAllCourts);Confirm downstream indexers/off-chain consumers are updated.
Also applies to: 403-404
379-385: Micro gas: unchecked i-- in swap-and-pop loop.Safe in 0.8.x and saves gas.
for (uint256 i = juror.courtIDs.length; i > 0; i--) { - if (juror.courtIDs[i - 1] == _courtID) { + unchecked { // i > 0 + if (juror.courtIDs[i - 1] == _courtID) { + } juror.courtIDs[i - 1] = juror.courtIDs[juror.courtIDs.length - 1]; juror.courtIDs.pop(); break; - } + } }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
contracts/src/arbitration/KlerosCoreBase.sol(4 hunks)contracts/src/arbitration/SortitionModuleBase.sol(5 hunks)contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol(1 hunks)contracts/src/arbitration/interfaces/ISortitionModule.sol(2 hunks)contracts/src/arbitration/university/SortitionModuleUniversity.sol(2 hunks)contracts/src/libraries/SortitionTrees.sol(1 hunks)contracts/src/test/SortitionModuleMock.sol(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (6)
- contracts/src/arbitration/university/SortitionModuleUniversity.sol
- contracts/src/arbitration/interfaces/ISortitionModule.sol
- contracts/src/arbitration/dispute-kits/DisputeKitClassicBase.sol
- contracts/src/arbitration/KlerosCoreBase.sol
- contracts/src/test/SortitionModuleMock.sol
- contracts/src/libraries/SortitionTrees.sol
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: SonarCloud
🔇 Additional comments (5)
contracts/src/arbitration/SortitionModuleBase.sol (5)
10-10: Library extraction and typed keys look good.The switch to SortitionTrees with typed TreeKey/CourtID improves safety; mapping extension usage is idiomatic. No issues spotted here.
Also applies to: 17-19, 54-54
466-474: Confirm randomness domain separation in SortitionTrees.draw.Ensure the library mixes randomNumber, _coreDisputeID, and _nonce (e.g., keccak(randomNumber, _coreDisputeID, _nonce)) to prevent cross-dispute/court correlation and replay across nonces.
Would you like a quick property-based test that asserts different (_coreDisputeID, _nonce) pairs produce independent selections for a fixed tree?
481-484: stakeOf port looks correct.Path ID derivation + per-court lookup align with the library model.
494-507: getJurorBalance surface is coherent.The values reflect juror struct + library-based stake; no issues.
183-189: No issues found in SortitionModuleBase.createTree and related flows
Usage ofcreateTree,draw,setStakePenalty, and theStakeSetevent is consistent across contracts, libraries, interfaces, and tests;_extraDataToTreeKcorrectly defaults toDEFAULT_K = 6; no unintendedstakedPnkvslockedPnkassumptions detected.
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (5)
contracts/CHANGELOG.md (1)
19-20: Changelog wording mismatches the implementationThe library functions still take a
TreeKey(bytes32) forcreateTree(mapping(TreeKey => Tree), TreeKey _key, ...)anddraw(Tree storage, ...)returns(address,uint96); the “Replace bytes32 _key by uint96 courtID” line makes it sound like the library’s own signatures changed. Clarify that the external module APIs switched touint96 courtID, while the library keepsTreeKeyand providesCourtID.toTreeKey()for conversion.Apply:
-- Replace the `bytes32 _key` parameter in `SortitionTrees.createTree()` and `SortitionTrees.draw()` by `uint96 courtID` ([#2113](...)) -– Extract the sortition sum trees logic into a library `SortitionTrees` ([#2113](...)) +- Replace `bytes32` court keys with `uint96 courtID` in module-facing APIs; the library keeps `TreeKey (bytes32)` and adds `CourtID.toTreeKey()` to convert. ([#2113](...)) +- Extract the sortition-sum-tree logic into the `SortitionTrees` library. ([#2113](...))contracts/src/libraries/SortitionTrees.sol (4)
44-53: Document_randomNumberparam
draw()takes_randomNumberbut it’s undocumented in the NatSpec.-/// @param _nonce Nonce to hash with random number. +/// @param _nonce Nonce to hash with random number. +/// @param _randomNumber External randomness seed (e.g., from RNG/VRF).
167-183: Tiny typo in comments and minor clarity“substract” -> “subtract”. Also consider early-returning when
_value == 0to skip the loop, but optional.- /// @param _plusOrMinus Whether to add (true) or substract (false). - /// @param _value The value to add or substract. + /// @param _plusOrMinus Whether to add (true) or subtract (false). + /// @param _value The value to add or subtract.
36-42: Consider returning the created Tree reference or emitting an event (optional)Returning
Tree storage(or emitting an event) can simplify call sites and aid debugging, but optional.
266-268: Add specific error for uninitialized tree (optional)If you adopt the guard in
set(), considererror TreeNotInitialized();to distinguish fromKMustBeGreaterThanOne().
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
contracts/CHANGELOG.md(1 hunks)contracts/src/libraries/SortitionTrees.sol(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: SonarCloud
🔇 Additional comments (2)
contracts/CHANGELOG.md (1)
21-22: Align pragma with toolchain bumpYou note Hardhat/Foundry set to 0.8.30, but
SortitionTrees.solusespragma ^0.8.24. Bump the library to^0.8.30to avoid mismatched compiler behaviors (e.g., IR pipeline) and ensure consistent bytecode.Do you want me to raise a tiny follow-up PR to bump the pragma across the new library and any dependents?
contracts/src/libraries/SortitionTrees.sol (1)
54-61: Prefer reverting when no stake exists rather than returning zero (behavioral choice)Current behavior returns
(address(0), 0)if root sum is zero. If consumers assume a revert on “no jurors,” define a dedicated error (e.g.,NothingToDraw()). If returning zero is intentional, document it in NatSpec.Do downstream modules (e.g., DisputeKitClassicBase) expect a revert or a sentinel value here?
|
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (5)
contracts/src/test/SortitionTreesMock.sol (2)
9-13: Consistent library invocation style.You’re mixing method-style (
trees.createTree(...)) and free-function calls (SortitionTrees.set(...),SortitionTrees.stakeOf(...)). For readability, pick one style (e.g., addusing SortitionTrees for SortitionTrees.Tree;and calltrees[key].set(...)/trees[key].stakeOf(...)).Example:
+ using SortitionTrees for SortitionTrees.Tree; @@ - bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID); - SortitionTrees.set(trees[key], _value, stakePathID); + bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID); + trees[key].set(_value, stakePathID); @@ - bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID); - return SortitionTrees.stakeOf(trees[key], stakePathID); + bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID); + return trees[key].stakeOf(stakePathID);Also applies to: 22-33, 46-51
98-114: Direct struct field access is fine for a mock; keep it strictly test-only.Accessing
nodes,stackandIDsToNodeIndexescouples tests to internals that may change. If these structs evolve, tests will break. Consider adding thin getters in the library or clearly marking these helpers as test-only and keeping them out of prod code.Also applies to: 116-134
contracts/test/sortition/index.ts (3)
7-7: Strongly type signers.Avoid
any[]. UseHardhatEthersSigner[]for better IDE/type safety.Apply:
- let accounts: any[]; + import type { HardhatEthersSigner } from "@nomicfoundation/hardhat-ethers/signers"; + let accounts: HardhatEthersSigner[];
279-283: Use bigint for courtID in empty-tree draw assertion.
courtIDis decoded asbigint.- expect(courtID).to.equal(0); + expectBnEq(courtID, 0n);
316-334: Reduce flakiness in weighted draw test.With only 100 draws, deterministic seeds could still yield borderline counts. Increase draws and set a tolerance band (e.g., >90%) to reduce false negatives.
- const numDraws = 100; + const numDraws = 512; @@ - expect(draws[getTestAddress(1).toLowerCase()]).to.be.greaterThan(numDraws * 0.8); // At least 80% for high stake + expect(draws[getTestAddress(1).toLowerCase()]).to.be.greaterThan(numDraws * 0.9); // Heavier weight dominates
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
contracts/src/libraries/SortitionTrees.sol(1 hunks)contracts/src/test/SortitionTreesMock.sol(1 hunks)contracts/test/sortition/index.ts(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- contracts/src/libraries/SortitionTrees.sol
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: SonarCloud
- GitHub Check: Mend Security Check
🔇 Additional comments (3)
contracts/src/test/SortitionTreesMock.sol (1)
41-44: Drop the type-mismatch suggestion
Thedrawfunction in contracts/src/libraries/SortitionTrees.sol is defined to return(address drawnAddress, uint96 fromSubcourtID), so the external wrapper’s signature already matches and will compile as-is.Likely an incorrect or invalid review comment.
contracts/test/sortition/index.ts (2)
510-519: Clarify expectations for fromSubcourtID.The checks
include([0,1])suggest ambiguity. Ifdrawshould return the exact subcourt used, assert deterministically per setup. Otherwise, document the intended behavior for hierarchical draws.Would you confirm whether
fromSubcourtIDshould always equal the requested_courtIDin this mock (no parent aggregation)? If yes, replace these with explicit equals.
3-3: Fix TypeChain import path and use a type-only import– import { SortitionTreesMock } from "../../typechain-types"; + import type { SortitionTreesMock } from "../../../typechain-types";Ensure the
typechain-typesdirectory lives at your project root (i.e. three levels up fromcontracts/test/sortition/index.ts).
| /// @dev Set stake for a juror across multiple courts (for testing hierarchy) | ||
| function setStakeInHierarchy(uint96[] calldata _courtIDs, address _account, uint256 _value) external { | ||
| for (uint256 i = 0; i < _courtIDs.length; i++) { | ||
| this.set(_courtIDs[i], _account, _value); | ||
| } | ||
| } | ||
|
|
||
| /// @dev Get stakes of a juror across multiple courts | ||
| function getStakesAcrossCourts( | ||
| address _account, | ||
| uint96[] calldata _courtIDs | ||
| ) external view returns (uint256[] memory stakes) { | ||
| stakes = new uint256[](_courtIDs.length); | ||
| for (uint256 i = 0; i < _courtIDs.length; i++) { | ||
| stakes[i] = this.stakeOf(_courtIDs[i], _account); | ||
| } | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Avoid external self-calls; make set/stakeOf callable internally.
Using this.set/this.stakeOf triggers external calls, adds gas and complicates reentrancy analysis (even in tests). Prefer internal/public calls.
Apply:
- function set(uint96 _courtID, address _account, uint256 _value) external {
+ function set(uint96 _courtID, address _account, uint256 _value) public {
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID);
SortitionTrees.set(trees[key], _value, stakePathID);
}
- function stakeOf(uint96 _courtID, address _account) external view returns (uint256) {
+ function stakeOf(uint96 _courtID, address _account) public view returns (uint256) {
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID);
return SortitionTrees.stakeOf(trees[key], stakePathID);
}- function setStakeInHierarchy(uint96[] calldata _courtIDs, address _account, uint256 _value) external {
+ function setStakeInHierarchy(uint96[] calldata _courtIDs, address _account, uint256 _value) external {
for (uint256 i = 0; i < _courtIDs.length; i++) {
- this.set(_courtIDs[i], _account, _value);
+ set(_courtIDs[i], _account, _value);
}
}- ) external view returns (uint256[] memory stakes) {
+ ) external view returns (uint256[] memory stakes) {
stakes = new uint256[](_courtIDs.length);
for (uint256 i = 0; i < _courtIDs.length; i++) {
- stakes[i] = this.stakeOf(_courtIDs[i], _account);
+ stakes[i] = stakeOf(_courtIDs[i], _account);
}
}Also applies to: 47-51
🤖 Prompt for AI Agents
In contracts/src/test/SortitionTreesMock.sol around lines 57-73 (and also apply
same fix to 47-51), replace external self-calls this.set(...) and
this.stakeOf(...) with direct internal calls set(...) and stakeOf(...), and
update the declarations of set and stakeOf so they are callable internally
(change visibility from external to public or internal as appropriate); ensure
any required state/return semantics remain unchanged and adjust tests if
visibility change affects imports or mocks.
| // Helper function to create a test juror address | ||
| const getTestAddress = (index: number): string => accounts[index % accounts.length].address; | ||
|
|
There was a problem hiding this comment.
Add BigInt helpers; many numeric assertions compare bigint to number.
Most contract calls return bigint. Numerous expect(...).to.equal(<number>) will fail. Introduce helpers and use them across tests.
Apply:
// Helper function to create a test juror address
const getTestAddress = (index: number): string => accounts[index % accounts.length].address;
+ // BigInt helpers for clean assertions
+ const bn = (x: number | bigint) => BigInt(x);
+ const sumBn = (arr: Array<number | bigint>) => arr.reduce<bigint>((acc, v) => acc + bn(v), 0n);
+ const expectBnEq = (actual: bigint, expected: number | bigint) => expect(actual).to.equal(bn(expected));Then, fix representative hotspots (apply pattern throughout):
- expect(unpackedCourtID).to.equal(testCase.courtID);
+ expectBnEq(unpackedCourtID, testCase.courtID);
@@
- expect(nodes[0]).to.equal(0); // Root starts at 0
+ expectBnEq(nodes[0], 0n); // Root starts at 0
@@
- const stake = 100;
+ const stake = 100n;
@@
- expect(await sortitionTree.stakeOf(0, juror)).to.equal(stake);
- expect(await sortitionTree.getRootSum(0)).to.equal(stake);
+ expectBnEq(await sortitionTree.stakeOf(0, juror), stake);
+ expectBnEq(await sortitionTree.getRootSum(0), stake);
@@
- const expectedTotal = stakes.reduce((sum, stake) => sum + stake, 0);
- expect(await sortitionTree.getRootSum(0)).to.equal(expectedTotal);
+ const expectedTotal = sumBn(stakes);
+ expectBnEq(await sortitionTree.getRootSum(0), expectedTotal);
@@
- expect(await sortitionTree.getNodeIndex(0, juror)).to.equal(0); // Should be cleared
+ expectBnEq(await sortitionTree.getNodeIndex(0, juror), 0n); // Should be cleared
@@
- expect(courtID).to.equal(0);
+ expectBnEq(courtID, 0n);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Helper function to create a test juror address | |
| const getTestAddress = (index: number): string => accounts[index % accounts.length].address; | |
| // Helper function to create a test juror address | |
| const getTestAddress = (index: number): string => | |
| accounts[index % accounts.length].address; | |
| // BigInt helpers for clean assertions | |
| const bn = (x: number | bigint): bigint => BigInt(x); | |
| const sumBn = (arr: Array<number | bigint>): bigint => | |
| arr.reduce<bigint>((acc, v) => acc + bn(v), 0n); | |
| const expectBnEq = (actual: bigint, expected: number | bigint): void => | |
| expect(actual).to.equal(bn(expected)); |
🤖 Prompt for AI Agents
In contracts/test/sortition/index.ts around lines 15-17 add small BigInt helper
functions at the top of the test file (e.g. const toBig = (n: number) =>
BigInt(n); and const fromBig = (b: bigint) => Number(b); or similar) and replace
numeric assertions that compare contract-returned bigints to plain numbers with
the helper (e.g. expect(result).to.equal(toBig(5)) or
expect(fromBig(result)).to.equal(5)); add the helpers next to getTestAddress and
update the representative hotspots in this file to use toBig()/fromBig()
consistently so bigint vs number mismatches are fixed.
| it("Should increase stake values correctly", async function () { | ||
| const juror = getTestAddress(0); | ||
| const oldStake = await sortitionTree.stakeOf(0, juror); | ||
| const newStake = 250; | ||
|
|
||
| await sortitionTree.set(0, juror, newStake); | ||
|
|
||
| expect(await sortitionTree.stakeOf(0, juror)).to.equal(newStake); | ||
|
|
||
| // Root should reflect the change | ||
| const rootSum = await sortitionTree.getRootSum(0); | ||
| expect(rootSum).to.equal(200 + newStake); | ||
| }); |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Fix arithmetic with mixed types in assertions.
Expressions like 200 + newStake are numbers, but rootSum is bigint. Convert operands to bigint.
Apply:
- expect(rootSum).to.equal(200 + newStake);
+ expectBnEq(rootSum, 200n + bn(newStake));
@@
- expect(await sortitionTree.getRootSum(0)).to.equal(200 + newStake);
+ expectBnEq(await sortitionTree.getRootSum(0), 200n + bn(newStake));Also applies to: 205-213
🤖 Prompt for AI Agents
In contracts/test/sortition/index.ts around lines 191 to 203, the assertion uses
mixed numeric types (e.g., 200 + newStake) while rootSum is a bigint; update the
assertion to perform the arithmetic with bigints (convert numeric literals and
newStake to BigInt before adding) so the comparison uses the same type, and
apply the same conversion fixes to the similar assertions around lines 205-213.
| it("Should handle boundary values correctly", async function () { | ||
| await sortitionTree.createTree(0, 2); | ||
| const juror = getTestAddress(0); | ||
|
|
||
| // Test with maximum stake value | ||
| const maxStake = ethers.MaxUint256; | ||
|
|
||
| // This might fail due to gas limits, but should not revert due to overflow | ||
| // Note: In practice, stakes would be much smaller | ||
| try { | ||
| await sortitionTree.set(0, juror, maxStake); | ||
| expect(await sortitionTree.stakeOf(0, juror)).to.equal(maxStake); | ||
| } catch (error) { | ||
| // Expected to fail due to gas limits, not due to overflow | ||
| expect(error).to.match(/gas/i); | ||
| } | ||
| }); |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Make the max-stake test deterministic.
Catching any error and matching /gas/i is brittle; revert messages and error types vary. If the goal is “no arithmetic overflow”, assert success. If you want to prove it can handle large values, pick a very large but reasonable stake (e.g., 2n**200n) to avoid potential env-specific gas issues.
Apply:
- // This might fail due to gas limits, but should not revert due to overflow
- // Note: In practice, stakes would be much smaller
- try {
- await sortitionTree.set(0, juror, maxStake);
- expect(await sortitionTree.stakeOf(0, juror)).to.equal(maxStake);
- } catch (error) {
- // Expected to fail due to gas limits, not due to overflow
- expect(error).to.match(/gas/i);
- }
+ // Assert no overflow on very large values.
+ await sortitionTree.set(0, juror, maxStake);
+ expect(await sortitionTree.stakeOf(0, juror)).to.equal(maxStake);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| it("Should handle boundary values correctly", async function () { | |
| await sortitionTree.createTree(0, 2); | |
| const juror = getTestAddress(0); | |
| // Test with maximum stake value | |
| const maxStake = ethers.MaxUint256; | |
| // This might fail due to gas limits, but should not revert due to overflow | |
| // Note: In practice, stakes would be much smaller | |
| try { | |
| await sortitionTree.set(0, juror, maxStake); | |
| expect(await sortitionTree.stakeOf(0, juror)).to.equal(maxStake); | |
| } catch (error) { | |
| // Expected to fail due to gas limits, not due to overflow | |
| expect(error).to.match(/gas/i); | |
| } | |
| }); | |
| it("Should handle boundary values correctly", async function () { | |
| await sortitionTree.createTree(0, 2); | |
| const juror = getTestAddress(0); | |
| // Test with maximum stake value | |
| const maxStake = ethers.MaxUint256; | |
| // Assert no overflow on very large values. | |
| await sortitionTree.set(0, juror, maxStake); | |
| expect(await sortitionTree.stakeOf(0, juror)).to.equal(maxStake); | |
| }); |
| // Verify tree integrity | ||
| let expectedTotal = 0; | ||
| const finalStakes = new Map(); | ||
|
|
||
| // Calculate expected final state | ||
| for (const op of operations) { | ||
| const prevStake = finalStakes.get(op.juror) || 0; | ||
| expectedTotal = expectedTotal - prevStake + op.stake; | ||
| finalStakes.set(op.juror, op.stake); | ||
| } | ||
|
|
||
| // Verify actual state matches expected | ||
| expect(await sortitionTree.getRootSum(0)).to.equal(expectedTotal); | ||
|
|
||
| for (const [juror, expectedStake] of finalStakes) { | ||
| expect(await sortitionTree.stakeOf(0, juror)).to.equal(expectedStake); | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use bigint consistently in stress/rapid-change sections.
expectedTotal and map values are numbers while on-chain values are bigint.
Apply:
- let expectedTotal = 0;
- const finalStakes = new Map();
+ let expectedTotal = 0n;
+ const finalStakes = new Map<string, bigint>();
@@
- const prevStake = finalStakes.get(op.juror) || 0;
- expectedTotal = expectedTotal - prevStake + op.stake;
- finalStakes.set(op.juror, op.stake);
+ const prevStake = finalStakes.get(op.juror) ?? 0n;
+ expectedTotal = expectedTotal - prevStake + bn(op.stake);
+ finalStakes.set(op.juror, bn(op.stake));- const stakeSequence = [100, 0, 200, 300, 0, 150, 0, 500];
+ const stakeSequence = [100n, 0n, 200n, 300n, 0n, 150n, 0n, 500n];
@@
- expect(await sortitionTree.stakeOf(0, juror)).to.equal(stake);
- expect(await sortitionTree.getRootSum(0)).to.equal(stake);
+ expectBnEq(await sortitionTree.stakeOf(0, juror), stake);
+ expectBnEq(await sortitionTree.getRootSum(0), stake);Also applies to: 688-695
🤖 Prompt for AI Agents
In contracts/test/sortition/index.ts around lines 665-681 (and also apply same
changes to 688-695), the test mixes JS numbers with on-chain bigint values;
change expectedTotal to initialize as 0n, ensure finalStakes stores bigint
values (use BigInt(op.stake) if op.stake may be number), compute expectedTotal
using bigint arithmetic (subtract and add BigInt(prevStake) / BigInt(op.stake)
as needed), and when asserting convert or use bigint versions of expectedTotal
and expectedStake so comparisons against sortitionTree.getRootSum and stakeOf
use consistent bigint types.
Changes:
createTree()function, into the Dispute Kits through thedraw()function.PR-Codex overview
This PR focuses on refactoring the
SortitionModuleand related components to improve type safety and efficiency by replacingbytes32identifiers withuint96court IDs. It also extracts sorting logic into a library for better modularity.Detailed summary
createTreeanddrawfunctions to useuint96 _courtIDinstead ofbytes32 _key.SortitionSumTreestructure toSortitionTrees.Tree.SortitionTreesfor handling tree operations.SortitionTreeslibrary.SortitionTreesMock.Summary by CodeRabbit
New Features
Refactor
Tests