From 7c5b8ba24e09a958237b6657a2ba1fb80f84cb63 Mon Sep 17 00:00:00 2001
From: Schlagonia <schlagonia@gmail.com>
Date: Fri, 26 Jan 2024 18:56:57 -0700
Subject: [PATCH] chore: organize reporting

---
 src/TokenizedStrategy.sol    | 71 +++++++++++++++++++-----------------
 src/test/ProfitLocking.t.sol | 17 ++++-----
 2 files changed, 45 insertions(+), 43 deletions(-)

diff --git a/src/TokenizedStrategy.sol b/src/TokenizedStrategy.sol
index 4ce6ba1..33eff91 100644
--- a/src/TokenizedStrategy.sol
+++ b/src/TokenizedStrategy.sol
@@ -1120,24 +1120,32 @@ contract TokenizedStrategy {
             // We have a profit.
             unchecked {
                 profit = newTotalAssets - oldTotalAssets;
-                // Asses performance fees.
-                totalFees = (profit * S.performanceFee) / MAX_BPS;
             }
 
-            address protocolFeesRecipient;
+            // We need to get the equivalent amount of shares
+            // at the current PPS before any minting or burning.
+            sharesToLock = _convertToShares(S, profit);
+
+            // Cache the performance fee.
+            uint16 fee = S.performanceFee;
             uint256 totalFeeShares;
-            uint256 protocolFeeShares;
-            // If performance fees are 0 so will protocol fees.
-            if (totalFees != 0) {
-                // We need to get the shares to issue for the fees at
-                // current PPS before any minting or burning.
-                totalFeeShares = _convertToShares(S, totalFees);
-
-                // Get the config from the factory.
-                uint16 protocolFeeBps;
-                (protocolFeeBps, protocolFeesRecipient) = IFactory(FACTORY)
-                    .protocol_fee_config();
+            // If we are charging a performance fee
+            if (fee != 0) {
+                // Asses performance fees.
+                unchecked {
+                    // Get in `asset` for the event.
+                    totalFees = (profit * fee) / MAX_BPS;
+                    // And in shares for the payment.
+                    totalFeeShares = (sharesToLock * fee) / MAX_BPS;
+                }
+
+                // Get the protocol fee config from the factory.
+                (
+                    uint16 protocolFeeBps,
+                    address protocolFeesRecipient
+                ) = IFactory(FACTORY).protocol_fee_config();
 
+                uint256 protocolFeeShares;
                 // Check if there is a protocol fee to charge.
                 if (protocolFeeBps != 0) {
                     unchecked {
@@ -1148,14 +1156,26 @@ contract TokenizedStrategy {
                         // Need amount in underlying for event.
                         protocolFees = (totalFees * protocolFeeBps) / MAX_BPS;
                     }
+
+                    // Mint the protocol fees to the recipient.
+                    _mint(S, protocolFeesRecipient, protocolFeeShares);
+                }
+
+                // Mint the difference to the strategy fee recipient.
+                unchecked {
+                    _mint(
+                        S,
+                        S.performanceFeeRecipient,
+                        totalFeeShares - protocolFeeShares
+                    );
                 }
             }
 
-            // we have a net profit. Check if we are locking profit.
+            // Check if we are locking profit.
             if (_profitMaxUnlockTime != 0) {
                 // lock (profit - fees)
                 unchecked {
-                    sharesToLock = _convertToShares(S, profit - totalFees);
+                    sharesToLock -= totalFeeShares;
                 }
 
                 // If we are burning more than re-locking.
@@ -1171,35 +1191,20 @@ contract TokenizedStrategy {
                     }
                 }
             }
-
-            // Mint fees shares to recipients.
-            if (totalFeeShares != 0) {
-                unchecked {
-                    _mint(
-                        S,
-                        S.performanceFeeRecipient,
-                        totalFeeShares - protocolFeeShares
-                    );
-                }
-
-                if (protocolFeeShares != 0) {
-                    _mint(S, protocolFeesRecipient, protocolFeeShares);
-                }
-            }
         } else {
             // Expect we have a loss.
             unchecked {
                 loss = oldTotalAssets - newTotalAssets;
             }
 
-            // Check in case else was due to being equal.
+            // Check in case `else` was due to being equal.
             if (loss != 0) {
                 // We will try and burn the unlocked shares and as much from any
                 // pending profit still unlocking to offset the loss to prevent any PPS decline post report.
                 sharesToBurn = Math.min(
                     // Cannot burn more than we have.
                     S.balances[address(this)],
-                    // Try and burn both the shares unlocked and the amount for the loss.
+                    // Try and burn both the shares already unlocked and the amount for the loss.
                     _convertToShares(S, loss) + sharesToBurn
                 );
             }
diff --git a/src/test/ProfitLocking.t.sol b/src/test/ProfitLocking.t.sol
index ce2d2ee..b294e94 100644
--- a/src/test/ProfitLocking.t.sol
+++ b/src/test/ProfitLocking.t.sol
@@ -780,9 +780,9 @@ contract ProfitLockingTest is Setup {
 
         uint256 newAmount = _amount + profit;
 
-        uint256 secondExpectedSharesForFees = strategy.convertToShares(
-            expectedProtocolFee + expectedPerformanceFee
-        );
+        uint256 secondExpectedSharesForFees = (strategy.convertToShares(
+            profit
+        ) * performanceFee) / MAX_BPS;
 
         createAndCheckProfit(
             strategy,
@@ -920,9 +920,9 @@ contract ProfitLockingTest is Setup {
 
         uint256 newAmount = _amount + profit;
 
-        uint256 secondExpectedSharesForFees = strategy.convertToShares(
-            expectedPerformanceFee + expectedProtocolFee
-        );
+        uint256 secondExpectedSharesForFees = (strategy.convertToShares(
+            profit
+        ) * performanceFee) / MAX_BPS;
 
         createAndCheckProfit(
             strategy,
@@ -938,10 +938,7 @@ contract ProfitLockingTest is Setup {
             0,
             newAmount -
                 ((profit - totalExpectedFees) / 2) +
-                strategy.previewWithdraw(
-                    profit - (expectedProtocolFee + expectedPerformanceFee)
-                ) +
-                secondExpectedSharesForFees
+                strategy.convertToShares(profit)
         );
 
         increaseTimeAndCheckBuffer(strategy, profitMaxUnlockTime, 0);