Skip to content

Commit a12e966

Browse files
jaybuidljaybuidl
committed
test: coverage of SortitionTrees and optimization of stakePathID packing/unpacking
1 parent e43b335 commit a12e966

File tree

3 files changed

+857
-41
lines changed

3 files changed

+857
-41
lines changed

‎contracts/src/libraries/SortitionTrees.sol

Lines changed: 5 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -205,26 +205,8 @@ library SortitionTrees {
205205
function toStakePathID(address _account, uint96 _courtID) internal pure returns (bytes32 stakePathID) {
206206
assembly {
207207
// solium-disable-line security/no-inline-assembly
208-
let ptr := mload(0x40)
209-
210-
// Write account address (first 20 bytes)
211-
for {
212-
let i := 0x00
213-
} lt(i, 0x14) {
214-
i := add(i, 0x01)
215-
} {
216-
mstore8(add(ptr, i), byte(add(0x0c, i), _account))
217-
}
218-
219-
// Write court ID (last 12 bytes)
220-
for {
221-
let i := 0x14
222-
} lt(i, 0x20) {
223-
i := add(i, 0x01)
224-
} {
225-
mstore8(add(ptr, i), byte(i, _courtID))
226-
}
227-
stakePathID := mload(ptr)
208+
// Pack address (20 bytes) and courtID (12 bytes) into a single bytes32
209+
stakePathID := or(shl(96, _account), _courtID)
228210
}
229211
}
230212

@@ -235,27 +217,9 @@ library SortitionTrees {
235217
function toAccountAndCourtID(bytes32 _stakePathID) internal pure returns (address account, uint96 courtID) {
236218
assembly {
237219
// solium-disable-line security/no-inline-assembly
238-
let ptr := mload(0x40)
239-
240-
// Read account address (first 20 bytes)
241-
for {
242-
let i := 0x00
243-
} lt(i, 0x14) {
244-
i := add(i, 0x01)
245-
} {
246-
mstore8(add(add(ptr, 0x0c), i), byte(i, _stakePathID))
247-
}
248-
account := mload(ptr)
249-
250-
// Read court ID (last 12 bytes)
251-
for {
252-
let i := 0x00
253-
} lt(i, 0x0c) {
254-
i := add(i, 0x01)
255-
} {
256-
mstore8(add(add(ptr, 0x14), i), byte(add(i, 0x14), _stakePathID))
257-
}
258-
courtID := mload(ptr)
220+
// Unpack address (first 20 bytes) and courtID (last 12 bytes) from the stake path ID
221+
account := shr(96, _stakePathID) // Right shift by 96 bits to get the address
222+
courtID := and(_stakePathID, 0xffffffffffffffffffffffff) // Mask the lower 96 bits to get the court ID
259223
}
260224
}
261225

‎contracts/src/test/SortitionTreesMock.sol

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity ^0.8.24;
3+
4+
import "../libraries/SortitionTrees.sol";
5+
6+
/// @title SortitionTreesMock
7+
/// @dev Test contract to expose SortitionTrees library functions for testing
8+
contract SortitionTreesMock {
9+
using SortitionTrees for mapping(TreeKey => SortitionTrees.Tree);
10+
11+
// Storage for multiple test trees
12+
mapping(TreeKey => SortitionTrees.Tree) public trees;
13+
14+
// Court hierarchy helpers (for testing parent-child relationships)
15+
mapping(uint96 => uint96[]) public childCourts;
16+
mapping(uint96 => uint96) public parentCourt;
17+
18+
// ************************************* //
19+
// * Main Functions * //
20+
// ************************************* //
21+
22+
/// @dev Create a sortition sum tree for a court
23+
function createTree(uint96 _courtID, uint256 _k) external {
24+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
25+
trees.createTree(key, _k);
26+
}
27+
28+
/// @dev Set stake for a juror in a specific court
29+
function set(uint96 _courtID, address _account, uint256 _value) external {
30+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
31+
bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID);
32+
SortitionTrees.set(trees[key], _value, stakePathID);
33+
}
34+
35+
/// @dev Draw a juror from a court's tree
36+
function draw(
37+
uint96 _courtID,
38+
uint256 _disputeID,
39+
uint256 _nonce,
40+
uint256 _randomNumber
41+
) external view returns (address drawnAddress, uint96 fromSubcourtID) {
42+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
43+
return SortitionTrees.draw(trees[key], _disputeID, _nonce, _randomNumber);
44+
}
45+
46+
/// @dev Get stake of a juror in a specific court
47+
function stakeOf(uint96 _courtID, address _account) external view returns (uint256) {
48+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
49+
bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID);
50+
return SortitionTrees.stakeOf(trees[key], stakePathID);
51+
}
52+
53+
// ************************************* //
54+
// * Multi-Court Operations * //
55+
// ************************************* //
56+
57+
/// @dev Set stake for a juror across multiple courts (for testing hierarchy)
58+
function setStakeInHierarchy(uint96[] calldata _courtIDs, address _account, uint256 _value) external {
59+
for (uint256 i = 0; i < _courtIDs.length; i++) {
60+
this.set(_courtIDs[i], _account, _value);
61+
}
62+
}
63+
64+
/// @dev Get stakes of a juror across multiple courts
65+
function getStakesAcrossCourts(
66+
address _account,
67+
uint96[] calldata _courtIDs
68+
) external view returns (uint256[] memory stakes) {
69+
stakes = new uint256[](_courtIDs.length);
70+
for (uint256 i = 0; i < _courtIDs.length; i++) {
71+
stakes[i] = this.stakeOf(_courtIDs[i], _account);
72+
}
73+
}
74+
75+
// ************************************* //
76+
// * Court Hierarchy Setup * //
77+
// ************************************* //
78+
79+
/// @dev Set parent court for testing hierarchy
80+
function setParentCourt(uint96 _childCourt, uint96 _parentCourt) external {
81+
parentCourt[_childCourt] = _parentCourt;
82+
}
83+
84+
/// @dev Add child court for testing hierarchy
85+
function addChildCourt(uint96 _parentCourt, uint96 _childCourt) external {
86+
childCourts[_parentCourt].push(_childCourt);
87+
}
88+
89+
/// @dev Get child courts
90+
function getChildCourts(uint96 _parentCourt) external view returns (uint96[] memory) {
91+
return childCourts[_parentCourt];
92+
}
93+
94+
// ************************************* //
95+
// * Tree State Inspection * //
96+
// ************************************* //
97+
98+
/// @dev Get all nodes in a tree
99+
function getTreeNodes(uint96 _courtID) external view returns (uint256[] memory) {
100+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
101+
return trees[key].nodes;
102+
}
103+
104+
/// @dev Get tree K value
105+
function getTreeK(uint96 _courtID) external view returns (uint256) {
106+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
107+
return trees[key].K;
108+
}
109+
110+
/// @dev Get tree stack
111+
function getTreeStack(uint96 _courtID) external view returns (uint256[] memory) {
112+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
113+
return trees[key].stack;
114+
}
115+
116+
/// @dev Get node index for a juror in a court
117+
function getNodeIndex(uint96 _courtID, address _account) external view returns (uint256) {
118+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
119+
bytes32 stakePathID = SortitionTrees.toStakePathID(_account, _courtID);
120+
return trees[key].IDsToNodeIndexes[stakePathID];
121+
}
122+
123+
/// @dev Check if a court tree exists
124+
function courtExists(uint96 _courtID) external view returns (bool) {
125+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
126+
return trees[key].K != 0;
127+
}
128+
129+
/// @dev Get the root sum (total stakes) of a court
130+
function getRootSum(uint96 _courtID) external view returns (uint256) {
131+
TreeKey key = CourtID.wrap(_courtID).toTreeKey();
132+
if (trees[key].nodes.length == 0) return 0;
133+
return trees[key].nodes[0];
134+
}
135+
136+
// ************************************* //
137+
// * Utility Functions * //
138+
// ************************************* //
139+
140+
/// @dev Test function to pack address and court ID
141+
function testToStakePathID(address _account, uint96 _courtID) external pure returns (bytes32) {
142+
return SortitionTrees.toStakePathID(_account, _courtID);
143+
}
144+
145+
/// @dev Test function to unpack stake path ID
146+
function testToAccountAndCourtID(bytes32 _stakePathID) external pure returns (address account, uint96 courtID) {
147+
return SortitionTrees.toAccountAndCourtID(_stakePathID);
148+
}
149+
150+
/// @dev Test function to convert court ID to tree key
151+
function testToTreeKey(uint96 _courtID) external pure returns (TreeKey) {
152+
return CourtID.wrap(_courtID).toTreeKey();
153+
}
154+
}

0 commit comments

Comments
 (0)