add deposit handler and share calculator

src/interfaces/omni/IRateProvider.sol

@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: BSD-3-Clause
+pragma solidity ^0.8.24;
+
+interface IRateProvider {
+    /**
+     * @notice Returns the rate for a given asset
+     * @param asset The address of the asset
+     * @return The rate of the asset
+     */
+    function rate(address asset) external view returns (uint256);
+
+    /**
+     * @notice Converts an amount of an asset to its equivalent in the unit of account
+     * @param asset The address of the asset to convert
+     * @param amount The amount of the asset to convert
+     * @return The equivalent amount in the unit of account
+     */
+    function convert(address asset, uint256 amount) external view returns (uint256);
+}

src/omni/DepositHandler.sol

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.24;
+
+import {IERC20} from "lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
+import {IERC20Metadata} from "lib/openzeppelin-contracts/contracts/token/ERC20/extensions/IERC20Metadata.sol";
+import {Math} from "lib/openzeppelin-contracts/contracts/utils/math/Math.sol";
+import {AccessControlUpgradeable} from "lib/openzeppelin-contracts-upgradeable/contracts/access/AccessControlUpgradeable.sol";
+import {Initializable} from "lib/openzeppelin-contracts-upgradeable/contracts/proxy/utils/Initializable.sol";
+import {OmniVault} from "./OmniVault.sol";
+import {ReentrancyGuardUpgradeable} from "lib/openzeppelin-contracts-upgradeable/contracts/utils/ReentrancyGuardUpgradeable.sol";
+import {IRateProvider} from "../interfaces/omni/IRateProvider.sol";
+
+
+contract DepositHandler is Initializable, AccessControlUpgradeable, ReentrancyGuardUpgradeable {
+
+    // State variables
+    OmniVault public vault;
+    IRateProvider public vaultRateProvider;
+
+    // Events
+    event Deposit(
+        address indexed receiver,
+        address indexed depositAsset,
+        uint256 depositAmount,
+        uint256 shareAmount
+    );
+
+    constructor() {
+        _disableInitializers();
+    }
+
+    function initialize(address _owner, address _vault, address _vaultRateProvider) public initializer {
+        __AccessControl_init();
+        _grantRole(DEFAULT_ADMIN_ROLE, _owner);
+
+        vault = OmniVault(payable(_vault));
+        vaultRateProvider = IRateProvider(_vaultRateProvider);
+    }
+
+    function deposit(IERC20 depositAsset, uint256 depositAmount, address receiver)
+        public
+        payable
+        nonReentrant
+        returns (uint256 shares)
+    {
+        uint256 rateInShares = vaultRateProvider.convert(address(depositAsset), depositAmount);
+
+        vault.enter(msg.sender, depositAsset, depositAmount, receiver, shares);
+    }
+}

src/omni/OmniVault.sol

@@ -0,0 +1,122 @@
+// SPDX-License-Identifier: BSD-3-Clause
+pragma solidity ^0.8.24;
+
+
+import {Address} from "@openzeppelin/contracts/utils/Address.sol";
+import {ERC20Upgradeable} from
+    "lib/openzeppelin-contracts-upgradeable/contracts/token/ERC20/ERC20Upgradeable.sol";
+import {console} from "lib/forge-std/src/console.sol";
+import {IERC20} from "lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
+import {AccessControlUpgradeable} from "lib/openzeppelin-contracts-upgradeable/contracts/access/AccessControlUpgradeable.sol";
+import {Initializable} from "lib/openzeppelin-contracts-upgradeable/contracts/proxy/utils/Initializable.sol";
+import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+
+
+contract OmniVault is ERC20Upgradeable, AccessControlUpgradeable {
+    using Address for address;
+    using SafeERC20 for IERC20;
+
+
+    // ========================================= CONSTANTS =========================================
+
+    bytes32 public constant MANAGER_ROLE = keccak256("MANAGER_ROLE");
+    bytes32 public constant ENTER_ROLE = keccak256("ENTER_ROLE");
+    bytes32 public constant EXIT_ROLE = keccak256("EXIT_ROLE");
+
+    //============================== EVENTS ===============================
+
+    event Enter(address indexed from, address indexed asset, uint256 amount, address indexed to, uint256 shares);
+    event Exit(address indexed to, address indexed asset, uint256 amount, address indexed from, uint256 shares);
+
+    //============================== INITIALIZER ===============================
+
+    constructor() {
+        _disableInitializers();
+    }
+
+    function initialize(address _owner, string memory _name, string memory _symbol) 
+        public 
+        initializer 
+    {
+        __ERC20_init(_name, _symbol);
+        __AccessControl_init();
+        _grantRole(DEFAULT_ADMIN_ROLE, _owner);
+        _grantRole(MANAGER_ROLE, _owner);
+        _grantRole(ENTER_ROLE, _owner);
+        _grantRole(EXIT_ROLE, _owner);
+    }
+
+    //============================== MANAGE ===============================
+
+    /**
+     * @notice Allows manager to make an arbitrary function call from this contract.
+     * @dev Callable by MANAGER_ROLE.
+     */
+    function manage(address target, bytes calldata data, uint256 value)
+        external
+        onlyRole(MANAGER_ROLE)
+        returns (bytes memory result)
+    {
+        result = target.functionCallWithValue(data, value);
+    }
+
+    /**
+     * @notice Allows manager to make arbitrary function calls from this contract.
+     * @dev Callable by MANAGER_ROLE.
+     */
+    function manage(address[] calldata targets, bytes[] calldata data, uint256[] calldata values)
+        external
+        onlyRole(MANAGER_ROLE)
+        returns (bytes[] memory results)
+    {
+        uint256 targetsLength = targets.length;
+        results = new bytes[](targetsLength);
+        for (uint256 i; i < targetsLength; ++i) {
+            results[i] = targets[i].functionCallWithValue(data[i], values[i]);
+        }
+    }
+
+    //============================== ENTER ===============================
+
+    /**
+     * @notice Allows minter to mint shares, in exchange for assets.
+     * @dev If assetAmount is zero, no assets are transferred in.
+     * @dev Callable by ENTER_ROLE.
+     */
+    function enter(address from, IERC20 asset, uint256 assetAmount, address to, uint256 shareAmount)
+        external
+        onlyRole(ENTER_ROLE)
+    {
+        // Transfer assets in
+        if (assetAmount > 0) asset.safeTransferFrom(from, address(this), assetAmount);
+
+        // Mint shares.
+        _mint(to, shareAmount);
+
+        emit Enter(from, address(asset), assetAmount, to, shareAmount);
+    }
+
+    //============================== EXIT ===============================
+
+    /**
+     * @notice Allows burner to burn shares, in exchange for assets.
+     * @dev If assetAmount is zero, no assets are transferred out.
+     * @dev Callable by EXIT_ROLE.
+     */
+    function exit(address to, IERC20 asset, uint256 assetAmount, address from, uint256 shareAmount)
+        external
+        onlyRole(EXIT_ROLE)
+    {
+        // Burn shares.
+        _burn(from, shareAmount);
+
+        // Transfer assets out.
+        if (assetAmount > 0) asset.safeTransfer(to, assetAmount);
+
+        emit Exit(to, address(asset), assetAmount, from, shareAmount);
+    }
+
+    //============================== RECEIVE ===============================
+
+    receive() external payable {}
+}

src/omni/VaultRateProvider.sol

@@ -0,0 +1,68 @@
+// SPDX-License-Identifier: BSD-3-Clause
+pragma solidity ^0.8.24;
+
+import {AccessControlUpgradeable} from "lib/openzeppelin-contracts-upgradeable/contracts/access/AccessControlUpgradeable.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import {OmniVault} from "./OmniVault.sol";
+import {IRateProvider} from "../interfaces/omni/IRateProvider.sol";
+import {Math} from "@openzeppelin/contracts/utils/math/Math.sol";
+import {IERC20Metadata} from "lib/openzeppelin-contracts/contracts/token/ERC20/extensions/IERC20Metadata.sol";
+
+
+interface ITvlModule {
+    function getTvl() external view returns (uint256);
+}
+
+
+contract VaultRateProvider is AccessControlUpgradeable {
+    using SafeERC20 for IERC20;
+
+    bytes32 public constant MANAGER_ROLE = keccak256("MANAGER_ROLE");
+
+    OmniVault public vault;
+    ITvlModule[] public tvlModules;
+    IRateProvider public assetRateProvider;
+
+    /**
+     * @dev One unit of Unit of Account Asset.
+     */
+    uint256 unit;
+    uint256 public vaultRate;
+
+
+    function initialize(
+        address _owner,
+        address _vault,
+        address _assetRateProvider,
+        uint256 _unit
+    ) public initializer {
+        __AccessControl_init();
+        _grantRole(DEFAULT_ADMIN_ROLE, _owner);
+        _grantRole(MANAGER_ROLE, _owner);
+
+        vault = OmniVault(payable(_vault));
+        assetRateProvider = IRateProvider(_assetRateProvider);
+        unit = _unit;
+    }
+
+    function rate(address asset) public view returns (uint256) {
+
+        uint256 unitOfAccountRateForAsset = assetRateProvider.rate(asset);
+
+        uint256 calculatedRate = Math.mulDiv(unitOfAccountRateForAsset, unit, vaultRate, Math.Rounding.Floor);
+
+        return calculatedRate;
+    }
+
+    function convert(address asset, uint256 amount) public view returns (uint256) {
+
+        // TODO: Check and verify the math in this function to ensure accuracy
+
+        uint256 unitOfAccountRateForAsset = assetRateProvider.rate(asset);
+
+        uint256 amountInUnitOfAccount = Math.mulDiv(amount, unitOfAccountRateForAsset, 10 ** IERC20Metadata(asset).decimals(), Math.Rounding.Floor);
+
+        return Math.mulDiv(amountInUnitOfAccount, unit, vaultRate, Math.Rounding.Floor);
+    }
+}

test/omni/OmniVault.t.sol

@@ -2,7 +2,7 @@
 pragma solidity ^0.8.24;
 
 import {Test} from "forge-std/Test.sol";
-import {OmniVault} from "src/OmniVault.sol";
+import {OmniVault} from "src/omni/OmniVault.sol";
 import {ERC20} from "lib/openzeppelin-contracts/contracts/token/ERC20/ERC20.sol";
 import {console} from "forge-std/console.sol";
 import {TimelockControllerUpgradeable} from "lib/openzeppelin-contracts-upgradeable/contracts/governance/TimelockControllerUpgradeable.sol";