refactor: use balances for a single user instead of multi-user

Author: maxpushka

src/vault/TradingVault.sol

@@ -16,7 +16,8 @@ contract TradingVault is ITradingVault, IAuthorizableV2, ReentrancyGuard, Ownabl
     /// that may not return a boolean value on success.
     using SafeERC20 for IERC20;
 
-    mapping(address user => mapping(address token => uint256 balance)) internal _balances;
+    mapping(address token => uint256 balance) internal _balances;
+    mapping(address token => uint256 balance) internal _broker_balances;
     uint256 internal _nonce;
 
     address public broker;
@@ -34,14 +35,14 @@ contract TradingVault is ITradingVault, IAuthorizableV2, ReentrancyGuard, Ownabl
 
     // ---------- View functions ----------
 
-    function balanceOf(address user, address token) external view returns (uint256) {
-        return _balances[user][token];
+    function balanceOf(address token) external view returns (uint256) {
+        return _balances[token];
     }
 
-    function balancesOfTokens(address user, address[] calldata tokens) external view returns (uint256[] memory) {
+    function balancesOfTokens(address[] calldata tokens) external view returns (uint256[] memory) {
         uint256[] memory balances = new uint256[](tokens.length);
         for (uint256 i = 0; i < tokens.length; i++) {
-            balances[i] = _balances[user][tokens[i]];
+            balances[i] = _balances[tokens[i]];
         }
         return balances;
     }
@@ -65,7 +66,6 @@ contract TradingVault is ITradingVault, IAuthorizableV2, ReentrancyGuard, Ownabl
 
     function deposit(ITradingStructs.Intent calldata intent) external payable notZeroAddress(intent.trader) {
         address sender = msg.sender;
-        address recipient = intent.trader;
         require(_nonce == intent.nonce, NonceMismatch(_nonce, intent.nonce));
 
         _nonce++;
@@ -77,22 +77,21 @@ contract TradingVault is ITradingVault, IAuthorizableV2, ReentrancyGuard, Ownabl
 
             if (token == address(0)) {
                 require(msg.value == amount, IncorrectValue());
-                _balances[recipient][address(0)] += amount;
+                _balances[address(0)] += amount;
             } else {
                 require(msg.value == 0, IncorrectValue());
-                _balances[recipient][token] += amount;
+                _balances[token] += amount;
                 IERC20(token).safeTransferFrom(sender, address(this), amount);
             }
 
-            emit Deposited(recipient, token, amount);
+            emit Deposited(intent.trader, token, amount);
         }
     }
 
     function withdraw(
         ITradingStructs.Intent calldata intent,
         bytes calldata additionalAuthData
     ) external notZeroAddress(intent.trader) {
-        address account = intent.trader;
         require(_nonce == intent.nonce, NonceMismatch(_nonce, intent.nonce));
 
         bytes memory authData = abi.encodePacked(
@@ -115,11 +114,12 @@ contract TradingVault is ITradingVault, IAuthorizableV2, ReentrancyGuard, Ownabl
         for (uint256 i = 0; i < intent.allocations.length; i++) {
             address asset = intent.allocations[i].asset;
             uint256 amount = intent.allocations[i].amount;
-            uint256 currentBalance = _balances[account][asset];
+            uint256 currentBalance = _balances[asset];
             require(currentBalance >= amount, InsufficientBalance(asset, amount, currentBalance));
 
-            _balances[account][asset] -= amount;
+            _balances[asset] -= amount;
 
+            address account = intent.trader;
             if (asset == address(0)) {
                 /// @dev using `call` instead of `transfer` to overcome 2300 gas ceiling that could make it revert with some AA wallets
                 (bool success, ) = account.call{value: amount}("");
@@ -174,28 +174,38 @@ contract TradingVault is ITradingVault, IAuthorizableV2, ReentrancyGuard, Ownabl
 
     // ---------- Internal functions ----------
 
+    function _getBalanceMapping(address user) private view returns (mapping(address => uint256) storage) {
+        return user == broker ? _broker_balances : _balances;
+    }
+
     function _checkAndVaultSwap(
         address sender,
         address receiver,
         ITradingStructs.Allocation memory alloc
     ) internal virtual {
-        uint256 balance = _balances[sender][alloc.asset];
+        mapping(address => uint256) storage senderBalances = _getBalanceMapping(sender);
+        mapping(address => uint256) storage receiverBalances = _getBalanceMapping(receiver);
+
+        uint256 balance = senderBalances[alloc.asset];
         require(balance >= alloc.amount, InsufficientBalance(alloc.asset, alloc.amount, balance));
 
-        _balances[sender][alloc.asset] -= alloc.amount;
-        _balances[receiver][alloc.asset] += alloc.amount;
+        senderBalances[alloc.asset] -= alloc.amount;
+        receiverBalances[alloc.asset] += alloc.amount;
     }
 
     function _checkAndVaultSendAccount(
         address sender,
         address receiver,
         ITradingStructs.Allocation memory alloc
     ) internal virtual {
-        uint256 balance = _balances[sender][alloc.asset];
+        mapping(address => uint256) storage senderBalances = _getBalanceMapping(sender);
+        mapping(address => uint256) storage receiverBalances = _getBalanceMapping(receiver);
+
+        uint256 balance = senderBalances[alloc.asset];
         require(balance >= alloc.amount, InsufficientBalance(alloc.asset, alloc.amount, balance));
 
-        _balances[sender][alloc.asset] -= alloc.amount;
-        _balances[receiver][alloc.asset] += alloc.amount;
+        senderBalances[alloc.asset] -= alloc.amount;
+        receiverBalances[alloc.asset] += alloc.amount;
     }
 
     function _accountSendVault(
@@ -204,7 +214,9 @@ contract TradingVault is ITradingVault, IAuthorizableV2, ReentrancyGuard, Ownabl
         ITradingStructs.Allocation memory alloc
     ) internal virtual {
         IERC20(alloc.asset).safeTransferFrom(sender, address(this), alloc.amount);
-        _balances[receiver][alloc.asset] += alloc.amount;
+
+        mapping(address => uint256) storage receiverBalances = _getBalanceMapping(receiver);
+        receiverBalances[alloc.asset] += alloc.amount;
     }
 
     function _accountSwap(address sender, address receiver, ITradingStructs.Allocation memory alloc) internal virtual {