ghost-dao-contracts/src/TreasuryExtender.sol

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import {IERC20} from "@openzeppelin-contracts/token/ERC20/IERC20.sol";
import {SafeERC20} from "@openzeppelin-contracts/token/ERC20/utils/SafeERC20.sol";

import {GhostAccessControlled} from "./types/GhostAccessControlled.sol";

import {ITreasury} from "./interfaces/ITreasury.sol";
import {IAllocator} from "./interfaces/IAllocator.sol";
import {ITreasuryExtender} from "./interfaces/ITreasuryExtender.sol";
import {IGhostAuthority} from "./interfaces/IGhostAuthority.sol";

contract TreasuryExtender is GhostAccessControlled, ITreasuryExtender {
    using SafeERC20 for IERC20;

    ITreasury public immutable treasury; // forge-lint: disable-line(screaming-snake-case-immutable)
    IAllocator[] public allocators;
    mapping(IAllocator => mapping(uint256 => AllocatorData)) public allocatorData;

    constructor(address treasuryAddress, address authorityAddress)
        GhostAccessControlled(IGhostAuthority(authorityAddress))
    {
        treasury = ITreasury(treasuryAddress);
        allocators.push(IAllocator(address(0)));
    }

    function _allocatorActivated(IAllocator.AllocatorStatus status) internal pure {
        if (IAllocator.AllocatorStatus.ACTIVATED != status) revert AllocatorNotActivated();
    }

    function _allocatorOffline(IAllocator.AllocatorStatus status) internal pure {
        if (IAllocator.AllocatorStatus.OFFLINE != status) revert AllocatorNotOffline();
    }

    function _onlyAllocator(
        IAllocator byStateId,
        address sender,
        uint256 id
    ) internal pure {
        if (IAllocator(sender) != byStateId) revert OnlyAllocator(id, sender);
    }

    function registerDeposit(address newAllocator) external override onlyGuardian {
        IAllocator allocator = IAllocator(newAllocator);
        uint256 id = allocators.length;
        allocators.push(allocator);
        allocator.addId(id);
        emit NewDepositRegistered(
            newAllocator,
            address(allocator.tokens()[allocator.tokenIds(id)]),
            id
        );
    }

    function setAllocatorLimits(
        uint256 id,
        AllocatorLimits calldata limits
    ) external override onlyGuardian {
        IAllocator allocator = allocators[id];
        _allocatorOffline(allocator.status());
        allocatorData[allocator][id].limits = limits;
        emit AllocatorLimitsChanged(id, limits.allocated, limits.loss);
    }

    function report(
        uint256 id,
        uint128 gain,
        uint128 loss
    ) external override {
        IAllocator allocator = allocators[id];
        AllocatorData storage data = allocatorData[allocator][id];
        AllocatorPerformance memory perf = data.performance;
        IAllocator.AllocatorStatus status = allocator.status();

        _onlyAllocator(allocator, msg.sender, id);
        if (status == IAllocator.AllocatorStatus.OFFLINE) revert AllocatorOffline();

        if (gain >= loss) {
            if (loss == type(uint128).max) {
                AllocatorData storage newAllocatorData =
                    allocatorData[allocators[allocators.length - 1]][id];

                newAllocatorData.holdings.allocated = data.holdings.allocated;
                newAllocatorData.performance.gain = data.performance.gain;
                data.holdings.allocated = 0;

                perf.gain = 0;
                perf.loss = 0;

                emit AllocatorReportedMigration(id);
            } else {
                perf.gain += gain;
                emit AllocatorReportedGain(id, gain);
            }
        } else {
            data.holdings.allocated -= loss;
            perf.loss += loss;
            emit AllocatorReportedLoss(id, loss);
        }
        data.performance = perf;
    }

    function requestFundsFromTreasury(
        uint256 id,
        uint256 amount
    ) external override onlyGuardian {
        IAllocator allocator = allocators[id];
        AllocatorData memory data = allocatorData[allocator][id];
        address token = address(allocator.tokens()[allocator.tokenIds(id)]);
        uint256 value = treasury.tokenValue(token, amount);

        _allocatorActivated(allocator.status());
        _allocatorBelowLimit(data, amount);

        treasury.manage(token, amount);
        allocatorData[allocator][id].holdings.allocated += amount;

        IERC20(token).safeTransfer(address(allocator), amount);
        emit AllocatorFunded(id, amount, value);
    }

    function returnFundsToTreasury(
        uint256 id,
        uint256 amount
    ) external override onlyGuardian {
        IAllocator allocator = allocators[id];
        uint256 allocated = allocatorData[allocator][id].holdings.allocated;
        uint128 gain = allocatorData[allocator][id].performance.gain;
        address token = address(allocator.tokens()[allocator.tokenIds(id)]);

        if (amount > allocated) {
            amount -= allocated;
            if (amount > gain) {
                amount = allocated + gain;
                gain = 0;
            } else {
                // forge-lint: disable-next-line(unsafe-typecast)
                gain -= uint128(amount);
                amount += allocated;
            }
            allocated = 0;
        } else {
            allocated -= amount;
        }

        uint256 value = treasury.tokenValue(token, amount);
        _allowTreasuryWithdrawal(IERC20(token));
        IERC20(token).safeTransferFrom(address(allocator), address(this), amount);

        allocatorData[allocator][id].holdings.allocated = allocated;
        if (allocated == 0) allocatorData[allocator][id].performance.gain = gain;

        assert(treasury.deposit(token, amount, value) == 0);
        emit AllocatedWithdrawal(id, amount, value);
    }

    function returnRewardsToTreasury(
        uint256 id,
        address token,
        uint256 amount
    ) external override {
        _returnRewardsToTreasury(allocators[id], IERC20(token), amount);
    }

    function returnRewardsToTreasury(
        address allocatorAddress,
        address token,
        uint256 amount
    ) external override {
        _returnRewardsToTreasury(IAllocator(allocatorAddress), IERC20(token), amount);
    }

    function getAllocatorById(uint256 id)
        external
        view
        override
        returns (address)
    {
        return address(allocators[id]);
    }

    function getTotalAllocatorCount()
        external
        view
        returns (uint256)
    {
        return allocators.length;
    }

    function getAllocatorLimits(uint256 id)
        external
        view
        override
        returns (AllocatorLimits memory)
    {
        return allocatorData[allocators[id]][id].limits;
    }

    function getAllocatorPerformance(uint256 id)
        external
        view
        override
        returns (AllocatorPerformance memory)
    {
        return allocatorData[allocators[id]][id].performance;
    }

    function getAllocatorAllocated(uint256 id)
        external
        view
        override
        returns (uint256)
    {
        return allocatorData[allocators[id]][id].holdings.allocated;
    }

    function _returnRewardsToTreasury(
        IAllocator allocator,
        IERC20 token,
        uint256 amount
    ) internal onlyGuardian {
        uint256 balance = token.balanceOf(address(allocator));
        amount = balance < amount ? balance : amount;
        uint256 value = treasury.tokenValue(address(token), amount);

        _allowTreasuryWithdrawal(token);

        token.safeTransferFrom(address(allocator), address(this), amount);
        assert(treasury.deposit(address(token), amount, value) == 0);
        emit AllocatorRewardsWithdrawal(address(allocator), amount, value);
    }

    function _allowTreasuryWithdrawal(IERC20 token) internal {
        if (token.allowance(address(this), address(treasury)) == 0) {
            token.approve(address(treasury), type(uint256).max);
        }
    }

    function _allocatorBelowLimit(
        AllocatorData memory data,
        uint256 amount
    ) internal pure {
        uint256 newAllocated = data.holdings.allocated + amount;
        if (newAllocated > data.limits.allocated) {
            revert AllocatorMaxAllocation(newAllocated, data.limits.allocated);
        }
    }
}