From 73279cf1fd34b441a8433f0a42103bf33a867987 Mon Sep 17 00:00:00 2001
From: eagle <aparokshambhargav@gmail.com>
Date: Thu, 16 Jan 2025 17:20:20 +0530
Subject: [PATCH] feat: update rebalancePool function

---
 src/interfaces/IAssetPool.sol |  6 +++---
 src/protocol/AssetPool.sol    | 32 +++++++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/interfaces/IAssetPool.sol b/src/interfaces/IAssetPool.sol
index 1f20f9a..bad48e5 100644
--- a/src/interfaces/IAssetPool.sol
+++ b/src/interfaces/IAssetPool.sol
@@ -23,13 +23,13 @@ interface IAssetPool {
     event BurnRequested(address indexed user, uint256 xTokenAmount, uint256 indexed cycleIndex);
     event BurnCancelled(address indexed user, uint256 amount, uint256 indexed cycleIndex);
     event ReserveWithdrawn(address indexed user, uint256 amount, uint256 indexed cycleIndex);
-    event Rebalanced(address indexed lp, uint256 amount, bool isDeposit, uint256 indexed cycleIndex);
+    event Rebalanced(address indexed lp, uint256 rebalancePrice, uint256 amount, bool isDeposit, uint256 indexed cycleIndex);
     event CycleStarted(uint256 indexed cycleIndex, uint256 timestamp);
     event CycleTimeUpdated(uint256 newCycleTime);
     event RebalanceTimeUpdated(uint256 newRebalanceTime);
     event RebalanceInitiated(
         uint256 indexed cycleIndex,
-        uint256 spotPrice,
+        uint256 assetPrice,
         int256 netReserveDelta,
         int256 rebalanceAmount
     );
@@ -62,7 +62,7 @@ interface IAssetPool {
     //                                  LP ACTIONS
     // --------------------------------------------------------------------------------
     function initiateRebalance() external;
-    function rebalancePool(address lp, uint256 amount, bool isDeposit) external;
+    function rebalancePool(address lp, uint256 rebalancePrice, uint256 amount, bool isDeposit) external;
 
     // --------------------------------------------------------------------------------
     //                              GOVERNANCE ACTIONS
diff --git a/src/protocol/AssetPool.sol b/src/protocol/AssetPool.sol
index c10d75f..5327cf3 100644
--- a/src/protocol/AssetPool.sol
+++ b/src/protocol/AssetPool.sol
@@ -51,6 +51,7 @@ contract AssetPool is IAssetPool, Ownable, Pausable {
     mapping(address => uint256) public lastActionCycle;     // Last cycle user interacted with
 
     mapping(uint256 => uint256) public cycleRebalancePrice;  // price at which the rebalance was executed in a cycle
+    mapping(uint256 => uint256) private cycleWeightedSum;  // Weighted sum of rebalance prices
 
     // Constants
     uint256 private constant PRECISION = 1e18; // Precision for calculations
@@ -208,12 +209,21 @@ contract AssetPool is IAssetPool, Ownable, Pausable {
      * @notice Once LPs have traded off-chain, they deposit or withdraw stablecoins accordingly.
      * @param lp Address of the LP performing the final on-chain step
      * @param amount Amount of stablecoin they are sending to (or withdrawing from) the pool
+     * @param rebalancePrice Price at which the rebalance was executed
      * @param isDeposit True if depositing stable, false if withdrawing
+     *
+     * ToDo: lpLiquidty should be bassed on how much being asset being rebalanced during the cycle
      */
-    function rebalancePool(address lp, uint256 amount, bool isDeposit) external onlyLP {
+
+    function rebalancePool(address lp, uint256 rebalancePrice, uint256 amount, bool isDeposit) external onlyLP {
         require(cycleState == CycleState.REBALANCING, "Not in rebalancing");
         if (hasRebalanced[lp]) revert AlreadyRebalanced();
         if (block.timestamp > nextRebalanceEndDate) revert RebalancingExpired();
+        uint256 lpLiquidity = lpRegistry.getLPLiquidity(address(this), lp);
+
+        _validateRebalancing(lp, rebalancePrice, amount, isDeposit);
+
+        cycleWeightedSum[cycleIndex] += rebalancePrice * lpLiquidity;
 
         // Approve stablecoins if deposit
         if (isDeposit) {
@@ -221,17 +231,18 @@ contract AssetPool is IAssetPool, Ownable, Pausable {
             reserveToken.transferFrom(lp, address(this), amount);
         } else {
             // If withdrawing stable, ensure the pool has enough
-            if (amount > uint256(rebalanceAmount)) revert InvalidAmount();
             reserveToken.transfer(lp, amount);
         }
 
         hasRebalanced[lp] = true;
         rebalancedLPs++;
 
-        emit Rebalanced(lp, amount, isDeposit, cycleIndex);
+        emit Rebalanced(lp, rebalancePrice, amount, isDeposit, cycleIndex);
 
         // If all LPs have rebalanced, start next cycle
         if (rebalancedLPs == lpRegistry.getLPCount(address(this))) {
+            uint256 totalLiquidity = lpRegistry.getTotalLPLiquidity(address(this));
+            cycleRebalancePrice[cycleIndex] = cycleWeightedSum[cycleIndex] / totalLiquidity;
             _startNewCycle();
         }
     }
@@ -267,8 +278,19 @@ contract AssetPool is IAssetPool, Ownable, Pausable {
     }
 
     // Internal functions
-    function _validateRebalancing() internal view {
-        require(lpRegistry.getLPLiquidity(address(this), msg.sender) > 0, "Insufficient LP liquidity");
+    function _validateRebalancing(address lp, uint256 rebalancePrice, uint256 amount, bool isDeposit) internal view {
+        uint256 lpLiquidity = lpRegistry.getLPLiquidity(address(this), lp);
+        require(lpLiquidity > 0, "Insufficient LP liquidity");
+
+        // Check if the rebalance direction aligns with the rebalanceAmount
+        if (rebalanceAmount > 0 && !isDeposit) revert("Expected deposit, got withdrawal");
+        if (rebalanceAmount < 0 && isDeposit) revert("Expected withdrawal, got deposit");
+
+        // Calculate the expected amount based on LP's liquidity share
+        uint256 expectedAmount = uint256(rebalanceAmount > 0 ? rebalanceAmount : -rebalanceAmount) * lpLiquidity / lpRegistry.getTotalLPLiquidity(address(this));
+        require(amount == expectedAmount, "Amount mismatch with liquidity share");
+
+        
     }
 
     // View functions