Skip to content

Commit

Permalink
chore: add getLocks & getLocksByOwner
Browse files Browse the repository at this point in the history
  • Loading branch information
pyk committed Jul 1, 2024
1 parent c21b6c5 commit fb9fc08
Show file tree
Hide file tree
Showing 3 changed files with 149 additions and 38 deletions.
76 changes: 40 additions & 36 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -47,50 +47,54 @@ Example output:
```
$ forge test
[⠊] Compiling...
[] Compiling 1 files with Solc 0.8.23
[] Solc 0.8.23 finished in 1.81s
[] Compiling 1 files with Solc 0.8.23
[] Solc 0.8.23 finished in 1.71s
Compiler run successful!
Ran 8 tests for test/Whitelist.t.sol:RewardDistributionTest
[PASS] test_disableWhitelist_InvalidAction() (gas: 20255)
[PASS] test_disableWhitelist_Unauthorized() (gas: 13677)
[PASS] test_disableWhitelist_Valid() (gas: 19431)
[PASS] test_lock_InvalidAction() (gas: 115236)
[PASS] test_lock_Valid() (gas: 222819)
[PASS] test_setRoot_InvalidAction() (gas: 15605)
[PASS] test_setRoot_Unauthorized() (gas: 13717)
[PASS] test_setRoot_Valid() (gas: 19647)
Suite result: ok. 8 passed; 0 failed; 0 skipped; finished in 10.53ms (848.79µs CPU time)
Ran 1 test for test/LlamaLocker.t.sol:LlamaLockerTest
[PASS] test_renounceOwnership_InvalidAction() (gas: 13344)
Suite result: ok. 1 passed; 0 failed; 0 skipped; finished in 10.54ms (480.79µs CPU time)
[PASS] test_renounceOwnership_InvalidAction() (gas: 13389)
Suite result: ok. 1 passed; 0 failed; 0 skipped; finished in 1.44ms (93.38µs CPU time)
Ran 4 tests for test/AddRewardTokens.t.sol:AddRewardTokensTest
[PASS] test_addRewardTokens_InvalidRewardToken() (gas: 92126)
[PASS] test_addRewardTokens_InvalidRewardTokenCount() (gas: 13931)
[PASS] test_addRewardTokens_Unauthorized() (gas: 14168)
[PASS] test_addRewardTokens_Valid() (gas: 133254)
Suite result: ok. 4 passed; 0 failed; 0 skipped; finished in 9.61ms (987.75µs CPU time)
[PASS] test_addRewardTokens_InvalidRewardToken() (gas: 92195)
[PASS] test_addRewardTokens_InvalidRewardTokenCount() (gas: 13954)
[PASS] test_addRewardTokens_Unauthorized() (gas: 14191)
[PASS] test_addRewardTokens_Valid() (gas: 133255)
Suite result: ok. 4 passed; 0 failed; 0 skipped; finished in 1.60ms (179.63µs CPU time)
Ran 5 tests for test/RewardDistribution.t.sol:RewardDistributionTest
[PASS] test_distributeRewardToken_Claimables() (gas: 971608)
[PASS] test_distributeRewardToken_InvalidRewardAmount() (gas: 134452)
[PASS] test_distributeRewardToken_InvalidRewardToken() (gas: 17533)
[PASS] test_distributeRewardToken_InvalidTotalShares() (gas: 136542)
[PASS] test_distributeRewardToken_Unauthorized() (gas: 13834)
Suite result: ok. 5 passed; 0 failed; 0 skipped; finished in 10.59ms (1.80ms CPU time)
Ran 8 tests for test/Whitelist.t.sol:RewardDistributionTest
[PASS] test_disableWhitelist_InvalidAction() (gas: 20211)
[PASS] test_disableWhitelist_Unauthorized() (gas: 13655)
[PASS] test_disableWhitelist_Valid() (gas: 19409)
[PASS] test_lock_InvalidAction() (gas: 115192)
[PASS] test_lock_Valid() (gas: 225023)
[PASS] test_setRoot_InvalidAction() (gas: 15583)
[PASS] test_setRoot_Unauthorized() (gas: 13695)
[PASS] test_setRoot_Valid() (gas: 19625)
Suite result: ok. 8 passed; 0 failed; 0 skipped; finished in 1.65ms (751.29µs CPU time)
Ran 6 tests for test/LockMechanism.t.sol:LockMechanismTest
[PASS] test_lock_InvalidTokenCount() (gas: 11476)
[PASS] test_lock_Valid() (gas: 449598)
[PASS] test_unlock_InvalidLockOwner() (gas: 223593)
[PASS] test_unlock_InvalidTokenCount() (gas: 9220)
[PASS] test_unlock_InvalidUnlockWindow() (gas: 243582)
[PASS] test_unlock_ValidUnlockWindow() (gas: 216046)
Suite result: ok. 6 passed; 0 failed; 0 skipped; finished in 9.69ms (1.18ms CPU time)
Ran 5 test suites in 719.53ms (50.95ms CPU time): 24 tests passed, 0 failed, 0 skipped (24 total tests)
[PASS] test_lock_InvalidTokenCount() (gas: 11454)
[PASS] test_lock_Valid() (gas: 496076)
[PASS] test_unlock_InvalidLockOwner() (gas: 225991)
[PASS] test_unlock_InvalidTokenCount() (gas: 9265)
[PASS] test_unlock_InvalidUnlockWindow() (gas: 246840)
[PASS] test_unlock_ValidUnlockWindow() (gas: 222618)
Suite result: ok. 6 passed; 0 failed; 0 skipped; finished in 1.71ms (784.38µs CPU time)
Ran 5 tests for test/RewardDistribution.t.sol:RewardDistributionTest
[PASS] test_distributeRewardToken_Claimables() (gas: 1020171)
[PASS] test_distributeRewardToken_InvalidRewardAmount() (gas: 134541)
[PASS] test_distributeRewardToken_InvalidRewardToken() (gas: 17599)
[PASS] test_distributeRewardToken_InvalidTotalShares() (gas: 136631)
[PASS] test_distributeRewardToken_Unauthorized() (gas: 13900)
Suite result: ok. 5 passed; 0 failed; 0 skipped; finished in 2.06ms (1.10ms CPU time)
Ran 1 test for test/OffchainQuery.sol:LockMechanismTest
[PASS] test_getLocks() (gas: 11133290)
Suite result: ok. 1 passed; 0 failed; 0 skipped; finished in 11.34ms (9.86ms CPU time)
Ran 6 test suites in 161.35ms (19.81ms CPU time): 25 tests passed, 0 failed, 0 skipped (25 total tests)
```

## Generate Merkle Tree
Expand Down
49 changes: 47 additions & 2 deletions src/LlamaLocker.sol
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step, ReentrancyGuard {
struct NFTLock {
address owner;
uint256 lockedAt;
uint256 tokenId;
}

struct Claimable {
Expand Down Expand Up @@ -102,7 +103,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step, ReentrancyGuard {

function _lockNFT(address owner_, uint256 tokenId_) private {
nft.safeTransferFrom(owner_, address(this), tokenId_);
locks[tokenId_] = NFTLock({owner: owner_, lockedAt: block.timestamp});
locks[tokenId_] = NFTLock({owner: owner_, lockedAt: block.timestamp, tokenId: tokenId_});
}

function _lock(address owner_, uint256[] calldata tokenIds_) private {
Expand Down Expand Up @@ -155,7 +156,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step, ReentrancyGuard {
uint256 modulo = lockedDurationInEpoch % LOCK_DURATION_IN_EPOCH;
if (modulo != 0) revert InvalidUnlockWindow();

locks[tokenId_] = NFTLock({owner: address(0), lockedAt: 0});
locks[tokenId_] = NFTLock({owner: address(0), lockedAt: 0, tokenId: 0});
nft.safeTransferFrom(address(this), owner_, tokenId_);
}

Expand Down Expand Up @@ -247,4 +248,48 @@ contract LlamaLocker is ERC721Holder, Ownable2Step, ReentrancyGuard {
results[i] = Claimable({token: token, amount: amount});
}
}

/// @dev Used for offchain only
function getLocks() external view returns (NFTLock[] memory results) {
uint256 lockedCount = 0;
uint256 maxSupply = 1111;

for (uint256 _tokenId = 0; _tokenId < maxSupply; _tokenId++) {
NFTLock memory locked = locks[_tokenId];
if (locked.owner == address(0)) continue;
lockedCount++;
}

results = new NFTLock[](lockedCount);
uint256 index = 0;
for (uint256 _tokenId = 0; _tokenId < maxSupply; _tokenId++) {
NFTLock memory locked = locks[_tokenId];
if (locked.owner == address(0)) continue;
results[index] = locked;
index++;
}
}

/// @dev Used for offchain only
function getLocksByOwner(address owner_) external view returns (NFTLock[] memory results) {
uint256 lockedCount = 0;
uint256 maxSupply = 1111;

for (uint256 _tokenId = 0; _tokenId < maxSupply; _tokenId++) {
NFTLock memory locked = locks[_tokenId];
if (locked.owner == address(0)) continue;
if (locked.owner != owner_) continue;
lockedCount++;
}

results = new NFTLock[](lockedCount);
uint256 index = 0;
for (uint256 _tokenId = 0; _tokenId < maxSupply; _tokenId++) {
NFTLock memory locked = locks[_tokenId];
if (locked.owner == address(0)) continue;
if (locked.owner != owner_) continue;
results[index] = locked;
index++;
}
}
}
62 changes: 62 additions & 0 deletions test/OffchainQuery.sol
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
// SPDX-License-Identifier: MIT
pragma solidity 0.8.23;

import {Test} from "forge-std/Test.sol";

import {LlamaLocker} from "../src/LlamaLocker.sol";
import {MockNFT} from "./MockNFT.sol";

contract LockMechanismTest is Test {
MockNFT public nft;
LlamaLocker public locker;

address public admin = makeAddr("admin");
address public alice = makeAddr("alice");
address public bob = makeAddr("bob");

function setUp() public {
nft = new MockNFT();
locker =
new LlamaLocker(admin, address(nft), 0x308b32884ae9d1e08a3e1d00ac6934fa79bf7ae30358cff55dda90c8aebedf9c);

vm.startPrank(admin);
locker.disableWhitelist();
vm.stopPrank();
}

function test_getLocks() public {
uint256 tokenId = nft.mint(alice);
uint256[] memory tokenIds = new uint256[](1);
tokenIds[0] = tokenId;

vm.startPrank(alice);
nft.setApprovalForAll(address(locker), true);
vm.stopPrank();

vm.startPrank(bob);
nft.setApprovalForAll(address(locker), true);
vm.stopPrank();

vm.startPrank(alice);
locker.lock(tokenIds);
vm.stopPrank();

tokenId = nft.mint(bob);
tokenIds = new uint256[](1);
tokenIds[0] = tokenId;

vm.startPrank(bob);
locker.lock(tokenIds);
vm.stopPrank();

LlamaLocker.NFTLock[] memory locks = locker.getLocks();
assertEq(locks[0].owner, alice, "invalid lock alice");
assertEq(locks[1].owner, bob, "invalid lock bob");

locks = locker.getLocksByOwner(alice);
assertEq(locks[0].owner, alice, "invalid lock alice");

locks = locker.getLocksByOwner(bob);
assertEq(locks[0].owner, bob, "invalid lock bob");
}
}

0 comments on commit fb9fc08

Please sign in to comment.