create token balancer

contracts/erc4626/TokenBalancer.sol

@@ -0,0 +1,181 @@
+//SPDX-License-Identifier: MIT
+pragma solidity 0.8.24;
+
+import {AccessControl} from "@openzeppelin/contracts/access/AccessControl.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {IERC20Metadata} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
+
+import {IPyth} from "@pythnetwork/pyth-sdk-solidity/IPyth.sol";
+import {PythStructs} from "@pythnetwork/pyth-sdk-solidity/PythStructs.sol";
+
+import {ISaucerSwap} from "./interfaces/ISaucerSwap.sol";
+
+import "../common/safe-HTS/SafeHTS.sol";
+
+/**
+ * @title Token Balancer
+ *
+ * The contract that helps to maintain reward token balances.
+ */
+abstract contract TokenBalancer is AccessControl {
+    mapping(address => uint256 allocationPercentage) internal targetPercentages;
+
+    mapping(address => bytes32 priceId) public priceIds;
+
+    mapping(address => address[]) public swapPaths;
+
+    mapping(address => uint256) public tokenPrices;
+
+    // Saucer Swap
+    ISaucerSwap public saucerSwap;
+
+    // Oracle
+    IPyth public pyth;
+
+    /**
+     * @dev Initializes contract with passed parameters.
+     *
+     * @param _pyth The address of the Pyth oracle contract.
+     * @param _saucerSwap The address of Saucer Swap contract.
+     * @param tokens The reward tokens.
+     * @param allocationPercentage The allocation percentages for rebalances.
+     * @param _priceIds The Pyth price ids to fetch prices.
+     */
+    function __TokenBalancer_init(
+        address _pyth,
+        address _saucerSwap,
+        address[] memory tokens,
+        uint256[] memory allocationPercentage,
+        bytes32[] memory _priceIds
+    ) internal {
+        saucerSwap = ISaucerSwap(_saucerSwap);
+        pyth = IPyth(_pyth);
+
+        address whbar = saucerSwap.WHBAR();
+
+        uint256 tokensSize = tokens.length;
+        for (uint256 i = 0; i < tokensSize; i++) {
+            targetPercentages[tokens[i]] = allocationPercentage[i];
+            priceIds[tokens[i]] = _priceIds[i];
+            swapPaths[tokens[i]] = [whbar, tokens[i]];
+            tokenPrices[tokens[i]] = _getPrice(tokens[i]);
+        }
+    }
+
+    /**
+     * @dev Gets token price and calculate one dollar in any token.
+     *
+     * @param token The token address.
+     */
+    function _getPrice(address token) public view returns (uint256 oneDollarInHbar) {
+        PythStructs.Price memory price = pyth.getPrice(priceIds[token]);
+
+        uint256 decimals = IERC20Metadata(token).decimals();
+
+        uint256 hbarPrice8Decimals = (uint(uint64(price.price)) * (18 ** decimals)) /
+            (18 ** uint8(uint32(-1 * price.expo)));
+        oneDollarInHbar = ((18 ** decimals) * (18 ** decimals)) / hbarPrice8Decimals;
+    }
+
+    /**
+     * @dev Updates price.
+     *
+     * @param pythPriceUpdate The pyth price update.
+     */
+    function update(bytes[] calldata pythPriceUpdate) public payable {
+        uint updateFee = pyth.getUpdateFee(pythPriceUpdate);
+        pyth.updatePriceFeeds{value: updateFee}(pythPriceUpdate);
+    }
+
+    /**
+     * @dev Rebalances reward balances.
+     */
+    function rebalance(address[] calldata _rewardTokens) external {
+        uint256 rewardTokensSize = _rewardTokens.length;
+        uint256[] memory prices;
+        for (uint256 i = 0; i < rewardTokensSize; i++) {
+            prices[i] = tokenPrices[_rewardTokens[i]];
+        }
+
+        uint256[] memory swapAmounts = _rebalance(prices, _rewardTokens);
+
+        _swapExtraRewardSupplyToTransitionToken(_rewardTokens);
+
+        uint256 swapsCount = swapAmounts.length;
+        for (uint256 i = 0; i < swapsCount; i++) {
+            saucerSwap.swapExactETHForTokens(
+                swapAmounts[i],
+                swapPaths[_rewardTokens[i]],
+                address(this),
+                block.timestamp
+            );
+        }
+    }
+
+    /**
+     * @dev Swaps extra reward balance to WHBAR token for future rebalance.
+     *
+     */
+    function _swapExtraRewardSupplyToTransitionToken(address[] calldata _rewardTokens) public {
+        for (uint256 i = 0; i < _rewardTokens.length; i++) {
+            address token = _rewardTokens[i];
+            uint256 tokenBalance = IERC20(token).balanceOf(address(this));
+            uint256 tokenPrice = tokenPrices[token];
+            uint256 totalValue = tokenBalance * tokenPrice;
+            uint256 targetValue = (totalValue * targetPercentages[token]) / 10000;
+            uint256 targetQuantity = targetValue / tokenPrice;
+
+            if (tokenBalance > targetQuantity) {
+                uint256 excessQuantity = tokenBalance - targetQuantity;
+
+                // Approve token transfer to SaucerSwap
+                IERC20(token).approve(address(saucerSwap), excessQuantity);
+
+                // Perform the swap
+                saucerSwap.swapExactTokensForETH(
+                    excessQuantity,
+                    0, // Accept any amount of ETH
+                    swapPaths[token],
+                    address(this),
+                    block.timestamp
+                );
+            }
+        }
+    }
+
+    function _rebalance(
+        uint256[] memory _tokenPrices,
+        address[] calldata _rewardTokens
+    ) public view returns (uint256[] memory) {
+        require(_tokenPrices.length == _rewardTokens.length, "Token prices array length mismatch");
+
+        uint256 totalValue;
+        uint256[] memory tokenBalances = new uint256[](_rewardTokens.length);
+
+        // Calculate total value in the contract
+        for (uint256 i = 0; i < _rewardTokens.length; i++) {
+            tokenBalances[i] = IERC20(_rewardTokens[i]).balanceOf(address(this));
+            totalValue += tokenBalances[i] * _tokenPrices[i];
+        }
+
+        // Array to store the amounts to swap
+        uint256[] memory swapAmounts = new uint256[](_rewardTokens.length);
+
+        // Calculate target values and swap amounts
+        for (uint256 i = 0; i < _rewardTokens.length; i++) {
+            uint256 targetValue = (totalValue * targetPercentages[_rewardTokens[i]]) / 10000;
+            uint256 targetQuantity = targetValue / _tokenPrices[i];
+
+            swapAmounts[i] = targetQuantity - tokenBalances[i];
+        }
+
+        return swapAmounts;
+    }
+
+    // Utility function to update target percentages
+    function setTargetPercentage(address token, uint256 percentage) external {
+        require(percentage < 10000, "Percentage exceeds 100%");
+        require(token != address(0), "Invalid token address");
+        targetPercentages[token] = percentage;
+    }
+}

contracts/erc4626/Vault.sol

@@ -3,13 +3,14 @@ pragma solidity 0.8.24;
 pragma abicoder v2;
 
 import {ERC20} from "./ERC20.sol";
-import {IERC4626} from "./IERC4626.sol";
+import {IERC4626} from "./interfaces/IERC4626.sol";
 import {IHRC} from "../common/hedera/IHRC.sol";
 
 import {FeeConfiguration} from "../common/FeeConfiguration.sol";
+import {TokenBalancer} from "./TokenBalancer.sol";
 
-import {FixedPointMathLib} from "./FixedPointMathLib.sol";
-import {SafeTransferLib} from "./SafeTransferLib.sol";
+import {FixedPointMathLib} from "./libraries/FixedPointMathLib.sol";
+import {SafeTransferLib} from "./libraries/SafeTransferLib.sol";
 
 import {ReentrancyGuard} from "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
 import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
@@ -22,7 +23,7 @@ import "../common/safe-HTS/IHederaTokenService.sol";
  *
  * The contract which represents a custom Vault with Hedera HTS support.
  */
-contract HederaVault is IERC4626, FeeConfiguration, Ownable, ReentrancyGuard {
+contract HederaVault is IERC4626, FeeConfiguration, TokenBalancer, Ownable, ReentrancyGuard {
     using SafeTransferLib for ERC20;
     using FixedPointMathLib for uint256;
     using Bits for uint256;
@@ -92,11 +93,18 @@ contract HederaVault is IERC4626, FeeConfiguration, Ownable, ReentrancyGuard {
         string memory _symbol,
         FeeConfig memory _feeConfig,
         address _vaultRewardController,
-        address _feeConfigController
+        address _feeConfigController,
+        address _pyth,
+        address _saucerSwap,
+        address[] memory _rewardTokens,
+        uint256[] memory allocationPercentage,
+        bytes32[] memory _priceIds
     ) payable ERC20(_name, _symbol, _underlying.decimals()) Ownable(msg.sender) {
         __FeeConfiguration_init(_feeConfig, _vaultRewardController, _feeConfigController);
+        __TokenBalancer_init(_pyth, _saucerSwap, _rewardTokens, allocationPercentage, _priceIds);
 
         asset = _underlying;
+        _rewardTokens = rewardTokens;
 
         _createTokenWithContractAsOwner(_name, _symbol, _underlying);
     }

contracts/erc4626/IERC4626.sol → contracts/erc4626/interfaces/IERC4626.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 pragma solidity 0.8.24;
 
-import {ERC20} from "./ERC20.sol";
+import {ERC20} from "../ERC20.sol";
 
 abstract contract IERC4626 is ERC20 {
     /*///////////////////////////////////////////////////////////////

contracts/erc4626/interfaces/ISaucerSwap.sol

@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.24;
+
+/**
+ * @title Saucer Swap
+ */
+interface ISaucerSwap {
+    function getReserves() external view returns (uint112 reserve0, uint112 reserve1, uint32 blockTimestampLast);
+
+    function getPair(address tokenA, address tokenB) external view returns (address pair);
+
+    function WHBAR() external pure returns (address);
+
+    function swapExactETHForTokens(
+        uint amountOutMin,
+        address[] calldata path,
+        address to,
+        uint deadline
+    ) external payable returns (uint[] memory amounts);
+
+    function swapTokensForExactTokens(
+        uint amountOut,
+        uint amountInMax,
+        address[] calldata path,
+        address to,
+        uint deadline
+    ) external returns (uint[] memory amounts);
+
+    function swapExactTokensForETH(
+        uint amountIn,
+        uint amountOutMin,
+        address[] calldata path,
+        address to,
+        uint deadline
+    ) external returns (uint[] memory amounts);
+}

contracts/erc4626/SafeTransferLib.sol → contracts/erc4626/libraries/SafeTransferLib.sol

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 pragma solidity >=0.8.0;
 
-import {ERC20} from "./ERC20.sol";
+import {ERC20} from "../ERC20.sol";
 
 /// @notice Safe ETH and ERC20 transfer library that gracefully handles missing return values.
 /// @author Solmate (https://github.com/transmissions11/solmate/blob/main/src/utils/SafeTransferLib.sol)

contracts/erc4626/mocks/MockOracle.sol

@@ -0,0 +1,81 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity 0.8.24;
+
+import {IPyth} from "@pythnetwork/pyth-sdk-solidity/IPyth.sol";
+import {PythStructs} from "@pythnetwork/pyth-sdk-solidity/PythStructs.sol";
+
+contract MockPyth is IPyth {
+    mapping(bytes32 => PythStructs.Price) public prices;
+
+    // Function to set mock prices
+    function setPrice(bytes32 id, int64 price, uint64 conf, int32 expo, uint publishTime) public {
+        prices[id] = PythStructs.Price({price: price, conf: conf, expo: expo, publishTime: publishTime});
+    }
+
+    function getValidTimePeriod() external pure override returns (uint validTimePeriod) {
+        return 3600; // 1 hour for example
+    }
+
+    function getPrice(bytes32 id) external view override returns (PythStructs.Price memory price) {
+        return prices[id];
+    }
+
+    function getEmaPrice(bytes32 id) external view override returns (PythStructs.Price memory price) {
+        return prices[id]; // Just returning the same price for simplicity
+    }
+
+    function getPriceUnsafe(bytes32 id) external view override returns (PythStructs.Price memory price) {
+        return prices[id];
+    }
+
+    function getPriceNoOlderThan(bytes32 id, uint age) external view override returns (PythStructs.Price memory price) {
+        require(block.timestamp - prices[id].publishTime <= age, "Price is too old");
+        return prices[id];
+    }
+
+    function getEmaPriceUnsafe(bytes32 id) external view override returns (PythStructs.Price memory price) {
+        return prices[id];
+    }
+
+    function getEmaPriceNoOlderThan(
+        bytes32 id,
+        uint age
+    ) external view override returns (PythStructs.Price memory price) {
+        require(block.timestamp - prices[id].publishTime <= age, "Price is too old");
+        return prices[id];
+    }
+
+    function updatePriceFeeds(bytes[] calldata updateData) external payable override {
+        // Mock implementation
+    }
+
+    function updatePriceFeedsIfNecessary(
+        bytes[] calldata updateData,
+        bytes32[] calldata priceIds,
+        uint64[] calldata publishTimes
+    ) external payable override {
+        // Mock implementation
+    }
+
+    function getUpdateFee(bytes[] calldata updateData) external view override returns (uint feeAmount) {
+        return 0; // No fee for mock
+    }
+
+    function parsePriceFeedUpdates(
+        bytes[] calldata updateData,
+        bytes32[] calldata priceIds,
+        uint64 minPublishTime,
+        uint64 maxPublishTime
+    ) external payable override returns (PythStructs.PriceFeed[] memory priceFeeds) {
+        // Mock implementation
+    }
+
+    function parsePriceFeedUpdatesUnique(
+        bytes[] calldata updateData,
+        bytes32[] calldata priceIds,
+        uint64 minPublishTime,
+        uint64 maxPublishTime
+    ) external payable override returns (PythStructs.PriceFeed[] memory priceFeeds) {
+        // Mock implementation
+    }
+}