diff --git a/contracts/UnionLens.sol b/contracts/UnionLens.sol index 9397542e..ee714361 100644 --- a/contracts/UnionLens.sol +++ b/contracts/UnionLens.sol @@ -74,6 +74,34 @@ contract UnionLens { userInfo.accountBorrow = uToken.getBorrowed(user); } + function getBorrowerAddresses(address underlying, address account) public view returns (address[] memory) { + IUserManager userManager = IUserManager(marketRegistry.userManagers(underlying)); + + uint256 voucheeCount = userManager.getVoucheeCount(account); + address[] memory addresses = new address[](voucheeCount); + + for (uint256 i = 0; i < voucheeCount; i++) { + (address borrower, ) = userManager.vouchees(account, i); + addresses[i] = borrower; + } + + return addresses; + } + + function getStakerAddresses(address underlying, address account) public view returns (address[] memory) { + IUserManager userManager = IUserManager(marketRegistry.userManagers(underlying)); + + uint256 voucherCount = userManager.getVoucherCount(account); + address[] memory addresses = new address[](voucherCount); + + for (uint256 i = 0; i < voucherCount; i++) { + (address staker, , ,) = userManager.vouchers(account, i); + addresses[i] = staker; + } + + return addresses; + } + function getVouchInfo( address underlying, address staker, diff --git a/contracts/mocks/UserManagerMock.sol b/contracts/mocks/UserManagerMock.sol index f2fddc3e..163b1e0f 100644 --- a/contracts/mocks/UserManagerMock.sol +++ b/contracts/mocks/UserManagerMock.sol @@ -1,11 +1,16 @@ //SPDX-License-Identifier: MIT pragma solidity 0.8.16; +import {SafeCastUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/math/SafeCastUpgradeable.sol"; + /** * @title UserManager Contract * @dev Manages the Union members credit lines, and their vouchees and borrowers info. */ contract UserManagerMock { + using SafeCastUpgradeable for uint256; + using SafeCastUpgradeable for uint128; + uint256 public constant MAX_TRUST_LIMIT = 100; uint256 public constant MAX_STAKE_AMOUNT = 1000e18; @@ -19,8 +24,27 @@ contract UserManagerMock { uint256 public totalFrozenAmount; address public stakingToken; + struct Index { + bool isSet; + uint128 idx; + } + struct Vouch { + address staker; + uint96 trust; + uint96 locked; + uint64 lastUpdated; + } + struct Vouchee { + address borrower; + uint96 voucherIndex; + } + mapping(address => uint256) public balances; mapping(address => mapping(address => uint256)) public trust; + mapping(address => mapping(address => Index)) public voucherIndexes; + mapping(address => mapping(address => Index)) public voucheeIndexes; + mapping(address => Vouch[]) public vouchers; + mapping(address => Vouchee[]) public vouchees; constructor(address _stakingToken) { stakingToken = _stakingToken; @@ -75,8 +99,21 @@ contract UserManagerMock { function addMember(address account) public {} - function updateTrust(address borrower_, uint96 trustAmount) external { - trust[msg.sender][borrower_] = trustAmount; + function updateTrust(address borrower, uint96 trustAmount) external { + address staker = msg.sender; + trust[staker][borrower] = trustAmount; + Index memory index = voucherIndexes[borrower][staker]; + if (index.isSet) { + Vouch storage vouch = vouchers[borrower][index.idx]; + vouch.trust = trustAmount; + } else { + uint256 voucheeIndex = vouchees[staker].length; + uint256 voucherIndex = vouchers[borrower].length; + voucherIndexes[borrower][staker] = Index(true, voucherIndex.toUint128()); + voucheeIndexes[borrower][staker] = Index(true, voucheeIndex.toUint128()); + vouchers[borrower].push(Vouch(staker, trustAmount, 0, 0)); + vouchees[staker].push(Vouchee(borrower, voucherIndex.toUint96())); + } } function cancelVouch(address staker, address borrower) external {} @@ -111,7 +148,13 @@ contract UserManagerMock { function onRepayBorrow(address borrower, uint256 pastBlocks) public {} - function getVoucherCount(address borrower) external view returns (uint256) {} + function getVoucherCount(address borrower) external view returns (uint256) { + return vouchers[borrower].length; + } + + function getVoucheeCount(address borrower) external view returns (uint256) { + return vouchees[borrower].length; + } function setEffectiveCount(uint256 effectiveCount) external {} diff --git a/test/foundry/UnionLens.t.sol b/test/foundry/UnionLens.t.sol new file mode 100644 index 00000000..fd376389 --- /dev/null +++ b/test/foundry/UnionLens.t.sol @@ -0,0 +1,38 @@ +pragma solidity ^0.8.0; + +import {TestWrapper} from "./TestWrapper.sol"; +import {UnionLens} from "union-v2-contracts/UnionLens.sol"; + +contract TestUnionLens is TestWrapper { + UnionLens public unionLens; + address public constant ADMIN = address(0); + + function setUp() public virtual { + deployMocks(); + vm.startPrank(ADMIN); + marketRegistryMock.setUToken(address(daiMock), address(uTokenMock)); + marketRegistryMock.setUserManager(address(daiMock), address(userManagerMock)); + vm.stopPrank(); + unionLens = new UnionLens(marketRegistryMock); + } + + function testGetStakerAddresses() public { + address[] memory addresses = unionLens.getStakerAddresses(address(daiMock), address(1)); + assertEq(addresses.length, 0); + + userManagerMock.updateTrust(address(1), 1 ether); + addresses = unionLens.getStakerAddresses(address(daiMock), address(1)); + assertEq(addresses.length, 1); + } + + function testGetBorrowerAddresses() public { + address[] memory addresses = unionLens.getBorrowerAddresses(address(daiMock), address(1)); + assertEq(addresses.length, 0); + + vm.startPrank(address(1)); + userManagerMock.updateTrust(address(2), 1 ether); + addresses = unionLens.getBorrowerAddresses(address(daiMock), address(1)); + assertEq(addresses.length, 1); + vm.stopPrank(); + } +}