Use project composition to calculate fees

src/FeeCalculator.sol

@@ -11,6 +11,7 @@ import {SD59x18, sd, intoUint256} from "@prb/math/src/SD59x18.sol";
 
 import {IFeeCalculator, FeeDistribution} from "./interfaces/IFeeCalculator.sol";
 import "./interfaces/IPool.sol";
+import {VintageData, ITCO2} from "./interfaces/ITCO2.sol";
 
 /// @title FeeCalculator
 /// @author Neutral Labs Inc.
@@ -146,7 +147,7 @@ contract FeeCalculator is IFeeCalculator, Ownable {
     {
         require(depositAmount > 0, "depositAmount must be > 0");
 
-        uint256 feeAmount = getDepositFee(depositAmount, getTokenBalance(pool, tco2), getTotalSupply(pool));
+        uint256 feeAmount = getDepositFee(depositAmount, getProjectSupply(pool, tco2), getTotalSupply(pool));
 
         require(feeAmount <= depositAmount, "Fee must be lower or equal to deposit amount");
         require(feeAmount > 0, "Fee must be greater than 0");
@@ -195,22 +196,13 @@ contract FeeCalculator is IFeeCalculator, Ownable {
 
         require(redemptionAmount > 0, "redemptionAmount must be > 0");
 
-        uint256 feeAmount = getRedemptionFee(redemptionAmount, getTokenBalance(pool, tco2), getTotalSupply(pool));
+        uint256 feeAmount = getRedemptionFee(redemptionAmount, getProjectSupply(pool, tco2), getTotalSupply(pool));
 
         require(feeAmount <= redemptionAmount, "Fee must be lower or equal to redemption amount");
         require(feeAmount > 0, "Fee must be greater than 0");
         feeDistribution = calculateFeeShares(feeAmount);
     }
 
-    /// @notice Gets the balance of the TCO2 token in a given pool.
-    /// @param pool The address of the pool.
-    /// @param tco2 The address of the TCO2 token.
-    /// @return The balance of the TCO2 token in the pool.
-    function getTokenBalance(address pool, address tco2) private view returns (uint256) {
-        uint256 tokenBalance = IERC20(tco2).balanceOf(pool);
-        return tokenBalance;
-    }
-
     /// @notice Gets the total supply of a given pool.
     /// @param pool The address of the pool.
     /// @return The total supply of the pool.
@@ -219,6 +211,15 @@ contract FeeCalculator is IFeeCalculator, Ownable {
         return totalSupply;
     }
 
+    /// @notice Gets the total supply of a project in the pool.
+    /// @param pool The address of the pool.
+    /// @return The total supply of the pool.
+    function getProjectSupply(address pool, address tco2) private view returns (uint256) {
+        VintageData memory vData = ITCO2(tco2).getVintageData();
+        uint256 projectSupply = IPool(pool).totalPerProjectTCO2Supply(vData.projectTokenId);
+        return projectSupply;
+    }
+
     /// @notice Calculates the ratios for deposit fee calculation.
     /// @param amount The amount to be deposited.
     /// @param current The current balance of the pool.

src/interfaces/IPool.sol

@@ -10,7 +10,14 @@ pragma solidity ^0.8.13;
 /// @notice This interface defines methods exposed by the Pool
 interface IPool {
     /// @notice Exposes the total TCO2 supply, tracked as the aggregation of deposit,
-    /// redemmption and bridge actions
+    /// redemption and bridge actions
     /// @return supply Current supply
     function totalTCO2Supply() external view returns (uint256 supply);
+
+    /// @notice Exposes the total TCO2 supply of a project in a pool,
+    /// tracked as the aggregation of deposit, redemmption and bridge actions
+    /// @param projectTokenId The token id of the project as it's tracked
+    /// in the CarbonProjects contract
+    /// @return supply Current supply of a project in the pool
+    function totalPerProjectTCO2Supply(uint256 projectTokenId) external view returns (uint256 supply);
 }

src/interfaces/ITCO2.sol

@@ -0,0 +1,31 @@
+// SPDX-FileCopyrightText: 2024 Neutral Labs Inc.
+//
+// SPDX-License-Identifier: UNLICENSED
+
+// If you encounter a vulnerability or an issue, please contact <info@neutralx.com>
+pragma solidity ^0.8.13;
+
+struct VintageData {
+    /// @dev A human-readable string which differentiates this from other vintages in
+    /// the same project, and helps build the corresponding TCO2 name and symbol.
+    string name;
+    uint64 startTime; // UNIX timestamp
+    uint64 endTime; // UNIX timestamp
+    uint256 projectTokenId;
+    uint64 totalVintageQuantity;
+    bool isCorsiaCompliant;
+    bool isCCPcompliant;
+    string coBenefits;
+    string correspAdjustment;
+    string additionalCertification;
+    string uri;
+    string registry;
+}
+
+/// @title ITCO2
+/// @notice This interface defines methods exposed by the TCO2
+interface ITCO2 {
+    /// @notice Get the vintage data for the TCO2
+    /// @return vintageData Vintage data of the TCO2
+    function getVintageData() external view returns (VintageData memory vintageData);
+}

test/FeeCalculator.fuzzy.t.sol

@@ -45,7 +45,7 @@ contract FeeCalculatorTestFuzzy is Test {
 
         // Set up mock pool
         mockPool.setTotalSupply(1e12 * 1e18);
-        mockToken.setTokenBalance(address(mockPool), 1e9 * 1e18);
+        mockPool.setProjectSupply(1, 1e9 * 1e18);
 
         vm.expectRevert("Fee must be greater than 0");
         feeCalculator.calculateDepositFees(address(mockPool), address(mockToken), depositAmount);
@@ -65,7 +65,7 @@ contract FeeCalculatorTestFuzzy is Test {
 
         // Set up mock pool
         mockPool.setTotalSupply(total);
-        mockToken.setTokenBalance(address(mockPool), current);
+        mockPool.setProjectSupply(1, current);
 
         // Act
         try feeCalculator.calculateDepositFees(address(mockPool), address(mockToken), depositAmount) {}
@@ -114,7 +114,7 @@ contract FeeCalculatorTestFuzzy is Test {
 
         // Set up mock pool
         mockPool.setTotalSupply(total);
-        mockToken.setTokenBalance(address(mockPool), current);
+        mockPool.setProjectSupply(1, current);
         uint256 oneTimeFee = 0;
         bool oneTimeRedemptionFailed = false;
         uint256 multipleTimesRedemptionFailedCount = 0;
@@ -158,7 +158,7 @@ contract FeeCalculatorTestFuzzy is Test {
                 total -= redemption;
                 current -= redemption;
                 mockPool.setTotalSupply(total);
-                mockToken.setTokenBalance(address(mockPool), current);
+                mockPool.setProjectSupply(1, current);
             } catch Error(string memory reason) {
                 multipleTimesRedemptionFailedCount++;
                 assertTrue(
@@ -204,7 +204,7 @@ contract FeeCalculatorTestFuzzy is Test {
         uint256 multipleTimesDepositFailedCount = 0;
         // Set up mock pool
         mockPool.setTotalSupply(total);
-        mockToken.setTokenBalance(address(mockPool), current);
+        mockPool.setProjectSupply(1, current);
 
         uint256 oneTimeFee = 0;
 
@@ -236,7 +236,7 @@ contract FeeCalculatorTestFuzzy is Test {
                 total += deposit;
                 current += deposit;
                 mockPool.setTotalSupply(total);
-                mockToken.setTokenBalance(address(mockPool), current);
+                mockPool.setProjectSupply(1, current);
             } catch Error(string memory reason) {
                 multipleTimesDepositFailedCount++;
                 assertTrue(
@@ -280,7 +280,7 @@ contract FeeCalculatorTestFuzzy is Test {
         uint256 depositAmount = 100 * 1e18;
         // Set up mock pool
         mockPool.setTotalSupply(200 * 1e18);
-        mockToken.setTokenBalance(address(mockPool), 100 * 1e18);
+        mockPool.setProjectSupply(1, 100 * 1e18);
 
         // Act
         FeeDistribution memory feeDistribution =