HAMT-72: Refactor Community Vault - PurchaseTracker Relationship

scripts/DeployHamzaVault.s.sol

@@ -7,6 +7,7 @@ import "@baal/Baal.sol";
 import "@baal/BaalSummoner.sol";
 
 import "../src/CommunityVault.sol";
+import "../src/CommunityRewardsCalculator.sol";
 import "../src/tokens/GovernanceToken.sol";
 import "../src/GovernanceVault.sol";
 
@@ -211,6 +212,8 @@ contract DeployHamzaVault is Script {
         // 5) Deploy the Community Vault
         CommunityVault communityVault = new CommunityVault(hatsSecurityContextAddr);
         vault = payable(address(communityVault));
+        CommunityRewardsCalculator calculator = new CommunityRewardsCalculator();
+        communityVault.setCommunityRewardsCalculator(calculator);
 
         // 6) Summon the Baal DAO
         BaalSummoner summoner = BaalSummoner(BAAL_SUMMONER);
@@ -412,10 +415,10 @@ contract DeployHamzaVault is Script {
         bool autoRelease = config.readBool(".escrow.autoRelease");
 
         // 15) Deploy PurchaseTracker
-        PurchaseTracker purchaseTracker = new PurchaseTracker(securityContext, vault, lootTokenAddr);
+        PurchaseTracker purchaseTracker = new PurchaseTracker(securityContext, lootTokenAddr);
 
         //setPurchaseTracker in community vault
-        CommunityVault(vault).setPurchaseTracker(address(purchaseTracker), lootTokenAddr);
+        CommunityVault(vault).setPurchaseTracker(address(purchaseTracker));
 
         purchaseTrackerAddr = address(purchaseTracker);
 
@@ -438,15 +441,15 @@ contract DeployHamzaVault is Script {
 
     function logDeployedAddresses(
         address newBaalAddr,
-        address vault,
+        address communityVault,
         address govTokenAddr,
         address govVaultAddr,
         address timelockAddr
     ) internal view {
         console2.log("Owner One (from PRIVATE_KEY):", OWNER_ONE);
         console2.log("Owner Two (from config):     ", OWNER_TWO);
 
-        console2.log("CommunityVault deployed at:", vault);
+        console2.log("CommunityVault deployed at:", communityVault);
 
         console2.log("BaalSummoner at:", BAAL_SUMMONER);
         console2.log("Baal (Hamza Vault) deployed at:", newBaalAddr);

src/CommunityRewardsCalculator.sol

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import "@hamza-escrow/IPurchaseTracker.sol";
+import "./ICommunityRewardsCalculator.sol";
+
+/**
+ * @title CommunityRewardsCalculator
+ * @dev Contains the logic for calculating who gets what rewards, and for what reasons. The rewards are 
+ * distributed through the CommunityVault. 
+ */
+contract CommunityRewardsCalculator is ICommunityRewardsCalculator {
+
+    function getRewardsToDistribute(
+        address token, 
+        address[] calldata recipients,
+        IPurchaseTracker purchaseTracker
+    ) external returns (uint256[] memory) {
+        uint256[] memory amounts = new uint256[](recipients.length);
+
+        // for every purchase or sale made by the recipient, distribute 1 loot token
+        for (uint i=0; i<recipients.length; i++) {
+            uint256 totalPurchase = purchaseTracker.getPurchaseCount(recipients[i]);
+            uint256 totalSales = purchaseTracker.getSalesAmount(recipients[i]);
+            uint256 totalRewards = totalPurchase + totalSales;
+            amounts[i] = totalRewards;
+        }
+
+        return amounts;
+    }
+}

src/CommunityVault.sol

@@ -5,7 +5,7 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@hamza-escrow/security/HasSecurityContext.sol";
 import "@hamza-escrow/security/Roles.sol";
 import "@hamza-escrow/security/ISecurityContext.sol";
-import "./GovernanceVault.sol";      
+import "./ICommunityRewardsCalculator.sol";
 
 /**
  * @title CommunityVault
@@ -17,12 +17,18 @@ contract CommunityVault is HasSecurityContext {
     // Mapping to store token balances held in the community vault
     mapping(address => uint256) public tokenBalances;
 
+    // Keeps a count of already distributed rewards
+    mapping(address => mapping(address => uint256)) public rewardsDistributed;
+
     // Governance staking contract address
     address public governanceVault;
 
     // Address for purchase tracker 
     address public purchaseTracker;
 
+    // Address for rewards calculator 
+    ICommunityRewardsCalculator public rewardsCalculator;
+
     // Events
     event Deposit(address indexed token, address indexed from, uint256 amount);
     event Withdraw(address indexed token, address indexed to, uint256 amount);
@@ -91,35 +97,36 @@ contract CommunityVault is HasSecurityContext {
         address[] calldata recipients,
         uint256[] calldata amounts
     ) external onlyRole(Roles.SYSTEM_ROLE) {
-        require(recipients.length == amounts.length, "Mismatched arrays");
-
-        for (uint256 i = 0; i < recipients.length; i++) {
-            require(tokenBalances[token] >= amounts[i], "Insufficient balance");
-
-            if (token == address(0)) {
-                // ETH distribution
-                (bool success, ) = recipients[i].call{value: amounts[i]}("");
-                require(success, "ETH transfer failed");
-            } else {
-                // ERC20 distribution
-                IERC20(token).safeTransfer(recipients[i], amounts[i]);
-            }
+        _distribute(token, recipients, amounts);
+    }
 
-            tokenBalances[token] -= amounts[i];
+    /**
+     * @dev Distribute tokens or ETH from the community vault to multiple recipients, using the 
+     *      CommunityRewardsCalculator to calculate the amounts to reward each recipient.
+     * @param token The address of the token 
+     * @param recipients The array of recipient addresses
+     */
+    function distributeRewards(address token, address[] memory recipients) external onlyRole(Roles.SYSTEM_ROLE) {
+        _distributeRewards(token, recipients);
+    }
 
-            emit Distribute(token, recipients[i], amounts[i]);
-        }
+    /**
+     * @dev Allows a rightful recipient of rewards to claim rewards that have been allocated to them.
+     * @param token The address of the token 
+     */
+    function claimRewards(address token) external {
+        address[] memory recipients = new address[](1);
+        recipients[0] = msg.sender;
+        _distributeRewards(token, recipients);
     }
 
     /**
-    * @dev Set the governance vault address and grant it unlimited allowance for `lootToken`.
-    *      Must be called by an admin role or similar.
-    * @param vault The address of the governance vault
-    * @param lootToken The address of the ERC20 token for which you'd like to grant unlimited allowance
-    */
-    function setGovernanceVault(address vault, address lootToken)
-        external
-    {
+     * @dev Set the governance vault address and grant it unlimited allowance for `lootToken`.
+     *      Must be called by an admin role or similar.
+     * @param vault The address of the governance vault
+     * @param lootToken The address of the ERC20 token for which you'd like to grant unlimited allowance
+     */
+    function setGovernanceVault(address vault, address lootToken) external onlyRole(Roles.SYSTEM_ROLE)  {
         require(vault != address(0), "Invalid staking contract address");
         require(lootToken != address(0), "Invalid loot token address");
 
@@ -130,15 +137,20 @@ contract CommunityVault is HasSecurityContext {
         IERC20(lootToken).safeApprove(vault, type(uint256).max);
     }
 
-    function setPurchaseTracker(address _purchaseTracker, address lootToken) external {
+    /**
+     * @dev Sets the purchase tracker that is used to keep track of who has done what, in order to get rewards. 
+     */
+    function setPurchaseTracker(address _purchaseTracker) external onlyRole(Roles.SYSTEM_ROLE) {
         require(_purchaseTracker != address(0), "Invalid purchase tracker address");
-        require(lootToken != address(0), "Invalid loot token address");
 
         purchaseTracker = _purchaseTracker;
+    }
 
-        // Grant unlimited allowance to the purchase tracker
-        IERC20(lootToken).safeApprove(_purchaseTracker, 0);
-        IERC20(lootToken).safeApprove(_purchaseTracker, type(uint256).max);
+    /**
+     * @dev Sets the address of the contract which holds the logic for calculating how to divide up rewards. 
+     */
+    function setCommunityRewardsCalculator(ICommunityRewardsCalculator calculator) external onlyRole(Roles.SYSTEM_ROLE) {
+        rewardsCalculator = calculator;
     }
 
     /**
@@ -149,6 +161,47 @@ contract CommunityVault is HasSecurityContext {
         return tokenBalances[token];
     }
 
+    function _distributeRewards(address token, address[] memory recipients) internal {
+        if (address(rewardsCalculator) != address(0) && address(purchaseTracker) != address(0)) {
+
+            //get rewards to distribute
+            uint256[] memory amounts = rewardsCalculator.getRewardsToDistribute(
+                token, recipients, IPurchaseTracker(purchaseTracker)
+            );
+
+            _distribute(token, recipients, amounts);
+        }
+    }
+
+    function _distribute(
+        address token,
+        address[] memory recipients,
+        uint256[] memory amounts
+    ) internal {
+        require(recipients.length == amounts.length, "Mismatched arrays");
+
+        for (uint256 i = 0; i < recipients.length; i++) {
+            require(tokenBalances[token] >= amounts[i], "Insufficient balance");
+
+            if (token == address(0)) {
+                // ETH distribution
+                (bool success, ) = recipients[i].call{value: amounts[i]}("");
+                require(success, "ETH transfer failed");
+            } else {
+                // ERC20 distribution
+                IERC20(token).safeTransfer(recipients[i], amounts[i]);
+            }
+
+            // decrement balance 
+            tokenBalances[token] -= amounts[i];
+
+            // record the distribution
+            rewardsDistributed[token][recipients[i]] += amounts[i];
+
+            emit Distribute(token, recipients[i], amounts[i]);
+        }
+    }
+
     // Fallback to receive ETH
     receive() external payable {}
 }

src/ICommunityRewardsCalculator.sol

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import "@hamza-escrow/IPurchaseTracker.sol";
+
+/**
+ * @title ICommunityRewardsCalculator
+ * @dev Defines the logic for calculating who gets what rewards, and for what reasons. The rewards are 
+ * distributed through the CommunityVault. 
+ */
+interface ICommunityRewardsCalculator {
+    function getRewardsToDistribute(
+        address token, 
+        address[] calldata recipients,
+        IPurchaseTracker purchaseTracker
+    ) external returns (uint256[] memory);
+}

src/PurchaseTracker.sol

@@ -4,13 +4,14 @@ pragma solidity ^0.8.20;
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@hamza-escrow/security/HasSecurityContext.sol";
+import "@hamza-escrow/IPurchaseTracker.sol";
 
 /**
  * @title PurchaseTracker
  * @notice A singleton contract that records purchase data.
  *
  */
-contract PurchaseTracker is HasSecurityContext {
+contract PurchaseTracker is HasSecurityContext, IPurchaseTracker {
     using SafeERC20 for IERC20;
 
     // Mapping from buyer address to cumulative purchase count and total purchase amount.
@@ -20,16 +21,10 @@ contract PurchaseTracker is HasSecurityContext {
     // mapping for sellers 
     mapping(address => uint256) public totalSalesCount;
     mapping(address => uint256) public totalSalesAmount;
-
-    // mapping rewards distributed
-    mapping(address => uint256) public rewardsDistributed;
 
     // Store details about each purchase (keyed by the unique payment ID).
     mapping(bytes32 => Purchase) public purchases;
 
-    //Comunity Vault address
-    address public communityVault;
-
     // loot token 
     IERC20 public lootToken;
 
@@ -50,8 +45,7 @@ contract PurchaseTracker is HasSecurityContext {
         _;
     }
 
-    constructor(ISecurityContext securityContext, address _communityVault, address _lootToken) {
-        communityVault = _communityVault;
+    constructor(ISecurityContext securityContext, address _lootToken) {
         lootToken = IERC20(_lootToken);
         _setSecurityContext(securityContext);
     }
@@ -96,23 +90,20 @@ contract PurchaseTracker is HasSecurityContext {
 
         emit PurchaseRecorded(paymentId, buyer, amount);
     }
-
-    // distrubte reward from communtiy vault
-    function distributeReward(address recipient) external {
-        // for every purchase or sale made by the recipient, distribute 1 loot token
-        uint256 totalPurchase = totalPurchaseCount[recipient];
-        uint256 totalSales = totalSalesCount[recipient];
-        uint256 rewardsDist = rewardsDistributed[recipient];
-
-        uint256 totalRewards = totalPurchase + totalSales - rewardsDist;
 
-        require(totalRewards > 0, "PurchaseTracker: No rewards to distribute");
+    function getPurchaseCount(address recipient) external view returns (uint256) {
+        return totalPurchaseCount[recipient];
+    }
 
-        // transfer loot token from community vault to recipient
-        lootToken.safeTransferFrom(communityVault, recipient, totalRewards);
+    function getPurchaseAmount(address recipient) external view returns (uint256) {
+        return totalPurchaseAmount[recipient];
+    }
 
-        rewardsDistributed[recipient] += totalRewards;
+    function getSalesCount(address recipient) external view returns (uint256) {
+        return totalSalesCount[recipient];
     }
 
-
+    function getSalesAmount(address recipient) external view returns (uint256) {
+        return totalSalesAmount[recipient];
+    }
 }