From 6ccc530605bc38e16bdc6b9e33744b3a94526e57 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Mon, 16 Sep 2024 14:36:40 -0300
Subject: [PATCH 01/22] feat: adding the splitter contract at marketplace
 contracts

---
 contracts/SplitterContract.sol | 146 ++++++++++++
 test/SplitterContract.test.ts  | 400 +++++++++++++++++++++++++++++++++
 2 files changed, 546 insertions(+)
 create mode 100644 contracts/SplitterContract.sol
 create mode 100644 test/SplitterContract.test.ts

diff --git a/contracts/SplitterContract.sol b/contracts/SplitterContract.sol
new file mode 100644
index 0000000..55d6503
--- /dev/null
+++ b/contracts/SplitterContract.sol
@@ -0,0 +1,146 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import '@openzeppelin/contracts/token/ERC20/IERC20.sol';
+import '@openzeppelin/contracts/security/ReentrancyGuard.sol';
+
+contract ERC20Splitter is ReentrancyGuard {
+    mapping(address => mapping(address => uint256)) public balances;
+    mapping(address => address[]) private _userTokens;
+    mapping(address => mapping(address => bool)) private _hasToken;
+
+    /** Events **/
+
+    event Deposit(
+        address indexed depositor,
+        address[] tokenAddresses,
+        uint256[] amounts,
+        uint256[][] shares,
+        address[][] recipients
+    );
+    event Withdraw(address indexed user, address[] tokenAddresses, uint256[] amounts);
+
+    uint256 public constant MAX_SHARES = 10000;
+
+    /** External Functions **/
+
+    /// @notice Deposits ERC20 or native tokens and splits between recipients based on shares.
+    /// @param tokenAddresses Array of token addresses (use address(0) for native tokens).
+    /// @param amounts Array of amounts for each token.
+    /// @param shares Array of share percentages (out of 10000) for each recipient.
+    /// @param recipients Array of recipients for each token.
+    function deposit(
+        address[] calldata tokenAddresses,
+        uint256[] calldata amounts,
+        uint256[][] calldata shares,
+        address[][] calldata recipients
+    ) external payable nonReentrant {
+        require(tokenAddresses.length == amounts.length, 'ERC20Splitter: Invalid input lengths');
+        require(
+            tokenAddresses.length == shares.length && tokenAddresses.length == recipients.length,
+            'ERC20Splitter: Mismatched input sizes'
+        );
+
+        uint256 totalEthAmount = 0;
+
+        for (uint256 i = 0; i < tokenAddresses.length; i++) {
+            if (tokenAddresses[i] == address(0)) {
+                totalEthAmount += amounts[i];
+            }
+            _splitTokens(tokenAddresses[i], amounts[i], shares[i], recipients[i]);
+        }
+
+        require(msg.value == totalEthAmount, 'ERC20Splitter: Incorrect native token amount sent');
+
+        emit Deposit(msg.sender, tokenAddresses, amounts, shares, recipients);
+    }
+
+    /// @notice Withdraw all tokens that the caller is entitled to.
+    /// Tokens are automatically determined based on previous deposits.
+    function withdraw() external nonReentrant {
+        address[] storage userTokens = _userTokens[msg.sender];
+        require(userTokens.length > 0, 'ERC20Splitter: No tokens to withdraw');
+
+        address[] memory withdrawnTokens = new address[](userTokens.length);
+        uint256[] memory withdrawnAmounts = new uint256[](userTokens.length);
+
+        for (uint256 i = 0; i < userTokens.length; i++) {
+            address tokenAddress = userTokens[i];
+            uint256 amount = balances[tokenAddress][msg.sender];
+
+            if (amount > 0) {
+                balances[tokenAddress][msg.sender] = 0;
+
+                if (tokenAddress == address(0)) {
+                    (bool success, ) = msg.sender.call{ value: amount }('');
+                    require(success, 'ERC20Splitter: Failed to send Ether');
+                } else {
+                    require(tokenAddress != address(0), 'ERC20Splitter: Invalid token address');
+
+                    require(
+                        IERC20(tokenAddress).transferFrom(address(this), msg.sender, amount),
+                        'ERC20Splitter: TransferFrom failed'
+                    );
+                }
+
+                withdrawnTokens[i] = tokenAddress;
+                withdrawnAmounts[i] = amount;
+            }
+
+            delete _hasToken[msg.sender][tokenAddress];
+        }
+
+        delete _userTokens[msg.sender];
+
+        emit Withdraw(msg.sender, withdrawnTokens, withdrawnAmounts);
+    }
+
+    /** Internal Functions **/
+
+    /// @notice Internal function to split the tokens among recipients.
+    /// @param tokenAddress The address of the token being split (use address(0) for native tokens).
+    /// @param amount The amount of tokens to be split.
+    /// @param shares Array of share percentages (out of 10000) for each recipient.
+    /// @param recipients Array of recipients for the token.
+    function _splitTokens(
+        address tokenAddress,
+        uint256 amount,
+        uint256[] calldata shares,
+        address[] calldata recipients
+    ) internal {
+        require(shares.length == recipients.length, 'ERC20Splitter: Shares and recipients length mismatch');
+        require(amount > 0, 'ERC20Splitter: Amount must be greater than zero');
+
+        uint256 totalSharePercentage = 0;
+
+        for (uint256 i = 0; i < shares.length; i++) {
+            totalSharePercentage += shares[i];
+        }
+
+        require(totalSharePercentage == MAX_SHARES, 'ERC20Splitter: Shares must sum to 100%');
+
+        if (tokenAddress != address(0)) {
+            require(
+                IERC20(tokenAddress).transferFrom(msg.sender, address(this), amount),
+                'ERC20Splitter: Transfer failed'
+            );
+        }
+
+        for (uint256 i = 0; i < recipients.length; i++) {
+            uint256 recipientAmount = (amount * shares[i]) / MAX_SHARES;
+            balances[tokenAddress][recipients[i]] += recipientAmount;
+
+            _addTokenForUser(recipients[i], tokenAddress);
+        }
+    }
+
+    /// @notice Adds a token to the list of tokens a user has received (for automatic withdrawals).
+    /// @param recipient The recipient of the token.
+    /// @param tokenAddress The address of the token.
+    function _addTokenForUser(address recipient, address tokenAddress) internal {
+        if (!_hasToken[recipient][tokenAddress]) {
+            _userTokens[recipient].push(tokenAddress);
+            _hasToken[recipient][tokenAddress] = true;
+        }
+    }
+}
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
new file mode 100644
index 0000000..3686e6f
--- /dev/null
+++ b/test/SplitterContract.test.ts
@@ -0,0 +1,400 @@
+/* eslint-disable no-unexpected-multiline */
+import { ethers, network } from 'hardhat'
+import { loadFixture } from '@nomicfoundation/hardhat-network-helpers'
+import { expect } from 'chai'
+import { MockERC20, ERC20Splitter } from '../typechain-types'
+import { AddressZero } from '../utils/constants'
+
+describe('ERC20Splitter', () => {
+  let splitter: ERC20Splitter
+  let mockERC20: MockERC20
+  let owner: Awaited<ReturnType<typeof ethers.getSigner>>
+  let recipient1: Awaited<ReturnType<typeof ethers.getSigner>>
+  let recipient2: Awaited<ReturnType<typeof ethers.getSigner>>
+  let recipient3: Awaited<ReturnType<typeof ethers.getSigner>>
+
+  const tokenAmount = ethers.parseEther('100')
+  const ethAmount = ethers.parseEther('1')
+
+  before(async function () {
+    // prettier-ignore
+    [owner, recipient1, recipient2, recipient3] = await ethers.getSigners()
+  })
+
+  async function deploySplitterContracts() {
+    const MockERC20 = await ethers.getContractFactory('MockERC20')
+    const ERC20Splitter = await ethers.getContractFactory('ERC20Splitter')
+
+    const mockERC20 = await MockERC20.deploy()
+    await mockERC20.waitForDeployment()
+
+    const splitter = await ERC20Splitter.deploy()
+    await splitter.waitForDeployment()
+
+    return { mockERC20, splitter }
+  }
+
+  beforeEach(async () => {
+    const contracts = await loadFixture(deploySplitterContracts)
+    mockERC20 = contracts.mockERC20
+    splitter = contracts.splitter
+
+    // Mint tokens to the owner
+    await mockERC20.connect(owner).mint(owner, ethers.parseEther('1000'))
+
+    const splitterAddress = await splitter.getAddress()
+
+    await network.provider.send('hardhat_setBalance', [
+      splitterAddress,
+      ethers.toQuantity(ethers.parseEther('2')), // Setting 2 Ether
+    ])
+
+    await network.provider.request({
+      method: 'hardhat_impersonateAccount',
+      params: [splitterAddress],
+    })
+    const splitterSigner = await ethers.getSigner(splitterAddress)
+
+    await mockERC20.connect(splitterSigner).approve(splitterAddress, ethers.MaxUint256)
+
+    await network.provider.request({
+      method: 'hardhat_stopImpersonatingAccount',
+      params: [splitterAddress],
+    })
+  })
+
+  describe('Main Functions', async () => {
+    describe('Deposit', async () => {
+      beforeEach(async () => {
+        await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
+      })
+
+      it('Should deposit ERC20 tokens and split them between recipients', async () => {
+        const shares = [[5000, 3000, 2000]] // 50%, 30%, 20%
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await expect(
+          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], shares, recipients),
+        ).to.emit(splitter, 'Deposit')
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('50'))
+        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(ethers.parseEther('30'))
+        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('20'))
+      })
+
+      it('Should deposit native tokens (ETH) and split them between recipients', async () => {
+        const shares = [[5000, 3000, 2000]]
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await expect(
+          splitter.connect(owner).deposit([AddressZero], [ethAmount], shares, recipients, {
+            value: ethAmount,
+          }),
+        ).to.emit(splitter, 'Deposit')
+
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(ethers.parseEther('0.5'))
+        expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(ethers.parseEther('0.3'))
+        expect(await splitter.balances(AddressZero, recipient3.address)).to.equal(ethers.parseEther('0.2'))
+      })
+
+      it('Should revert if shares do not sum to 100%', async () => {
+        const invalidShares = [[4000, 4000, 1000]] // Sums to 90%
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await expect(
+          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], invalidShares, recipients),
+        ).to.be.revertedWith('ERC20Splitter: Shares must sum to 100%')
+      })
+
+      it('Should revert if the number of shares and recipients do not match', async () => {
+        const invalidShares = [[5000, 3000]] // Only 2 shares
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]] // 3 recipients
+
+        await expect(
+          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], invalidShares, recipients),
+        ).to.be.revertedWith('ERC20Splitter: Shares and recipients length mismatch')
+      })
+
+      it('Should handle multiple native token (ETH) deposits in a single transaction', async () => {
+        const ethShares = [
+          [5000, 5000],
+          [6000, 4000],
+        ]
+        const ethRecipients1 = [recipient1.address, recipient2.address] // Recipients for first ETH deposit
+        const ethRecipients2 = [recipient2.address, recipient3.address] // Recipients for second ETH deposit
+
+        const ethAmount1 = ethers.parseEther('1') // First ETH deposit (1 ETH)
+        const ethAmount2 = ethers.parseEther('2') // Second ETH deposit (2 ETH)
+
+        await expect(
+          splitter
+            .connect(owner)
+            .deposit(
+              [AddressZero, AddressZero],
+              [ethAmount1, ethAmount2],
+              [ethShares[0], ethShares[1]],
+              [ethRecipients1, ethRecipients2],
+              { value: ethAmount1 + ethAmount2 },
+            ),
+        ).to.emit(splitter, 'Deposit')
+
+        // Check balances for recipient1 (50% of 1 ETH)
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(ethers.parseEther('0.5'))
+
+        // Check balances for recipient2 (50% of 1 ETH + 60% of 2 ETH = 0.5 + 1.2 = 1.7 ETH)
+        expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(ethers.parseEther('1.7'))
+
+        // Check balances for recipient3 (40% of 2 ETH = 0.8 ETH)
+        expect(await splitter.balances(AddressZero, recipient3.address)).to.equal(ethers.parseEther('0.8'))
+      })
+
+      it('Should handle both native token (ETH) and ERC-20 deposits in a single transaction', async () => {
+        const ethShares = [[5000, 5000]]
+        const erc20Shares = [[6000, 4000]]
+
+        const ethRecipients = [recipient1.address, recipient2.address]
+        const erc20Recipients = [recipient2.address, recipient3.address]
+
+        const ethAmount = ethers.parseEther('1') // ETH deposit (1 ETH)
+        const erc20Amount = ethers.parseEther('100') // ERC-20 deposit (100 tokens)
+
+        await mockERC20.connect(owner).approve(splitter.getAddress(), erc20Amount)
+
+        await expect(
+          splitter
+            .connect(owner)
+            .deposit(
+              [AddressZero, mockERC20.getAddress()],
+              [ethAmount, erc20Amount],
+              [ethShares[0], erc20Shares[0]],
+              [ethRecipients, erc20Recipients],
+              { value: ethAmount },
+            ),
+        ).to.emit(splitter, 'Deposit')
+
+        // Check balances for recipient1 (50% of 1 ETH)
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(ethers.parseEther('0.5'))
+
+        // Check balances for recipient2 (50% of 1 ETH + 60% of 100 ERC-20 tokens = 0.5 ETH + 60 tokens)
+        expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(ethers.parseEther('0.5'))
+        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(ethers.parseEther('60'))
+
+        // Check balances for recipient3 (40% of 100 ERC-20 tokens = 40 tokens)
+        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('40'))
+      })
+    })
+
+    describe('Withdraw', async () => {
+      beforeEach(async () => {
+        const shares = [[5000, 3000, 2000]]
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
+        await splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], shares, recipients)
+      })
+
+      it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
+        await expect(splitter.connect(recipient1).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
+
+        expect(await splitter.balances(await mockERC20.getAddress(), recipient1.address)).to.equal(0)
+      })
+
+      it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
+        const shares = [[5000, 3000, 2000]]
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await splitter.connect(owner).deposit([AddressZero], [ethAmount], shares, recipients, {
+          value: ethAmount,
+        })
+
+        await expect(splitter.connect(recipient1).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient1.address,
+            [await mockERC20.getAddress(), AddressZero], // Expect both ERC-20 and native token
+            [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
+          )
+
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
+      })
+    })
+
+    describe('Withdraw ERC-20 and Native Tokens', async () => {
+      beforeEach(async () => {
+        const shares = [[5000, 3000, 2000]]
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
+        await splitter.connect(owner).deposit([await mockERC20.getAddress()], [tokenAmount], shares, recipients)
+      })
+
+      it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
+        await expect(splitter.connect(recipient1).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
+      })
+
+      it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
+        const shares = [[5000, 3000, 2000]]
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await splitter.connect(owner).deposit([AddressZero], [ethAmount], shares, recipients, {
+          value: ethAmount,
+        })
+
+        await expect(splitter.connect(recipient1).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient1.address,
+            [await mockERC20.getAddress(), AddressZero], // Expect both ERC-20 and native token
+            [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
+          )
+
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
+      })
+    })
+
+    describe('Withdraw Only Native Tokens (ETH)', async () => {
+      beforeEach(async () => {
+        const shares = [[5000, 3000, 2000]]
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await splitter.connect(owner).deposit([AddressZero], [ethAmount], shares, recipients, {
+          value: ethAmount,
+        })
+      })
+
+      it('Should allow a recipient to withdraw only their split native tokens (ETH)', async () => {
+        await expect(splitter.connect(recipient1).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient1.address,
+            [AddressZero], // Expect only native token (ETH)
+            [ethers.parseEther('0.5')], // Expect 0.5 ETH (50% of 1 ETH)
+          )
+
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
+      })
+    })
+
+    describe('Deposit ETH for recipient1 and ERC-20 for other recipients', async () => {
+      beforeEach(async () => {
+        const ethShares = [[10000]] // 100% for recipient1 (ETH)
+        const erc20Shares = [[5000, 5000]] // 50% for recipient2, 50% for recipient3 (ERC-20)
+        const ethRecipients = [[recipient1.address]] // Only recipient1 gets ETH
+        const erc20Recipients = [
+          [recipient2.address, recipient3.address], // recipient2 and recipient3 get ERC-20 tokens
+        ]
+        await splitter.connect(owner).deposit([AddressZero], [ethAmount], ethShares, ethRecipients, {
+          value: ethAmount,
+        })
+
+        // Then, deposit ERC-20 tokens for recipient2 and recipient3
+        await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
+        await splitter
+          .connect(owner)
+          .deposit([await mockERC20.getAddress()], [tokenAmount], erc20Shares, erc20Recipients)
+      })
+
+      it('Should allow recipient1 to withdraw only their ETH and other recipients to withdraw their ERC-20 tokens', async () => {
+        await expect(splitter.connect(recipient1).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient1.address,
+            [AddressZero], // Only native token (ETH)
+            [ethers.parseEther('1')], // Full 1 ETH
+          )
+
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
+
+        await expect(splitter.connect(recipient2).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient2.address,
+            [await mockERC20.getAddress()], // Only ERC-20 token
+            [ethers.parseEther('50')], // 50% of ERC-20 tokens
+          )
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(0)
+
+        await expect(splitter.connect(recipient3).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient3.address,
+            [await mockERC20.getAddress()], // Only ERC-20 token
+            [ethers.parseEther('50')], // 50% of ERC-20 tokens
+          )
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(0)
+      })
+    })
+    describe('Withdraw for both native tokens (ETH) and ERC-20 tokens multiples addresses 0', () => {
+      let ethShares, erc20Shares
+      let ethRecipients, erc20Recipients
+      let ethAmount, erc20Amount
+
+      beforeEach(async () => {
+        // Define shares and recipients for both ETH and ERC-20
+        ethShares = [[5000, 5000]] // 50%-50% for ETH
+        erc20Shares = [[6000, 4000]] // 60%-40% for ERC-20
+
+        ethRecipients = [recipient1.address, recipient2.address]
+        erc20Recipients = [recipient2.address, recipient3.address]
+
+        ethAmount = ethers.parseEther('1') // 1 ETH
+        erc20Amount = ethers.parseEther('100') // 100 ERC-20 tokens
+
+        await mockERC20.connect(owner).approve(splitter.getAddress(), erc20Amount)
+
+        await splitter
+          .connect(owner)
+          .deposit(
+            [AddressZero, mockERC20.getAddress()],
+            [ethAmount, erc20Amount],
+            [ethShares[0], erc20Shares[0]],
+            [ethRecipients, erc20Recipients],
+            { value: ethAmount },
+          )
+      })
+
+      it('Should allow recipient1 to withdraw only ETH', async () => {
+        await expect(splitter.connect(recipient1).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient1.address,
+            [AddressZero], // Only native token (ETH)
+            [ethers.parseEther('0.5')], // 50% of 1 ETH
+          )
+
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
+      })
+
+      it('Should allow recipient2 to withdraw both ETH and ERC-20 tokens', async () => {
+        await expect(splitter.connect(recipient2).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(
+            recipient2.address,
+            [AddressZero, await mockERC20.getAddress()], // First ETH, then ERC-20
+            [ethers.parseEther('0.5'), ethers.parseEther('60')], // 50% of 1 ETH and 60 ERC-20 tokens
+          )
+
+        expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(0)
+      })
+
+      it('Should allow recipient3 to withdraw only ERC-20 tokens', async () => {
+        await expect(splitter.connect(recipient3).withdraw())
+          .to.emit(splitter, 'Withdraw')
+          .withArgs(recipient3.address, [await mockERC20.getAddress()], [ethers.parseEther('40')])
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(0)
+      })
+    })
+  })
+})

From 06eb6aea4869fdb3c1774c0ccf235317faeb2050 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Mon, 16 Sep 2024 19:12:30 -0300
Subject: [PATCH 02/22] fix: Comments from pr

---
 ...SplitterContract.sol => ERC20Splitter.sol} |  76 ++++-----
 contracts/mocks/MaliciousERC20.sol            |  22 +++
 contracts/mocks/MaliciousRecipient.sol        |  15 ++
 test/SplitterContract.test.ts                 | 157 ++++++++++++++++--
 4 files changed, 218 insertions(+), 52 deletions(-)
 rename contracts/{SplitterContract.sol => ERC20Splitter.sol} (66%)
 create mode 100644 contracts/mocks/MaliciousERC20.sol
 create mode 100644 contracts/mocks/MaliciousRecipient.sol

diff --git a/contracts/SplitterContract.sol b/contracts/ERC20Splitter.sol
similarity index 66%
rename from contracts/SplitterContract.sol
rename to contracts/ERC20Splitter.sol
index 55d6503..f021ad9 100644
--- a/contracts/SplitterContract.sol
+++ b/contracts/ERC20Splitter.sol
@@ -1,13 +1,16 @@
-// SPDX-License-Identifier: MIT
-pragma solidity ^0.8.0;
+// SPDX-License-Identifier: CC0-1.0
+pragma solidity 0.8.9;
 
 import '@openzeppelin/contracts/token/ERC20/IERC20.sol';
 import '@openzeppelin/contracts/security/ReentrancyGuard.sol';
 
 contract ERC20Splitter is ReentrancyGuard {
+    // tokenAddress => userAddress => balance
     mapping(address => mapping(address => uint256)) public balances;
-    mapping(address => address[]) private _userTokens;
-    mapping(address => mapping(address => bool)) private _hasToken;
+    // userAddress => tokenAddress[]
+    mapping(address => address[]) private userTokens;
+    // tokenAddress => boolean
+    mapping(address => mapping(address => bool)) private hasToken;
 
     /** Events **/
 
@@ -15,12 +18,12 @@ contract ERC20Splitter is ReentrancyGuard {
         address indexed depositor,
         address[] tokenAddresses,
         uint256[] amounts,
-        uint256[][] shares,
+        uint16[][] shares,
         address[][] recipients
     );
-    event Withdraw(address indexed user, address[] tokenAddresses, uint256[] amounts);
+    event Withdraw(address indexed user, uint256[] amounts);
 
-    uint256 public constant MAX_SHARES = 10000;
+    uint16 public constant MAX_SHARES = 10000;
 
     /** External Functions **/
 
@@ -32,7 +35,7 @@ contract ERC20Splitter is ReentrancyGuard {
     function deposit(
         address[] calldata tokenAddresses,
         uint256[] calldata amounts,
-        uint256[][] calldata shares,
+        uint16[][] calldata shares,
         address[][] calldata recipients
     ) external payable nonReentrant {
         require(tokenAddresses.length == amounts.length, 'ERC20Splitter: Invalid input lengths');
@@ -58,41 +61,40 @@ contract ERC20Splitter is ReentrancyGuard {
     /// @notice Withdraw all tokens that the caller is entitled to.
     /// Tokens are automatically determined based on previous deposits.
     function withdraw() external nonReentrant {
-        address[] storage userTokens = _userTokens[msg.sender];
-        require(userTokens.length > 0, 'ERC20Splitter: No tokens to withdraw');
+        address payable to = payable(msg.sender);
+        address[] storage senderTokens = userTokens[to];
 
-        address[] memory withdrawnTokens = new address[](userTokens.length);
-        uint256[] memory withdrawnAmounts = new uint256[](userTokens.length);
-
-        for (uint256 i = 0; i < userTokens.length; i++) {
-            address tokenAddress = userTokens[i];
-            uint256 amount = balances[tokenAddress][msg.sender];
+        if (senderTokens.length == 0) {
+            return;
+        }
 
-            if (amount > 0) {
-                balances[tokenAddress][msg.sender] = 0;
+        uint256[] memory withdrawnAmounts = new uint256[](senderTokens.length);
 
-                if (tokenAddress == address(0)) {
-                    (bool success, ) = msg.sender.call{ value: amount }('');
-                    require(success, 'ERC20Splitter: Failed to send Ether');
-                } else {
-                    require(tokenAddress != address(0), 'ERC20Splitter: Invalid token address');
+        for (uint256 i = 0; i < senderTokens.length; i++) {
+            address tokenAddress = senderTokens[i];
+            uint256 amount = balances[tokenAddress][to];
 
-                    require(
-                        IERC20(tokenAddress).transferFrom(address(this), msg.sender, amount),
-                        'ERC20Splitter: TransferFrom failed'
-                    );
-                }
+            require(amount > 0, 'ERC20Splitter: Amount to withdraw must be greater than zero');
+            balances[tokenAddress][to] = 0;
 
-                withdrawnTokens[i] = tokenAddress;
-                withdrawnAmounts[i] = amount;
+            if (tokenAddress == address(0)) {
+                (bool success, ) = to.call{ value: amount }('');
+                require(success, 'ERC20Splitter: Failed to send Ether');
+            } else {
+                require(
+                    IERC20(tokenAddress).transferFrom(address(this), to, amount),
+                    'ERC20Splitter: TransferFrom failed'
+                );
             }
 
-            delete _hasToken[msg.sender][tokenAddress];
+            withdrawnAmounts[i] = amount;
+
+            delete hasToken[to][tokenAddress];
         }
 
-        delete _userTokens[msg.sender];
+        delete userTokens[to];
 
-        emit Withdraw(msg.sender, withdrawnTokens, withdrawnAmounts);
+        emit Withdraw(to, withdrawnAmounts);
     }
 
     /** Internal Functions **/
@@ -105,7 +107,7 @@ contract ERC20Splitter is ReentrancyGuard {
     function _splitTokens(
         address tokenAddress,
         uint256 amount,
-        uint256[] calldata shares,
+        uint16[] calldata shares,
         address[] calldata recipients
     ) internal {
         require(shares.length == recipients.length, 'ERC20Splitter: Shares and recipients length mismatch');
@@ -138,9 +140,9 @@ contract ERC20Splitter is ReentrancyGuard {
     /// @param recipient The recipient of the token.
     /// @param tokenAddress The address of the token.
     function _addTokenForUser(address recipient, address tokenAddress) internal {
-        if (!_hasToken[recipient][tokenAddress]) {
-            _userTokens[recipient].push(tokenAddress);
-            _hasToken[recipient][tokenAddress] = true;
+        if (!hasToken[recipient][tokenAddress]) {
+            userTokens[recipient].push(tokenAddress);
+            hasToken[recipient][tokenAddress] = true;
         }
     }
 }
diff --git a/contracts/mocks/MaliciousERC20.sol b/contracts/mocks/MaliciousERC20.sol
new file mode 100644
index 0000000..dc655c5
--- /dev/null
+++ b/contracts/mocks/MaliciousERC20.sol
@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: CC0-1.0
+
+pragma solidity ^0.8.0;
+
+import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
+
+contract MaliciousERC20 is ERC20 {
+    constructor() ERC20("MaliciousToken", "MTK") {}
+
+    function transferFrom(
+        address sender,
+        address recipient,
+        uint256 amount
+    ) public override returns (bool) {
+        return false;
+    }
+
+    // Mint function for testing purposes
+    function mint(address to, uint256 amount) external {
+        _mint(to, amount);
+    }
+}
diff --git a/contracts/mocks/MaliciousRecipient.sol b/contracts/mocks/MaliciousRecipient.sol
new file mode 100644
index 0000000..2bd7ac7
--- /dev/null
+++ b/contracts/mocks/MaliciousRecipient.sol
@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: CC0-1.0
+
+// contracts/MaliciousRecipient.sol
+pragma solidity ^0.8.0;
+
+contract MaliciousRecipient {
+    // Fallback function that reverts when receiving Ether
+    fallback() external payable {
+        revert("MaliciousRecipient: Reverting on receive");
+    }
+
+    receive() external payable {
+        revert("MaliciousRecipient: Reverting on receive");
+    }
+}
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 3686e6f..af97e36 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -2,42 +2,54 @@
 import { ethers, network } from 'hardhat'
 import { loadFixture } from '@nomicfoundation/hardhat-network-helpers'
 import { expect } from 'chai'
-import { MockERC20, ERC20Splitter } from '../typechain-types'
+import { MockERC20, ERC20Splitter, MaliciousRecipient, MaliciousERC20 } from '../typechain-types'
 import { AddressZero } from '../utils/constants'
 
 describe('ERC20Splitter', () => {
   let splitter: ERC20Splitter
   let mockERC20: MockERC20
+  let maliciousERC20: MaliciousERC20
   let owner: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient1: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient2: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient3: Awaited<ReturnType<typeof ethers.getSigner>>
+  let anotherUser: Awaited<ReturnType<typeof ethers.getSigner>>
+  let maliciousRecipient: MaliciousRecipient
 
   const tokenAmount = ethers.parseEther('100')
   const ethAmount = ethers.parseEther('1')
 
   before(async function () {
     // prettier-ignore
-    [owner, recipient1, recipient2, recipient3] = await ethers.getSigners()
+    [owner, recipient1, recipient2, recipient3, anotherUser] = await ethers.getSigners()
   })
 
   async function deploySplitterContracts() {
     const MockERC20 = await ethers.getContractFactory('MockERC20')
     const ERC20Splitter = await ethers.getContractFactory('ERC20Splitter')
 
+    const MaliciousRecipientFactory = await ethers.getContractFactory('MaliciousRecipient')
+    maliciousRecipient = await MaliciousRecipientFactory.deploy()
+    await maliciousRecipient.waitForDeployment()
+
+    const MaliciousERC20Factory = await ethers.getContractFactory('MaliciousERC20')
+    maliciousERC20 = await MaliciousERC20Factory.deploy()
+    await maliciousERC20.waitForDeployment()
+
     const mockERC20 = await MockERC20.deploy()
     await mockERC20.waitForDeployment()
 
     const splitter = await ERC20Splitter.deploy()
     await splitter.waitForDeployment()
 
-    return { mockERC20, splitter }
+    return { mockERC20, splitter, maliciousERC20 }
   }
 
   beforeEach(async () => {
     const contracts = await loadFixture(deploySplitterContracts)
     mockERC20 = contracts.mockERC20
     splitter = contracts.splitter
+    maliciousERC20 = contracts.maliciousERC20
 
     // Mint tokens to the owner
     await mockERC20.connect(owner).mint(owner, ethers.parseEther('1000'))
@@ -61,6 +73,9 @@ describe('ERC20Splitter', () => {
       method: 'hardhat_stopImpersonatingAccount',
       params: [splitterAddress],
     })
+
+    const tokenAmount = ethers.parseEther('100')
+    await maliciousERC20.mint(splitter, tokenAmount)
   })
 
   describe('Main Functions', async () => {
@@ -115,6 +130,91 @@ describe('ERC20Splitter', () => {
         ).to.be.revertedWith('ERC20Splitter: Shares and recipients length mismatch')
       })
 
+      it('Should revert when msg.value does not match the expected Ether amount', async () => {
+        const incorrectMsgValue = ethers.parseEther('1') // Incorrect Ether amount
+        const correctEtherAmount = ethers.parseEther('2') // Correct Ether amount to be split
+        const tokenAddresses = [ethers.ZeroAddress] // Using address(0) for Ether
+        const amounts = [correctEtherAmount] // Amount to split among recipients
+        const shares = [[5000, 3000, 2000]] // Shares summing up to 100%
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await expect(
+          splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, {
+            value: incorrectMsgValue, // Sending incorrect msg.value
+          }),
+        ).to.be.revertedWith('ERC20Splitter: Incorrect native token amount sent')
+      })
+      it('Should revert when amount is 0', async () => {
+        const incorrectMsgValue = ethers.parseEther('1') // Incorrect Ether amount
+        const correctEtherAmount = ethers.parseEther('2') // Correct Ether amount to be split
+        const tokenAddresses = [ethers.ZeroAddress] // Using address(0) for Ether
+        const amounts = [correctEtherAmount] // Amount to split among recipients
+        const shares = [[5000, 3000, 2000]] // Shares summing up to 100%
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await expect(
+          splitter.connect(owner).deposit(tokenAddresses, [0], shares, recipients, {
+            value: incorrectMsgValue, // Sending incorrect msg.value
+          }),
+        ).to.be.revertedWith('ERC20Splitter: Amount must be greater than zero')
+      })
+
+      it('Should revert when tokenAddresses and amounts lengths mismatch', async () => {
+        const tokenAddresses = [mockERC20.getAddress(), ethers.ZeroAddress]
+        const amounts = [ethers.parseEther('100')] // Length 1, intentional mismatch
+        const shares = [[5000, 3000, 2000]] // Correct length
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await expect(
+          splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, {
+            value: ethers.parseEther('0'), // No Ether sent
+          }),
+        ).to.be.revertedWith('ERC20Splitter: Invalid input lengths')
+      })
+
+      it('Should revert when tokenAddresses, shares, and recipients lengths mismatch', async () => {
+        const tokenAddresses = [mockERC20.getAddress(), ethers.ZeroAddress]
+        const amounts = [ethers.parseEther('100'), ethers.parseEther('2')]
+        const shares = [
+          [5000, 3000, 2000], // Length 1
+        ] // Length 1 (intentional mismatch)
+        const recipients = [
+          [recipient1.address, recipient2.address, recipient3.address],
+          [recipient1.address, recipient2.address, recipient3.address],
+        ] // Length 2
+
+        await expect(
+          splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, {
+            value: ethers.parseEther('2'),
+          }),
+        ).to.be.revertedWith('ERC20Splitter: Mismatched input sizes')
+      })
+
+      it('Should revert when shares and recipients lengths mismatch within sub-arrays', async () => {
+        const tokenAddresses = [mockERC20.getAddress()] // Length 1
+        const amounts = [ethers.parseEther('100')] // Length 1
+        const shares = [[5000, 3000, 2000]] // Length 1, sub-array length 3
+        const recipients = [
+          [recipient1.address, recipient2.address], // Length mismatch in sub-array
+        ] // Length 1, sub-array length 2
+
+        await expect(splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients)).to.be.revertedWith(
+          'ERC20Splitter: Shares and recipients length mismatch',
+        )
+      })
+
+      it('Should revert when ERC20 transferFrom fails during withdrawal', async () => {
+        const tokenAmount = ethers.parseEther('100')
+        const tokenAddresses = [await maliciousERC20.getAddress()]
+        const amounts = [tokenAmount]
+        const shares = [[10000]] // 100% share
+        const recipients = [[recipient1.address]]
+
+        await expect(splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients)).to.be.revertedWith(
+          'ERC20Splitter: Transfer failed',
+        )
+      })
+
       it('Should handle multiple native token (ETH) deposits in a single transaction', async () => {
         const ethShares = [
           [5000, 5000],
@@ -196,7 +296,7 @@ describe('ERC20Splitter', () => {
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
         await expect(splitter.connect(recipient1).withdraw())
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [ethers.parseEther('50')])
 
         expect(await splitter.balances(await mockERC20.getAddress(), recipient1.address)).to.equal(0)
       })
@@ -213,13 +313,47 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
-            [await mockERC20.getAddress(), AddressZero], // Expect both ERC-20 and native token
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
         expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
         expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
       })
+
+      it('Should handle withdraw() when user has no tokens', async () => {
+        await splitter.connect(anotherUser).withdraw()
+      })
+
+      it('Should revert when sending Ether to a recipient fails', async () => {
+        const malicious = await maliciousRecipient.getAddress()
+
+        await network.provider.request({
+          method: 'hardhat_impersonateAccount',
+          params: [malicious],
+        })
+
+        const ethAmount = ethers.parseEther('1')
+        const tokenAddresses = [ethers.ZeroAddress] // Ether represented by address zero
+        const amounts = [ethAmount]
+        const shares = [[10000]] // 100% share
+        const recipients = [[maliciousRecipient.getAddress()]]
+
+        await splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, {
+          value: ethAmount,
+        })
+
+        const maliciousSigner = await ethers.getSigner(malicious)
+
+        await network.provider.send('hardhat_setBalance', [
+          malicious,
+          ethers.toQuantity(ethers.parseEther('1')), // Setting 2 Ether
+        ])
+
+        // Attempt to withdraw as the malicious recipient
+        await expect(splitter.connect(maliciousSigner).withdraw()).to.be.revertedWith(
+          'ERC20Splitter: Failed to send Ether',
+        )
+      })
     })
 
     describe('Withdraw ERC-20 and Native Tokens', async () => {
@@ -234,7 +368,7 @@ describe('ERC20Splitter', () => {
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
         await expect(splitter.connect(recipient1).withdraw())
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [ethers.parseEther('50')])
 
         expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
       })
@@ -250,8 +384,7 @@ describe('ERC20Splitter', () => {
         await expect(splitter.connect(recipient1).withdraw())
           .to.emit(splitter, 'Withdraw')
           .withArgs(
-            recipient1.address,
-            [await mockERC20.getAddress(), AddressZero], // Expect both ERC-20 and native token
+            recipient1.address, // Expect both ERC-20 and native token
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
@@ -275,7 +408,6 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
-            [AddressZero], // Expect only native token (ETH)
             [ethers.parseEther('0.5')], // Expect 0.5 ETH (50% of 1 ETH)
           )
 
@@ -307,7 +439,6 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
-            [AddressZero], // Only native token (ETH)
             [ethers.parseEther('1')], // Full 1 ETH
           )
 
@@ -317,7 +448,6 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
-            [await mockERC20.getAddress()], // Only ERC-20 token
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
@@ -327,7 +457,6 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient3.address,
-            [await mockERC20.getAddress()], // Only ERC-20 token
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
@@ -368,7 +497,6 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
-            [AddressZero], // Only native token (ETH)
             [ethers.parseEther('0.5')], // 50% of 1 ETH
           )
 
@@ -380,7 +508,6 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
-            [AddressZero, await mockERC20.getAddress()], // First ETH, then ERC-20
             [ethers.parseEther('0.5'), ethers.parseEther('60')], // 50% of 1 ETH and 60 ERC-20 tokens
           )
 
@@ -391,7 +518,7 @@ describe('ERC20Splitter', () => {
       it('Should allow recipient3 to withdraw only ERC-20 tokens', async () => {
         await expect(splitter.connect(recipient3).withdraw())
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient3.address, [await mockERC20.getAddress()], [ethers.parseEther('40')])
+          .withArgs(recipient3.address, [ethers.parseEther('40')])
 
         expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(0)
       })

From 0de2f4476696eb94e660a0dc36d212c8908019df Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Mon, 16 Sep 2024 19:15:57 -0300
Subject: [PATCH 03/22] fix: Comments from pr

---
 contracts/mocks/MaliciousERC20.sol     | 2 +-
 contracts/mocks/MaliciousRecipient.sol | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/contracts/mocks/MaliciousERC20.sol b/contracts/mocks/MaliciousERC20.sol
index dc655c5..247cb64 100644
--- a/contracts/mocks/MaliciousERC20.sol
+++ b/contracts/mocks/MaliciousERC20.sol
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: CC0-1.0
 
-pragma solidity ^0.8.0;
+pragma solidity ^0.8.9;
 
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 
diff --git a/contracts/mocks/MaliciousRecipient.sol b/contracts/mocks/MaliciousRecipient.sol
index 2bd7ac7..89cf9f6 100644
--- a/contracts/mocks/MaliciousRecipient.sol
+++ b/contracts/mocks/MaliciousRecipient.sol
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: CC0-1.0
 
-// contracts/MaliciousRecipient.sol
-pragma solidity ^0.8.0;
+
+pragma solidity ^0.8.9;
 
 contract MaliciousRecipient {
     // Fallback function that reverts when receiving Ether

From dbd6b603bf32dce49f7e701c9ad48d3a3714b607 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Tue, 17 Sep 2024 13:43:26 -0300
Subject: [PATCH 04/22] fix: Bring branch test coverage to 100%

---
 contracts/ERC20Splitter.sol   |  1 -
 test/SplitterContract.test.ts | 45 ++++++++++++++++++++++++++---------
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index f021ad9..ae4f863 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -74,7 +74,6 @@ contract ERC20Splitter is ReentrancyGuard {
             address tokenAddress = senderTokens[i];
             uint256 amount = balances[tokenAddress][to];
 
-            require(amount > 0, 'ERC20Splitter: Amount to withdraw must be greater than zero');
             balances[tokenAddress][to] = 0;
 
             if (tokenAddress == address(0)) {
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index af97e36..7b0a015 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -2,13 +2,12 @@
 import { ethers, network } from 'hardhat'
 import { loadFixture } from '@nomicfoundation/hardhat-network-helpers'
 import { expect } from 'chai'
-import { MockERC20, ERC20Splitter, MaliciousRecipient, MaliciousERC20 } from '../typechain-types'
+import { MockERC20, ERC20Splitter, MaliciousRecipient } from '../typechain-types'
 import { AddressZero } from '../utils/constants'
 
 describe('ERC20Splitter', () => {
   let splitter: ERC20Splitter
   let mockERC20: MockERC20
-  let maliciousERC20: MaliciousERC20
   let owner: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient1: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient2: Awaited<ReturnType<typeof ethers.getSigner>>
@@ -32,24 +31,19 @@ describe('ERC20Splitter', () => {
     maliciousRecipient = await MaliciousRecipientFactory.deploy()
     await maliciousRecipient.waitForDeployment()
 
-    const MaliciousERC20Factory = await ethers.getContractFactory('MaliciousERC20')
-    maliciousERC20 = await MaliciousERC20Factory.deploy()
-    await maliciousERC20.waitForDeployment()
-
     const mockERC20 = await MockERC20.deploy()
     await mockERC20.waitForDeployment()
 
     const splitter = await ERC20Splitter.deploy()
     await splitter.waitForDeployment()
 
-    return { mockERC20, splitter, maliciousERC20 }
+    return { mockERC20, splitter }
   }
 
   beforeEach(async () => {
     const contracts = await loadFixture(deploySplitterContracts)
     mockERC20 = contracts.mockERC20
     splitter = contracts.splitter
-    maliciousERC20 = contracts.maliciousERC20
 
     // Mint tokens to the owner
     await mockERC20.connect(owner).mint(owner, ethers.parseEther('1000'))
@@ -75,7 +69,7 @@ describe('ERC20Splitter', () => {
     })
 
     const tokenAmount = ethers.parseEther('100')
-    await maliciousERC20.mint(splitter, tokenAmount)
+    await mockERC20.mint(splitter, tokenAmount)
   })
 
   describe('Main Functions', async () => {
@@ -203,13 +197,14 @@ describe('ERC20Splitter', () => {
         )
       })
 
-      it('Should revert when ERC20 transferFrom fails during withdrawal', async () => {
+      it('Should revert when ERC20 transferFrom fails during deposit', async () => {
         const tokenAmount = ethers.parseEther('100')
-        const tokenAddresses = [await maliciousERC20.getAddress()]
+        const tokenAddresses = [await mockERC20.getAddress()]
         const amounts = [tokenAmount]
         const shares = [[10000]] // 100% share
         const recipients = [[recipient1.address]]
 
+        await mockERC20.transferReverts(true, 0)
         await expect(splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients)).to.be.revertedWith(
           'ERC20Splitter: Transfer failed',
         )
@@ -354,6 +349,34 @@ describe('ERC20Splitter', () => {
           'ERC20Splitter: Failed to send Ether',
         )
       })
+      it('Should revert when ERC20 transferFrom fails during withdraw', async () => {
+        const mockERC20false = await mockERC20.getAddress()
+
+        await network.provider.request({
+          method: 'hardhat_impersonateAccount',
+          params: [mockERC20false],
+        })
+
+        const ethAmount = ethers.parseEther('1')
+        const tokenAddresses = [ethers.ZeroAddress] // Ether represented by address zero
+        const amounts = [ethAmount]
+        const shares = [[10000]] // 100% share
+        const recipients = [[recipient1.getAddress()]]
+
+        await splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, {
+          value: ethAmount,
+        })
+
+        await network.provider.send('hardhat_setBalance', [
+          mockERC20false,
+          ethers.toQuantity(ethers.parseEther('1')), // Setting 2 Ether
+        ])
+
+        await mockERC20.transferReverts(true, 0)
+
+        // Attempt to withdraw as the malicious recipient
+        await expect(splitter.connect(recipient1).withdraw()).to.be.revertedWith('ERC20Splitter: TransferFrom failed')
+      })
     })
 
     describe('Withdraw ERC-20 and Native Tokens', async () => {

From 0807f1e5ae1152ff780398c77f7e573882d0a2b4 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Tue, 17 Sep 2024 15:34:39 -0300
Subject: [PATCH 05/22] fix: fixing comments in PR

---
 contracts/ERC20Splitter.sol   | 10 +++---
 test/SplitterContract.test.ts | 61 ++++++++---------------------------
 2 files changed, 17 insertions(+), 54 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index ae4f863..50e1936 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -21,7 +21,7 @@ contract ERC20Splitter is ReentrancyGuard {
         uint16[][] shares,
         address[][] recipients
     );
-    event Withdraw(address indexed user, uint256[] amounts);
+    event Withdraw(address indexed user, address[] tokenAddresses, uint256[] amounts);
 
     uint16 public constant MAX_SHARES = 10000;
 
@@ -62,7 +62,7 @@ contract ERC20Splitter is ReentrancyGuard {
     /// Tokens are automatically determined based on previous deposits.
     function withdraw() external nonReentrant {
         address payable to = payable(msg.sender);
-        address[] storage senderTokens = userTokens[to];
+        address[] storage senderTokens = userTokens[msg.sender];
 
         if (senderTokens.length == 0) {
             return;
@@ -77,8 +77,7 @@ contract ERC20Splitter is ReentrancyGuard {
             balances[tokenAddress][to] = 0;
 
             if (tokenAddress == address(0)) {
-                (bool success, ) = to.call{ value: amount }('');
-                require(success, 'ERC20Splitter: Failed to send Ether');
+                to.transfer(amount);
             } else {
                 require(
                     IERC20(tokenAddress).transferFrom(address(this), to, amount),
@@ -90,10 +89,9 @@ contract ERC20Splitter is ReentrancyGuard {
 
             delete hasToken[to][tokenAddress];
         }
+        emit Withdraw(to, userTokens[msg.sender], withdrawnAmounts);
 
         delete userTokens[to];
-
-        emit Withdraw(to, withdrawnAmounts);
     }
 
     /** Internal Functions **/
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 7b0a015..7c763a4 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -197,19 +197,6 @@ describe('ERC20Splitter', () => {
         )
       })
 
-      it('Should revert when ERC20 transferFrom fails during deposit', async () => {
-        const tokenAmount = ethers.parseEther('100')
-        const tokenAddresses = [await mockERC20.getAddress()]
-        const amounts = [tokenAmount]
-        const shares = [[10000]] // 100% share
-        const recipients = [[recipient1.address]]
-
-        await mockERC20.transferReverts(true, 0)
-        await expect(splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients)).to.be.revertedWith(
-          'ERC20Splitter: Transfer failed',
-        )
-      })
-
       it('Should handle multiple native token (ETH) deposits in a single transaction', async () => {
         const ethShares = [
           [5000, 5000],
@@ -291,9 +278,9 @@ describe('ERC20Splitter', () => {
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
         await expect(splitter.connect(recipient1).withdraw())
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
 
-        expect(await splitter.balances(await mockERC20.getAddress(), recipient1.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
       })
 
       it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
@@ -308,47 +295,18 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
+            [await mockERC20.getAddress(), AddressZero],
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
+        expect(await splitter.balances(await mockERC20.getAddress(), recipient1.address)).to.equal(0)
         expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
       })
 
       it('Should handle withdraw() when user has no tokens', async () => {
         await splitter.connect(anotherUser).withdraw()
       })
 
-      it('Should revert when sending Ether to a recipient fails', async () => {
-        const malicious = await maliciousRecipient.getAddress()
-
-        await network.provider.request({
-          method: 'hardhat_impersonateAccount',
-          params: [malicious],
-        })
-
-        const ethAmount = ethers.parseEther('1')
-        const tokenAddresses = [ethers.ZeroAddress] // Ether represented by address zero
-        const amounts = [ethAmount]
-        const shares = [[10000]] // 100% share
-        const recipients = [[maliciousRecipient.getAddress()]]
-
-        await splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, {
-          value: ethAmount,
-        })
-
-        const maliciousSigner = await ethers.getSigner(malicious)
-
-        await network.provider.send('hardhat_setBalance', [
-          malicious,
-          ethers.toQuantity(ethers.parseEther('1')), // Setting 2 Ether
-        ])
-
-        // Attempt to withdraw as the malicious recipient
-        await expect(splitter.connect(maliciousSigner).withdraw()).to.be.revertedWith(
-          'ERC20Splitter: Failed to send Ether',
-        )
-      })
       it('Should revert when ERC20 transferFrom fails during withdraw', async () => {
         const mockERC20false = await mockERC20.getAddress()
 
@@ -391,7 +349,7 @@ describe('ERC20Splitter', () => {
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
         await expect(splitter.connect(recipient1).withdraw())
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
 
         expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
       })
@@ -408,6 +366,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address, // Expect both ERC-20 and native token
+            [await mockERC20.getAddress(), AddressZero],
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
@@ -431,6 +390,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
+            [AddressZero],
             [ethers.parseEther('0.5')], // Expect 0.5 ETH (50% of 1 ETH)
           )
 
@@ -462,6 +422,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
+            [AddressZero],
             [ethers.parseEther('1')], // Full 1 ETH
           )
 
@@ -471,6 +432,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
+            [await mockERC20.getAddress()],
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
@@ -480,6 +442,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient3.address,
+            [await mockERC20.getAddress()],
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
@@ -520,6 +483,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
+            [AddressZero],
             [ethers.parseEther('0.5')], // 50% of 1 ETH
           )
 
@@ -531,6 +495,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
+            [AddressZero, await mockERC20.getAddress()],
             [ethers.parseEther('0.5'), ethers.parseEther('60')], // 50% of 1 ETH and 60 ERC-20 tokens
           )
 
@@ -541,7 +506,7 @@ describe('ERC20Splitter', () => {
       it('Should allow recipient3 to withdraw only ERC-20 tokens', async () => {
         await expect(splitter.connect(recipient3).withdraw())
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient3.address, [ethers.parseEther('40')])
+          .withArgs(recipient3.address, [await mockERC20.getAddress()], [ethers.parseEther('40')])
 
         expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(0)
       })

From 61bccc92e502a7859653bbcc4f46f0a5c8dba19d Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Tue, 17 Sep 2024 17:18:43 -0300
Subject: [PATCH 06/22] hardhattest

---
 hardhat.config.ts | 1 +
 1 file changed, 1 insertion(+)

diff --git a/hardhat.config.ts b/hardhat.config.ts
index 4a965d7..16a51e0 100644
--- a/hardhat.config.ts
+++ b/hardhat.config.ts
@@ -67,6 +67,7 @@ const BASE_CONFIG = {
       forking: {
         url: POLYGON_PROVIDER_URL,
         blockNumber: 55899875,
+        timeout: 9999900,
       },
     },
   },

From cb7631bb853473129792f9b8e5eb26c17e4235a7 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Wed, 18 Sep 2024 12:19:06 -0300
Subject: [PATCH 07/22] adding event for each recipient

---
 contracts/ERC20Splitter.sol   | 30 ++++++++++++++++---------
 test/SplitterContract.test.ts | 41 ++++++++++++++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index 50e1936..eeabdaf 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -9,7 +9,7 @@ contract ERC20Splitter is ReentrancyGuard {
     mapping(address => mapping(address => uint256)) public balances;
     // userAddress => tokenAddress[]
     mapping(address => address[]) private userTokens;
-    // tokenAddress => boolean
+    // tokenAddress => userAddress => boolean
     mapping(address => mapping(address => bool)) private hasToken;
 
     /** Events **/
@@ -23,6 +23,14 @@ contract ERC20Splitter is ReentrancyGuard {
     );
     event Withdraw(address indexed user, address[] tokenAddresses, uint256[] amounts);
 
+    event RecipientSplit(
+        address indexed depositor,
+        address indexed tokenAddress,
+        address indexed recipient,
+        uint256 amount,
+        uint16 sharePercentage
+    );
+
     uint16 public constant MAX_SHARES = 10000;
 
     /** External Functions **/
@@ -61,8 +69,8 @@ contract ERC20Splitter is ReentrancyGuard {
     /// @notice Withdraw all tokens that the caller is entitled to.
     /// Tokens are automatically determined based on previous deposits.
     function withdraw() external nonReentrant {
-        address payable to = payable(msg.sender);
-        address[] storage senderTokens = userTokens[msg.sender];
+        address recipient = msg.sender;
+        address[] storage senderTokens = userTokens[recipient];
 
         if (senderTokens.length == 0) {
             return;
@@ -72,26 +80,26 @@ contract ERC20Splitter is ReentrancyGuard {
 
         for (uint256 i = 0; i < senderTokens.length; i++) {
             address tokenAddress = senderTokens[i];
-            uint256 amount = balances[tokenAddress][to];
+            uint256 amount = balances[tokenAddress][recipient];
 
-            balances[tokenAddress][to] = 0;
+            balances[tokenAddress][recipient] = 0;
 
             if (tokenAddress == address(0)) {
-                to.transfer(amount);
+                payable(msg.sender).transfer(amount);
             } else {
                 require(
-                    IERC20(tokenAddress).transferFrom(address(this), to, amount),
+                    IERC20(tokenAddress).transferFrom(address(this), recipient, amount),
                     'ERC20Splitter: TransferFrom failed'
                 );
             }
 
             withdrawnAmounts[i] = amount;
 
-            delete hasToken[to][tokenAddress];
+            delete hasToken[recipient][tokenAddress];
         }
-        emit Withdraw(to, userTokens[msg.sender], withdrawnAmounts);
+        emit Withdraw(recipient, userTokens[recipient], withdrawnAmounts);
 
-        delete userTokens[to];
+        delete userTokens[recipient];
     }
 
     /** Internal Functions **/
@@ -130,6 +138,8 @@ contract ERC20Splitter is ReentrancyGuard {
             balances[tokenAddress][recipients[i]] += recipientAmount;
 
             _addTokenForUser(recipients[i], tokenAddress);
+
+            emit RecipientSplit(msg.sender, tokenAddress, recipients[i], recipientAmount, shares[i]);
         }
     }
 
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 7c763a4..ad40501 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -264,6 +264,46 @@ describe('ERC20Splitter', () => {
         // Check balances for recipient3 (40% of 100 ERC-20 tokens = 40 tokens)
         expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('40'))
       })
+      it('Should emit RecipientSplit events for each recipient on deposit', async function () {
+        const mockAddress = await mockERC20.getAddress()
+        const tokenAmount = ethers.parseEther('100')
+        const ethAmount = ethers.parseEther('1')
+
+        const tokenAddresses = [await mockERC20.getAddress(), AddressZero]
+        const amounts = [tokenAmount, ethAmount]
+        const shares = [
+          [5000, 3000, 2000], // For ERC20 token
+          [7000, 2000, 1000], // For ETH
+        ]
+        const recipients = [
+          [recipient1.address, recipient2.address, recipient3.address],
+          [recipient1.address, recipient2.address, recipient3.address],
+        ]
+
+        await expect(splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, { value: ethAmount }))
+          .to.emit(splitter, 'Deposit')
+          .withArgs(owner.address, tokenAddresses, amounts, shares, recipients)
+          .and.to.emit(splitter, 'RecipientSplit')
+          .withArgs(owner.address, mockAddress, recipient1.address, ethers.parseEther('50'), 5000)
+          .and.to.emit(splitter, 'RecipientSplit')
+          .withArgs(owner.address, mockAddress, recipient2.address, ethers.parseEther('30'), 3000)
+          .and.to.emit(splitter, 'RecipientSplit')
+          .withArgs(owner.address, mockAddress, recipient3.address, ethers.parseEther('20'), 2000)
+          .and.to.emit(splitter, 'RecipientSplit')
+          .withArgs(owner.address, AddressZero, recipient1.address, ethers.parseEther('0.7'), 7000)
+          .and.to.emit(splitter, 'RecipientSplit')
+          .withArgs(owner.address, AddressZero, recipient2.address, ethers.parseEther('0.2'), 2000)
+          .and.to.emit(splitter, 'RecipientSplit')
+          .withArgs(owner.address, AddressZero, recipient3.address, ethers.parseEther('0.1'), 1000)
+
+        expect(await splitter.balances(mockAddress, recipient1.address)).to.equal(ethers.parseEther('50'))
+        expect(await splitter.balances(mockAddress, recipient2.address)).to.equal(ethers.parseEther('30'))
+        expect(await splitter.balances(mockAddress, recipient3.address)).to.equal(ethers.parseEther('20'))
+
+        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(ethers.parseEther('0.7'))
+        expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(ethers.parseEther('0.2'))
+        expect(await splitter.balances(AddressZero, recipient3.address)).to.equal(ethers.parseEther('0.1'))
+      })
     })
 
     describe('Withdraw', async () => {
@@ -332,7 +372,6 @@ describe('ERC20Splitter', () => {
 
         await mockERC20.transferReverts(true, 0)
 
-        // Attempt to withdraw as the malicious recipient
         await expect(splitter.connect(recipient1).withdraw()).to.be.revertedWith('ERC20Splitter: TransferFrom failed')
       })
     })

From 20391f1708144c8813d6980983636a27da116320 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Wed, 18 Sep 2024 12:20:50 -0300
Subject: [PATCH 08/22] adding event for each recipient

---
 contracts/ERC20Splitter.sol | 1 +
 1 file changed, 1 insertion(+)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index eeabdaf..4fc3a38 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -21,6 +21,7 @@ contract ERC20Splitter is ReentrancyGuard {
         uint16[][] shares,
         address[][] recipients
     );
+
     event Withdraw(address indexed user, address[] tokenAddresses, uint256[] amounts);
 
     event RecipientSplit(

From 71ebf221635968ed72a2a456419dd88d927f6286 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Wed, 18 Sep 2024 12:38:12 -0300
Subject: [PATCH 09/22] Remove contracts/mocks/MaliciousERC20.sol from the PR

---
 contracts/mocks/MaliciousERC20.sol | 22 ----------------------
 1 file changed, 22 deletions(-)
 delete mode 100644 contracts/mocks/MaliciousERC20.sol

diff --git a/contracts/mocks/MaliciousERC20.sol b/contracts/mocks/MaliciousERC20.sol
deleted file mode 100644
index 247cb64..0000000
--- a/contracts/mocks/MaliciousERC20.sol
+++ /dev/null
@@ -1,22 +0,0 @@
-// SPDX-License-Identifier: CC0-1.0
-
-pragma solidity ^0.8.9;
-
-import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
-
-contract MaliciousERC20 is ERC20 {
-    constructor() ERC20("MaliciousToken", "MTK") {}
-
-    function transferFrom(
-        address sender,
-        address recipient,
-        uint256 amount
-    ) public override returns (bool) {
-        return false;
-    }
-
-    // Mint function for testing purposes
-    function mint(address to, uint256 amount) external {
-        _mint(to, amount);
-    }
-}

From 17cb4c125a3917fa669454c9e26a9e22dc8ca3ce Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Wed, 18 Sep 2024 17:47:10 -0300
Subject: [PATCH 10/22] reducing gas cost

---
 contracts/ERC20Splitter.sol   |  28 +--------
 test/SplitterContract.test.ts | 103 +++++++++++++++++++++++++++++++++-
 2 files changed, 102 insertions(+), 29 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index 4fc3a38..d58bb87 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -9,13 +9,11 @@ contract ERC20Splitter is ReentrancyGuard {
     mapping(address => mapping(address => uint256)) public balances;
     // userAddress => tokenAddress[]
     mapping(address => address[]) private userTokens;
-    // tokenAddress => userAddress => boolean
-    mapping(address => mapping(address => bool)) private hasToken;
 
     /** Events **/
 
     event Deposit(
-        address indexed depositor,
+        address indexed user,
         address[] tokenAddresses,
         uint256[] amounts,
         uint16[][] shares,
@@ -24,14 +22,6 @@ contract ERC20Splitter is ReentrancyGuard {
 
     event Withdraw(address indexed user, address[] tokenAddresses, uint256[] amounts);
 
-    event RecipientSplit(
-        address indexed depositor,
-        address indexed tokenAddress,
-        address indexed recipient,
-        uint256 amount,
-        uint16 sharePercentage
-    );
-
     uint16 public constant MAX_SHARES = 10000;
 
     /** External Functions **/
@@ -95,8 +85,6 @@ contract ERC20Splitter is ReentrancyGuard {
             }
 
             withdrawnAmounts[i] = amount;
-
-            delete hasToken[recipient][tokenAddress];
         }
         emit Withdraw(recipient, userTokens[recipient], withdrawnAmounts);
 
@@ -137,20 +125,6 @@ contract ERC20Splitter is ReentrancyGuard {
         for (uint256 i = 0; i < recipients.length; i++) {
             uint256 recipientAmount = (amount * shares[i]) / MAX_SHARES;
             balances[tokenAddress][recipients[i]] += recipientAmount;
-
-            _addTokenForUser(recipients[i], tokenAddress);
-
-            emit RecipientSplit(msg.sender, tokenAddress, recipients[i], recipientAmount, shares[i]);
-        }
-    }
-
-    /// @notice Adds a token to the list of tokens a user has received (for automatic withdrawals).
-    /// @param recipient The recipient of the token.
-    /// @param tokenAddress The address of the token.
-    function _addTokenForUser(address recipient, address tokenAddress) internal {
-        if (!hasToken[recipient][tokenAddress]) {
-            userTokens[recipient].push(tokenAddress);
-            hasToken[recipient][tokenAddress] = true;
         }
     }
 }
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index ad40501..069f89d 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -8,10 +8,14 @@ import { AddressZero } from '../utils/constants'
 describe('ERC20Splitter', () => {
   let splitter: ERC20Splitter
   let mockERC20: MockERC20
+  let mockERC20_2: MockERC20
+  let mockERC20_3: MockERC20
+  let mockERC20_4: MockERC20
   let owner: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient1: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient2: Awaited<ReturnType<typeof ethers.getSigner>>
   let recipient3: Awaited<ReturnType<typeof ethers.getSigner>>
+  let recipient4: Awaited<ReturnType<typeof ethers.getSigner>>
   let anotherUser: Awaited<ReturnType<typeof ethers.getSigner>>
   let maliciousRecipient: MaliciousRecipient
 
@@ -20,11 +24,15 @@ describe('ERC20Splitter', () => {
 
   before(async function () {
     // prettier-ignore
-    [owner, recipient1, recipient2, recipient3, anotherUser] = await ethers.getSigners()
+    [owner, recipient1, recipient2, recipient3,recipient4, anotherUser] = await ethers.getSigners()
   })
 
   async function deploySplitterContracts() {
     const MockERC20 = await ethers.getContractFactory('MockERC20')
+    const MockERC20_2 = await ethers.getContractFactory('MockERC20')
+    const MockERC20_3 = await ethers.getContractFactory('MockERC20')
+    const MockERC20_4 = await ethers.getContractFactory('MockERC20')
+
     const ERC20Splitter = await ethers.getContractFactory('ERC20Splitter')
 
     const MaliciousRecipientFactory = await ethers.getContractFactory('MaliciousRecipient')
@@ -34,19 +42,34 @@ describe('ERC20Splitter', () => {
     const mockERC20 = await MockERC20.deploy()
     await mockERC20.waitForDeployment()
 
+    const mockERC20_2 = await MockERC20_2.deploy()
+    await mockERC20_2.waitForDeployment()
+
+    const mockERC20_3 = await MockERC20_3.deploy()
+    await mockERC20_3.waitForDeployment()
+
+    const mockERC20_4 = await MockERC20_4.deploy()
+    await mockERC20_4.waitForDeployment()
+
     const splitter = await ERC20Splitter.deploy()
     await splitter.waitForDeployment()
 
-    return { mockERC20, splitter }
+    return { mockERC20, mockERC20_2, mockERC20_3, mockERC20_4, splitter }
   }
 
   beforeEach(async () => {
     const contracts = await loadFixture(deploySplitterContracts)
     mockERC20 = contracts.mockERC20
+    mockERC20_2 = contracts.mockERC20_2
+    mockERC20_3 = contracts.mockERC20_3
+    mockERC20_4 = contracts.mockERC20_4
     splitter = contracts.splitter
 
     // Mint tokens to the owner
     await mockERC20.connect(owner).mint(owner, ethers.parseEther('1000'))
+    await mockERC20_2.connect(owner).mint(owner, ethers.parseEther('1000'))
+    await mockERC20_3.connect(owner).mint(owner, ethers.parseEther('1000'))
+    await mockERC20_4.connect(owner).mint(owner, ethers.parseEther('1000'))
 
     const splitterAddress = await splitter.getAddress()
 
@@ -62,6 +85,9 @@ describe('ERC20Splitter', () => {
     const splitterSigner = await ethers.getSigner(splitterAddress)
 
     await mockERC20.connect(splitterSigner).approve(splitterAddress, ethers.MaxUint256)
+    await mockERC20_2.connect(splitterSigner).approve(splitterAddress, ethers.MaxUint256)
+    await mockERC20_3.connect(splitterSigner).approve(splitterAddress, ethers.MaxUint256)
+    await mockERC20_4.connect(splitterSigner).approve(splitterAddress, ethers.MaxUint256)
 
     await network.provider.request({
       method: 'hardhat_stopImpersonatingAccount',
@@ -70,12 +96,85 @@ describe('ERC20Splitter', () => {
 
     const tokenAmount = ethers.parseEther('100')
     await mockERC20.mint(splitter, tokenAmount)
+    await mockERC20_2.mint(splitter, tokenAmount)
+    await mockERC20_3.mint(splitter, tokenAmount)
+    await mockERC20_4.mint(splitter, tokenAmount)
   })
 
   describe('Main Functions', async () => {
     describe('Deposit', async () => {
       beforeEach(async () => {
         await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
+        await mockERC20_2.connect(owner).approve(splitter.getAddress(), tokenAmount)
+        await mockERC20_3.connect(owner).approve(splitter.getAddress(), tokenAmount)
+        await mockERC20_4.connect(owner).approve(splitter.getAddress(), tokenAmount)
+      })
+
+      it('Should deposit ERC20 tokens for one recipient', async () => {
+        const shares = [[10000]] // 50%, 30%, 20%
+        const recipients = [[recipient1.address]]
+
+        await expect(
+          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], shares, recipients),
+        ).to.emit(splitter, 'Deposit')
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('100'))
+      })
+
+      it('Should deposit four ERC20 tokens and split them between recipients', async () => {
+        const tokenAmounts = [
+          ethers.parseEther('100'),
+          ethers.parseEther('100'),
+          ethers.parseEther('100'),
+          ethers.parseEther('100'),
+        ]
+        const shares = [[10000], [10000], [10000], [10000]]
+        const recipients = [[recipient1.address], [recipient2.address], [recipient3.address], [recipient4.address]]
+
+        await expect(
+          splitter
+            .connect(owner)
+            .deposit(
+              [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
+              tokenAmounts,
+              shares,
+              recipients,
+            ),
+        ).to.emit(splitter, 'Deposit')
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('100'))
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(ethers.parseEther('100'))
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient3.address)).to.equal(ethers.parseEther('100'))
+        expect(await splitter.balances(mockERC20_4.getAddress(), recipient4.address)).to.equal(ethers.parseEther('100'))
+      })
+
+      it.only('Should deposit three ERC20 tokens and split them between recipients', async () => {
+        const tokenAmounts = [ethers.parseEther('100'), ethers.parseEther('100'), ethers.parseEther('100')]
+        const shares = [
+          [5000, 3000, 2000],
+          [5000, 3000, 2000],
+          [5000, 3000, 2000],
+        ]
+        const recipients = [
+          [recipient1.address, recipient2.address, recipient3.address],
+          [recipient1.address, recipient2.address, recipient3.address],
+          [recipient1.address, recipient2.address, recipient3.address],
+        ]
+
+        await expect(
+          splitter
+            .connect(owner)
+            .deposit(
+              [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress()],
+              tokenAmounts,
+              shares,
+              recipients,
+            ),
+        ).to.emit(splitter, 'Deposit')
+
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('50'))
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(ethers.parseEther('30'))
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient3.address)).to.equal(ethers.parseEther('20'))
       })
 
       it('Should deposit ERC20 tokens and split them between recipients', async () => {

From 21cd544a410f8e563b0d42e043274f8892839370 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Wed, 18 Sep 2024 18:45:25 -0300
Subject: [PATCH 11/22] feat: gas

---
 contracts/ERC20Splitter.sol   |   2 +
 test/SplitterContract.test.ts | 133 +++++++++++++++++++++++-----------
 2 files changed, 94 insertions(+), 41 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index d58bb87..d593292 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -86,6 +86,7 @@ contract ERC20Splitter is ReentrancyGuard {
 
             withdrawnAmounts[i] = amount;
         }
+
         emit Withdraw(recipient, userTokens[recipient], withdrawnAmounts);
 
         delete userTokens[recipient];
@@ -125,6 +126,7 @@ contract ERC20Splitter is ReentrancyGuard {
         for (uint256 i = 0; i < recipients.length; i++) {
             uint256 recipientAmount = (amount * shares[i]) / MAX_SHARES;
             balances[tokenAddress][recipients[i]] += recipientAmount;
+            userTokens[recipients[i]].push(tokenAddress);
         }
     }
 }
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 069f89d..00738ed 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -148,7 +148,98 @@ describe('ERC20Splitter', () => {
         expect(await splitter.balances(mockERC20_4.getAddress(), recipient4.address)).to.equal(ethers.parseEther('100'))
       })
 
-      it.only('Should deposit three ERC20 tokens and split them between recipients', async () => {
+      it.only('Should deposit four ERC20 tokens and split them between recipients', async () => {
+        const tokenAmounts = [
+          ethers.parseEther('100'), // Amount for token 1
+          ethers.parseEther('60'), // Amount for token 2
+          ethers.parseEther('40'), // Amount for token 3
+          ethers.parseEther('80'), // Amount for token 4
+        ]
+
+        const shares = [
+          [5000, 3000, 2000], // Shares for token 1 (50%, 30%, 20%)
+          [5000, 3000, 2000], // Shares for token 2 (50%, 30%, 20%)
+          [5000, 3000, 2000], // Shares for token 3 (50%, 30%, 20%)
+          [5000, 3000, 2000], // Shares for token 4 (50%, 30%, 20%)
+        ]
+
+        const recipients = [
+          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 1
+          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 2
+          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 3
+          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 4
+        ]
+
+        // Call the deposit function with matching array lengths
+        await expect(
+          splitter.connect(owner).deposit(
+            [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
+            tokenAmounts, // The array of amounts for the 4 tokens
+            shares, // Provide separate shares for each token
+            recipients, // Recipients array matches in length with token addresses
+          ),
+        ).to.emit(splitter, 'Deposit')
+
+        // Expected balances based on deposit amounts and shares
+        const expectedRecipient1BalanceToken1 = ethers.parseEther('50') // 50% of 100 tokens (token 1)
+        const expectedRecipient2BalanceToken1 = ethers.parseEther('30') // 30% of 100 tokens (token 1)
+        const expectedRecipient3BalanceToken1 = ethers.parseEther('20') // 20% of 100 tokens (token 1)
+
+        const expectedRecipient1BalanceToken2 = ethers.parseEther('30') // 50% of 60 tokens (token 2)
+        const expectedRecipient2BalanceToken2 = ethers.parseEther('18') // 30% of 60 tokens (token 2)
+        const expectedRecipient3BalanceToken2 = ethers.parseEther('12') // 20% of 60 tokens (token 2)
+
+        const expectedRecipient1BalanceToken3 = ethers.parseEther('20') // 50% of 40 tokens (token 3)
+        const expectedRecipient2BalanceToken3 = ethers.parseEther('12') // 30% of 40 tokens (token 3)
+        const expectedRecipient3BalanceToken3 = ethers.parseEther('8') // 20% of 40 tokens (token 3)
+
+        const expectedRecipient1BalanceToken4 = ethers.parseEther('40') // 50% of 80 tokens (token 4)
+        const expectedRecipient2BalanceToken4 = ethers.parseEther('24') // 30% of 80 tokens (token 4)
+        const expectedRecipient3BalanceToken4 = ethers.parseEther('16') // 20% of 80 tokens (token 4)
+
+        // Checking balances after splitting for all tokens
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(
+          expectedRecipient1BalanceToken1,
+        )
+        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(
+          expectedRecipient2BalanceToken1,
+        )
+        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(
+          expectedRecipient3BalanceToken1,
+        )
+
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient1.address)).to.equal(
+          expectedRecipient1BalanceToken2,
+        )
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(
+          expectedRecipient2BalanceToken2,
+        )
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient3.address)).to.equal(
+          expectedRecipient3BalanceToken2,
+        )
+
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient1.address)).to.equal(
+          expectedRecipient1BalanceToken3,
+        )
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient2.address)).to.equal(
+          expectedRecipient2BalanceToken3,
+        )
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient3.address)).to.equal(
+          expectedRecipient3BalanceToken3,
+        )
+
+        expect(await splitter.balances(mockERC20_4.getAddress(), recipient1.address)).to.equal(
+          expectedRecipient1BalanceToken4,
+        )
+        expect(await splitter.balances(mockERC20_4.getAddress(), recipient2.address)).to.equal(
+          expectedRecipient2BalanceToken4,
+        )
+        expect(await splitter.balances(mockERC20_4.getAddress(), recipient3.address)).to.equal(
+          expectedRecipient3BalanceToken4,
+        )
+      })
+
+      it('Should deposit three ERC20 tokens and split them between recipients', async () => {
         const tokenAmounts = [ethers.parseEther('100'), ethers.parseEther('100'), ethers.parseEther('100')]
         const shares = [
           [5000, 3000, 2000],
@@ -363,46 +454,6 @@ describe('ERC20Splitter', () => {
         // Check balances for recipient3 (40% of 100 ERC-20 tokens = 40 tokens)
         expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('40'))
       })
-      it('Should emit RecipientSplit events for each recipient on deposit', async function () {
-        const mockAddress = await mockERC20.getAddress()
-        const tokenAmount = ethers.parseEther('100')
-        const ethAmount = ethers.parseEther('1')
-
-        const tokenAddresses = [await mockERC20.getAddress(), AddressZero]
-        const amounts = [tokenAmount, ethAmount]
-        const shares = [
-          [5000, 3000, 2000], // For ERC20 token
-          [7000, 2000, 1000], // For ETH
-        ]
-        const recipients = [
-          [recipient1.address, recipient2.address, recipient3.address],
-          [recipient1.address, recipient2.address, recipient3.address],
-        ]
-
-        await expect(splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, { value: ethAmount }))
-          .to.emit(splitter, 'Deposit')
-          .withArgs(owner.address, tokenAddresses, amounts, shares, recipients)
-          .and.to.emit(splitter, 'RecipientSplit')
-          .withArgs(owner.address, mockAddress, recipient1.address, ethers.parseEther('50'), 5000)
-          .and.to.emit(splitter, 'RecipientSplit')
-          .withArgs(owner.address, mockAddress, recipient2.address, ethers.parseEther('30'), 3000)
-          .and.to.emit(splitter, 'RecipientSplit')
-          .withArgs(owner.address, mockAddress, recipient3.address, ethers.parseEther('20'), 2000)
-          .and.to.emit(splitter, 'RecipientSplit')
-          .withArgs(owner.address, AddressZero, recipient1.address, ethers.parseEther('0.7'), 7000)
-          .and.to.emit(splitter, 'RecipientSplit')
-          .withArgs(owner.address, AddressZero, recipient2.address, ethers.parseEther('0.2'), 2000)
-          .and.to.emit(splitter, 'RecipientSplit')
-          .withArgs(owner.address, AddressZero, recipient3.address, ethers.parseEther('0.1'), 1000)
-
-        expect(await splitter.balances(mockAddress, recipient1.address)).to.equal(ethers.parseEther('50'))
-        expect(await splitter.balances(mockAddress, recipient2.address)).to.equal(ethers.parseEther('30'))
-        expect(await splitter.balances(mockAddress, recipient3.address)).to.equal(ethers.parseEther('20'))
-
-        expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(ethers.parseEther('0.7'))
-        expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(ethers.parseEther('0.2'))
-        expect(await splitter.balances(AddressZero, recipient3.address)).to.equal(ethers.parseEther('0.1'))
-      })
     })
 
     describe('Withdraw', async () => {

From 4c239e127dd9641bc09e67d24c0e5639dc6d1194 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Wed, 18 Sep 2024 19:27:10 -0300
Subject: [PATCH 12/22] fix: adding push to the storage again

---
 test/SplitterContract.test.ts | 107 ++++++++++------------------------
 1 file changed, 32 insertions(+), 75 deletions(-)

diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 00738ed..9f92922 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -150,93 +150,50 @@ describe('ERC20Splitter', () => {
 
       it.only('Should deposit four ERC20 tokens and split them between recipients', async () => {
         const tokenAmounts = [
-          ethers.parseEther('100'), // Amount for token 1
-          ethers.parseEther('60'), // Amount for token 2
-          ethers.parseEther('40'), // Amount for token 3
-          ethers.parseEther('80'), // Amount for token 4
+          ethers.parseEther('100'),
+          ethers.parseEther('60'),
+          ethers.parseEther('40'),
+          ethers.parseEther('80'),
         ]
 
         const shares = [
-          [5000, 3000, 2000], // Shares for token 1 (50%, 30%, 20%)
-          [5000, 3000, 2000], // Shares for token 2 (50%, 30%, 20%)
-          [5000, 3000, 2000], // Shares for token 3 (50%, 30%, 20%)
-          [5000, 3000, 2000], // Shares for token 4 (50%, 30%, 20%)
+          [5000, 3000, 2000],
+          [5000, 3000, 2000],
+          [5000, 3000, 2000],
+          [5000, 3000, 2000],
         ]
 
         const recipients = [
-          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 1
-          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 2
-          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 3
-          [recipient1.address, recipient2.address, recipient3.address], // Recipients for token 4
+          [recipient1.address, recipient2.address, recipient3.address],
+          [recipient1.address, recipient2.address, recipient3.address],
+          [recipient1.address, recipient2.address, recipient3.address],
+          [recipient1.address, recipient2.address, recipient3.address],
         ]
 
-        // Call the deposit function with matching array lengths
         await expect(
-          splitter.connect(owner).deposit(
-            [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
-            tokenAmounts, // The array of amounts for the 4 tokens
-            shares, // Provide separate shares for each token
-            recipients, // Recipients array matches in length with token addresses
-          ),
+          splitter
+            .connect(owner)
+            .deposit(
+              [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
+              tokenAmounts,
+              shares,
+              recipients,
+            ),
         ).to.emit(splitter, 'Deposit')
 
-        // Expected balances based on deposit amounts and shares
-        const expectedRecipient1BalanceToken1 = ethers.parseEther('50') // 50% of 100 tokens (token 1)
-        const expectedRecipient2BalanceToken1 = ethers.parseEther('30') // 30% of 100 tokens (token 1)
-        const expectedRecipient3BalanceToken1 = ethers.parseEther('20') // 20% of 100 tokens (token 1)
-
-        const expectedRecipient1BalanceToken2 = ethers.parseEther('30') // 50% of 60 tokens (token 2)
-        const expectedRecipient2BalanceToken2 = ethers.parseEther('18') // 30% of 60 tokens (token 2)
-        const expectedRecipient3BalanceToken2 = ethers.parseEther('12') // 20% of 60 tokens (token 2)
-
-        const expectedRecipient1BalanceToken3 = ethers.parseEther('20') // 50% of 40 tokens (token 3)
-        const expectedRecipient2BalanceToken3 = ethers.parseEther('12') // 30% of 40 tokens (token 3)
-        const expectedRecipient3BalanceToken3 = ethers.parseEther('8') // 20% of 40 tokens (token 3)
-
-        const expectedRecipient1BalanceToken4 = ethers.parseEther('40') // 50% of 80 tokens (token 4)
-        const expectedRecipient2BalanceToken4 = ethers.parseEther('24') // 30% of 80 tokens (token 4)
-        const expectedRecipient3BalanceToken4 = ethers.parseEther('16') // 20% of 80 tokens (token 4)
-
-        // Checking balances after splitting for all tokens
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(
-          expectedRecipient1BalanceToken1,
-        )
-        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(
-          expectedRecipient2BalanceToken1,
-        )
-        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(
-          expectedRecipient3BalanceToken1,
-        )
-
-        expect(await splitter.balances(mockERC20_2.getAddress(), recipient1.address)).to.equal(
-          expectedRecipient1BalanceToken2,
-        )
-        expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(
-          expectedRecipient2BalanceToken2,
-        )
-        expect(await splitter.balances(mockERC20_2.getAddress(), recipient3.address)).to.equal(
-          expectedRecipient3BalanceToken2,
-        )
-
-        expect(await splitter.balances(mockERC20_3.getAddress(), recipient1.address)).to.equal(
-          expectedRecipient1BalanceToken3,
-        )
-        expect(await splitter.balances(mockERC20_3.getAddress(), recipient2.address)).to.equal(
-          expectedRecipient2BalanceToken3,
-        )
-        expect(await splitter.balances(mockERC20_3.getAddress(), recipient3.address)).to.equal(
-          expectedRecipient3BalanceToken3,
-        )
+        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('50'))
+        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(ethers.parseEther('30'))
+        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('20'))
 
-        expect(await splitter.balances(mockERC20_4.getAddress(), recipient1.address)).to.equal(
-          expectedRecipient1BalanceToken4,
-        )
-        expect(await splitter.balances(mockERC20_4.getAddress(), recipient2.address)).to.equal(
-          expectedRecipient2BalanceToken4,
-        )
-        expect(await splitter.balances(mockERC20_4.getAddress(), recipient3.address)).to.equal(
-          expectedRecipient3BalanceToken4,
-        )
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient1.address)).to.equal(ethers.parseEther('30'))
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(ethers.parseEther('18'))
+        expect(await splitter.balances(mockERC20_2.getAddress(), recipient3.address)).to.equal(ethers.parseEther('12'))
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient1.address)).to.equal(ethers.parseEther('20'))
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient2.address)).to.equal(ethers.parseEther('12'))
+        expect(await splitter.balances(mockERC20_3.getAddress(), recipient3.address)).to.equal(ethers.parseEther('8'))
+        expect(await splitter.balances(mockERC20_4.getAddress(), recipient1.address)).to.equal(ethers.parseEther('40'))
+        expect(await splitter.balances(mockERC20_4.getAddress(), recipient2.address)).to.equal(ethers.parseEther('24'))
+        expect(await splitter.balances(mockERC20_4.getAddress(), recipient3.address)).to.equal(ethers.parseEther('16'))
       })
 
       it('Should deposit three ERC20 tokens and split them between recipients', async () => {

From 432328e200eba7ab255bb0e0b40574a7475ce9c5 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Thu, 19 Sep 2024 10:08:06 -0300
Subject: [PATCH 13/22] fix: comments from PR

---
 contracts/ERC20Splitter.sol   | 13 ++++++-------
 test/SplitterContract.test.ts |  2 +-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index d593292..894bb46 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -60,8 +60,7 @@ contract ERC20Splitter is ReentrancyGuard {
     /// @notice Withdraw all tokens that the caller is entitled to.
     /// Tokens are automatically determined based on previous deposits.
     function withdraw() external nonReentrant {
-        address recipient = msg.sender;
-        address[] storage senderTokens = userTokens[recipient];
+        address[] storage senderTokens = userTokens[msg.sender];
 
         if (senderTokens.length == 0) {
             return;
@@ -71,15 +70,15 @@ contract ERC20Splitter is ReentrancyGuard {
 
         for (uint256 i = 0; i < senderTokens.length; i++) {
             address tokenAddress = senderTokens[i];
-            uint256 amount = balances[tokenAddress][recipient];
+            uint256 amount = balances[tokenAddress][msg.sender];
 
-            balances[tokenAddress][recipient] = 0;
+            delete balances[tokenAddress][msg.sender];
 
             if (tokenAddress == address(0)) {
                 payable(msg.sender).transfer(amount);
             } else {
                 require(
-                    IERC20(tokenAddress).transferFrom(address(this), recipient, amount),
+                    IERC20(tokenAddress).transferFrom(address(this), msg.sender, amount),
                     'ERC20Splitter: TransferFrom failed'
                 );
             }
@@ -87,9 +86,9 @@ contract ERC20Splitter is ReentrancyGuard {
             withdrawnAmounts[i] = amount;
         }
 
-        emit Withdraw(recipient, userTokens[recipient], withdrawnAmounts);
+        emit Withdraw(msg.sender, userTokens[msg.sender], withdrawnAmounts);
 
-        delete userTokens[recipient];
+        delete userTokens[msg.sender];
     }
 
     /** Internal Functions **/
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 9f92922..d283175 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -148,7 +148,7 @@ describe('ERC20Splitter', () => {
         expect(await splitter.balances(mockERC20_4.getAddress(), recipient4.address)).to.equal(ethers.parseEther('100'))
       })
 
-      it.only('Should deposit four ERC20 tokens and split them between recipients', async () => {
+      it('Should deposit four ERC20 tokens and split them between recipients', async () => {
         const tokenAmounts = [
           ethers.parseEther('100'),
           ethers.parseEther('60'),

From f489d77d7ee1dd2532c20dbce4a57f1991d63f60 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Thu, 19 Sep 2024 14:08:45 -0300
Subject: [PATCH 14/22] feat: withdraw with tokenAddresses as parameter

---
 contracts/ERC20Splitter.sol | 30 +++++++++++++-----------------
 1 file changed, 13 insertions(+), 17 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index 894bb46..7531d69 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -7,8 +7,6 @@ import '@openzeppelin/contracts/security/ReentrancyGuard.sol';
 contract ERC20Splitter is ReentrancyGuard {
     // tokenAddress => userAddress => balance
     mapping(address => mapping(address => uint256)) public balances;
-    // userAddress => tokenAddress[]
-    mapping(address => address[]) private userTokens;
 
     /** Events **/
 
@@ -59,36 +57,35 @@ contract ERC20Splitter is ReentrancyGuard {
 
     /// @notice Withdraw all tokens that the caller is entitled to.
     /// Tokens are automatically determined based on previous deposits.
-    function withdraw() external nonReentrant {
-        address[] storage senderTokens = userTokens[msg.sender];
+      function withdraw(address[] calldata tokenAddresses) external nonReentrant {
+        uint256 tokenCount = tokenAddresses.length;
+        require(tokenCount > 0, 'ERC20Splitter: No tokens specified');
 
-        if (senderTokens.length == 0) {
-            return;
-        }
-
-        uint256[] memory withdrawnAmounts = new uint256[](senderTokens.length);
+        uint256[] memory withdrawnAmounts = new uint256[](tokenCount);
 
-        for (uint256 i = 0; i < senderTokens.length; i++) {
-            address tokenAddress = senderTokens[i];
+        for (uint256 i = 0; i < tokenCount; i++) {
+            address tokenAddress = tokenAddresses[i];
             uint256 amount = balances[tokenAddress][msg.sender];
 
-            delete balances[tokenAddress][msg.sender];
+            if (amount == 0) {
+                continue; // Skip if no balance
+            }
+
+            balances[tokenAddress][msg.sender] = 0;
 
             if (tokenAddress == address(0)) {
                 payable(msg.sender).transfer(amount);
             } else {
                 require(
                     IERC20(tokenAddress).transferFrom(address(this), msg.sender, amount),
-                    'ERC20Splitter: TransferFrom failed'
+                    'ERC20Splitter: Transfer failed'
                 );
             }
 
             withdrawnAmounts[i] = amount;
         }
 
-        emit Withdraw(msg.sender, userTokens[msg.sender], withdrawnAmounts);
-
-        delete userTokens[msg.sender];
+        emit Withdraw(msg.sender, tokenAddresses, withdrawnAmounts);
     }
 
     /** Internal Functions **/
@@ -125,7 +122,6 @@ contract ERC20Splitter is ReentrancyGuard {
         for (uint256 i = 0; i < recipients.length; i++) {
             uint256 recipientAmount = (amount * shares[i]) / MAX_SHARES;
             balances[tokenAddress][recipients[i]] += recipientAmount;
-            userTokens[recipients[i]].push(tokenAddress);
         }
     }
 }

From 50305ff91289ffeeda140aff65da4b2a5717ee6a Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Thu, 19 Sep 2024 14:10:26 -0300
Subject: [PATCH 15/22] feat: withdraw with tokenAddresses as parameter

---
 contracts/ERC20Splitter.sol | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index 7531d69..609f9eb 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -57,7 +57,7 @@ contract ERC20Splitter is ReentrancyGuard {
 
     /// @notice Withdraw all tokens that the caller is entitled to.
     /// Tokens are automatically determined based on previous deposits.
-      function withdraw(address[] calldata tokenAddresses) external nonReentrant {
+    function withdraw(address[] calldata tokenAddresses) external nonReentrant {
         uint256 tokenCount = tokenAddresses.length;
         require(tokenCount > 0, 'ERC20Splitter: No tokens specified');
 
@@ -68,10 +68,10 @@ contract ERC20Splitter is ReentrancyGuard {
             uint256 amount = balances[tokenAddress][msg.sender];
 
             if (amount == 0) {
-                continue; // Skip if no balance
+                continue;
             }
 
-            balances[tokenAddress][msg.sender] = 0;
+            delete balances[tokenAddress][msg.sender];
 
             if (tokenAddress == address(0)) {
                 payable(msg.sender).transfer(amount);

From 31cb4fd8779436ccb83cf58256125b707c1cd6c0 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Thu, 19 Sep 2024 16:13:17 -0300
Subject: [PATCH 16/22] feat: pushing tests to cover withdraw with parameter

---
 test/SplitterContract.test.ts | 38 +++++++++++++++++++++++------------
 1 file changed, 25 insertions(+), 13 deletions(-)

diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index d283175..1b4a4e2 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -423,7 +423,8 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
-        await expect(splitter.connect(recipient1).withdraw())
+        const tokens = [await mockERC20.getAddress()]
+        await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
 
@@ -431,6 +432,7 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
+        const tokens = [await mockERC20.getAddress(), AddressZero]
         const shares = [[5000, 3000, 2000]]
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
@@ -438,7 +440,7 @@ describe('ERC20Splitter', () => {
           value: ethAmount,
         })
 
-        await expect(splitter.connect(recipient1).withdraw())
+        await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
@@ -451,10 +453,12 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should handle withdraw() when user has no tokens', async () => {
-        await splitter.connect(anotherUser).withdraw()
+        const tokens = [await mockERC20.getAddress(), AddressZero]
+        await splitter.connect(anotherUser).withdraw(tokens)
       })
 
       it('Should revert when ERC20 transferFrom fails during withdraw', async () => {
+        const tokens = [await mockERC20.getAddress()]
         const mockERC20false = await mockERC20.getAddress()
 
         await network.provider.request({
@@ -479,7 +483,7 @@ describe('ERC20Splitter', () => {
 
         await mockERC20.transferReverts(true, 0)
 
-        await expect(splitter.connect(recipient1).withdraw()).to.be.revertedWith('ERC20Splitter: TransferFrom failed')
+        await expect(splitter.connect(recipient1).withdraw(tokens)).to.be.revertedWith('ERC20Splitter: Transfer failed')
       })
     })
 
@@ -493,7 +497,8 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
-        await expect(splitter.connect(recipient1).withdraw())
+        const tokens = [await mockERC20.getAddress()]
+        await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
 
@@ -501,6 +506,7 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
+        const tokens = [await mockERC20.getAddress(), AddressZero]
         const shares = [[5000, 3000, 2000]]
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
@@ -508,7 +514,7 @@ describe('ERC20Splitter', () => {
           value: ethAmount,
         })
 
-        await expect(splitter.connect(recipient1).withdraw())
+        await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address, // Expect both ERC-20 and native token
@@ -532,7 +538,8 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow a recipient to withdraw only their split native tokens (ETH)', async () => {
-        await expect(splitter.connect(recipient1).withdraw())
+        const tokens = [AddressZero]
+        await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
@@ -564,7 +571,9 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow recipient1 to withdraw only their ETH and other recipients to withdraw their ERC-20 tokens', async () => {
-        await expect(splitter.connect(recipient1).withdraw())
+        const tokenEth = [AddressZero]
+        const tokenErc20 = [await mockERC20.getAddress()]
+        await expect(splitter.connect(recipient1).withdraw(tokenEth))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
@@ -574,7 +583,7 @@ describe('ERC20Splitter', () => {
 
         expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
 
-        await expect(splitter.connect(recipient2).withdraw())
+        await expect(splitter.connect(recipient2).withdraw(tokenErc20))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
@@ -584,7 +593,7 @@ describe('ERC20Splitter', () => {
 
         expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(0)
 
-        await expect(splitter.connect(recipient3).withdraw())
+        await expect(splitter.connect(recipient3).withdraw(tokenErc20))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient3.address,
@@ -625,7 +634,8 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow recipient1 to withdraw only ETH', async () => {
-        await expect(splitter.connect(recipient1).withdraw())
+        const tokens = [AddressZero]
+        await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
@@ -637,7 +647,8 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow recipient2 to withdraw both ETH and ERC-20 tokens', async () => {
-        await expect(splitter.connect(recipient2).withdraw())
+        const tokens = [AddressZero, await mockERC20.getAddress()]
+        await expect(splitter.connect(recipient2).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
@@ -650,7 +661,8 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow recipient3 to withdraw only ERC-20 tokens', async () => {
-        await expect(splitter.connect(recipient3).withdraw())
+        const tokens = [await mockERC20.getAddress()]
+        await expect(splitter.connect(recipient3).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(recipient3.address, [await mockERC20.getAddress()], [ethers.parseEther('40')])
 

From f91f56d8384232d69831ef90dc26a005f055ab75 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Thu, 19 Sep 2024 16:25:38 -0300
Subject: [PATCH 17/22] feat: bring coverage 100%

---
 test/SplitterContract.test.ts | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 1b4a4e2..c47389f 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -4,6 +4,7 @@ import { loadFixture } from '@nomicfoundation/hardhat-network-helpers'
 import { expect } from 'chai'
 import { MockERC20, ERC20Splitter, MaliciousRecipient } from '../typechain-types'
 import { AddressZero } from '../utils/constants'
+import { AddressLike, Typed } from 'ethers'
 
 describe('ERC20Splitter', () => {
   let splitter: ERC20Splitter
@@ -271,6 +272,17 @@ describe('ERC20Splitter', () => {
         ).to.be.revertedWith('ERC20Splitter: Shares and recipients length mismatch')
       })
 
+      it('Should revert if shares do not sum to 100%', async () => {
+        const invalidShares = [[4000, 4000, 2000]] // Sums to 90%
+        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
+
+        await mockERC20.transferReverts(true, 0)
+
+        await expect(
+          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], invalidShares, recipients),
+        ).to.be.revertedWith('ERC20Splitter: Transfer failed')
+      })
+
       it('Should revert when msg.value does not match the expected Ether amount', async () => {
         const incorrectMsgValue = ethers.parseEther('1') // Incorrect Ether amount
         const correctEtherAmount = ethers.parseEther('2') // Correct Ether amount to be split
@@ -485,6 +497,26 @@ describe('ERC20Splitter', () => {
 
         await expect(splitter.connect(recipient1).withdraw(tokens)).to.be.revertedWith('ERC20Splitter: Transfer failed')
       })
+
+      it('Should revert when ERC20 transferFrom fails during withdraw', async () => {
+        const tokens: AddressLike[] | Typed = []
+
+        const ethAmount = ethers.parseEther('1')
+        const tokenAddresses = [ethers.ZeroAddress] // Ether represented by address zero
+        const amounts = [ethAmount]
+        const shares = [[10000]] // 100% share
+        const recipients = [[recipient1.getAddress()]]
+
+        await splitter.connect(owner).deposit(tokenAddresses, amounts, shares, recipients, {
+          value: ethAmount,
+        })
+
+        await mockERC20.transferReverts(true, 0)
+
+        await expect(splitter.connect(recipient1).withdraw(tokens)).to.be.revertedWith(
+          'ERC20Splitter: No tokens specified',
+        )
+      })
     })
 
     describe('Withdraw ERC-20 and Native Tokens', async () => {

From 254c27af135e08991475570e5553fc67e75c0929 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Fri, 20 Sep 2024 11:24:58 -0300
Subject: [PATCH 18/22] fix: comments from PR, revert hardhat config

---
 contracts/ERC20Splitter.sol   |  7 +++++--
 hardhat.config.ts             |  2 +-
 test/SplitterContract.test.ts | 22 ++++++----------------
 3 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index 609f9eb..02a25af 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -57,6 +57,7 @@ contract ERC20Splitter is ReentrancyGuard {
 
     /// @notice Withdraw all tokens that the caller is entitled to.
     /// Tokens are automatically determined based on previous deposits.
+    /// @param tokenAddresses Array of token addresses (use address(0) for native tokens).
     function withdraw(address[] calldata tokenAddresses) external nonReentrant {
         uint256 tokenCount = tokenAddresses.length;
         require(tokenCount > 0, 'ERC20Splitter: No tokens specified');
@@ -68,7 +69,7 @@ contract ERC20Splitter is ReentrancyGuard {
             uint256 amount = balances[tokenAddress][msg.sender];
 
             if (amount == 0) {
-                continue;
+                return;
             }
 
             delete balances[tokenAddress][msg.sender];
@@ -102,7 +103,9 @@ contract ERC20Splitter is ReentrancyGuard {
         address[] calldata recipients
     ) internal {
         require(shares.length == recipients.length, 'ERC20Splitter: Shares and recipients length mismatch');
-        require(amount > 0, 'ERC20Splitter: Amount must be greater than zero');
+        if(amount == 0 ) {
+            return;
+        }
 
         uint256 totalSharePercentage = 0;
 
diff --git a/hardhat.config.ts b/hardhat.config.ts
index 16a51e0..2664c0f 100644
--- a/hardhat.config.ts
+++ b/hardhat.config.ts
@@ -67,7 +67,7 @@ const BASE_CONFIG = {
       forking: {
         url: POLYGON_PROVIDER_URL,
         blockNumber: 55899875,
-        timeout: 9999900,
+        timeout: 20000,
       },
     },
   },
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index c47389f..2d651b9 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -122,6 +122,12 @@ describe('ERC20Splitter', () => {
         expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('100'))
       })
 
+      it('Should handle deposit when user has no tokens', async () => {
+        const recipients = [[recipient1.address]]
+        const shares = [[10000]]
+        await splitter.connect(owner).deposit([mockERC20.getAddress()], [0], shares, recipients)
+      })
+
       it('Should deposit four ERC20 tokens and split them between recipients', async () => {
         const tokenAmounts = [
           ethers.parseEther('100'),
@@ -238,7 +244,6 @@ describe('ERC20Splitter', () => {
         expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(ethers.parseEther('30'))
         expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('20'))
       })
-
       it('Should deposit native tokens (ETH) and split them between recipients', async () => {
         const shares = [[5000, 3000, 2000]]
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
@@ -297,21 +302,6 @@ describe('ERC20Splitter', () => {
           }),
         ).to.be.revertedWith('ERC20Splitter: Incorrect native token amount sent')
       })
-      it('Should revert when amount is 0', async () => {
-        const incorrectMsgValue = ethers.parseEther('1') // Incorrect Ether amount
-        const correctEtherAmount = ethers.parseEther('2') // Correct Ether amount to be split
-        const tokenAddresses = [ethers.ZeroAddress] // Using address(0) for Ether
-        const amounts = [correctEtherAmount] // Amount to split among recipients
-        const shares = [[5000, 3000, 2000]] // Shares summing up to 100%
-        const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
-
-        await expect(
-          splitter.connect(owner).deposit(tokenAddresses, [0], shares, recipients, {
-            value: incorrectMsgValue, // Sending incorrect msg.value
-          }),
-        ).to.be.revertedWith('ERC20Splitter: Amount must be greater than zero')
-      })
-
       it('Should revert when tokenAddresses and amounts lengths mismatch', async () => {
         const tokenAddresses = [mockERC20.getAddress(), ethers.ZeroAddress]
         const amounts = [ethers.parseEther('100')] // Length 1, intentional mismatch

From b9c323509e0ebe48f17a5f23fea4f75758e562c5 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Fri, 20 Sep 2024 11:55:20 -0300
Subject: [PATCH 19/22] fix: comments from PR

---
 contracts/ERC20Splitter.sol | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index 02a25af..8a0f486 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -67,9 +67,9 @@ contract ERC20Splitter is ReentrancyGuard {
         for (uint256 i = 0; i < tokenCount; i++) {
             address tokenAddress = tokenAddresses[i];
             uint256 amount = balances[tokenAddress][msg.sender];
-
+            withdrawnAmounts[i] = amount;
             if (amount == 0) {
-                return;
+                continue;
             }
 
             delete balances[tokenAddress][msg.sender];
@@ -82,8 +82,6 @@ contract ERC20Splitter is ReentrancyGuard {
                     'ERC20Splitter: Transfer failed'
                 );
             }
-
-            withdrawnAmounts[i] = amount;
         }
 
         emit Withdraw(msg.sender, tokenAddresses, withdrawnAmounts);

From 534f2851f2a1260afd041c0a600fd371e3e4c5ea Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Fri, 20 Sep 2024 11:59:43 -0300
Subject: [PATCH 20/22] fix: comments from PR

---
 contracts/ERC20Splitter.sol | 1 +
 1 file changed, 1 insertion(+)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index 8a0f486..b245664 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -68,6 +68,7 @@ contract ERC20Splitter is ReentrancyGuard {
             address tokenAddress = tokenAddresses[i];
             uint256 amount = balances[tokenAddress][msg.sender];
             withdrawnAmounts[i] = amount;
+            
             if (amount == 0) {
                 continue;
             }

From 5d74170c5a6df7cf983c72f417cf5dcf9288d322 Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Fri, 20 Sep 2024 13:12:17 -0300
Subject: [PATCH 21/22] fix: improve the tests

---
 test/SplitterContract.test.ts | 122 +++++++++++++++++-----------------
 1 file changed, 62 insertions(+), 60 deletions(-)

diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index 2d651b9..bddaed3 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -22,6 +22,7 @@ describe('ERC20Splitter', () => {
 
   const tokenAmount = ethers.parseEther('100')
   const ethAmount = ethers.parseEther('1')
+  let mockERC20Address: string
 
   before(async function () {
     // prettier-ignore
@@ -65,6 +66,7 @@ describe('ERC20Splitter', () => {
     mockERC20_3 = contracts.mockERC20_3
     mockERC20_4 = contracts.mockERC20_4
     splitter = contracts.splitter
+    mockERC20Address = await mockERC20.getAddress() // Store the address
 
     // Mint tokens to the owner
     await mockERC20.connect(owner).mint(owner, ethers.parseEther('1000'))
@@ -115,17 +117,18 @@ describe('ERC20Splitter', () => {
         const shares = [[10000]] // 50%, 30%, 20%
         const recipients = [[recipient1.address]]
 
-        await expect(
-          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], shares, recipients),
-        ).to.emit(splitter, 'Deposit')
+        await expect(splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], shares, recipients)).to.emit(
+          splitter,
+          'Deposit',
+        )
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('100'))
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(ethers.parseEther('100'))
       })
 
       it('Should handle deposit when user has no tokens', async () => {
         const recipients = [[recipient1.address]]
         const shares = [[10000]]
-        await splitter.connect(owner).deposit([mockERC20.getAddress()], [0], shares, recipients)
+        await splitter.connect(owner).deposit([mockERC20Address], [0], shares, recipients)
       })
 
       it('Should deposit four ERC20 tokens and split them between recipients', async () => {
@@ -142,14 +145,14 @@ describe('ERC20Splitter', () => {
           splitter
             .connect(owner)
             .deposit(
-              [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
+              [mockERC20Address, mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
               tokenAmounts,
               shares,
               recipients,
             ),
         ).to.emit(splitter, 'Deposit')
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('100'))
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(ethers.parseEther('100'))
         expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(ethers.parseEther('100'))
         expect(await splitter.balances(mockERC20_3.getAddress(), recipient3.address)).to.equal(ethers.parseEther('100'))
         expect(await splitter.balances(mockERC20_4.getAddress(), recipient4.address)).to.equal(ethers.parseEther('100'))
@@ -181,16 +184,16 @@ describe('ERC20Splitter', () => {
           splitter
             .connect(owner)
             .deposit(
-              [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
+              [mockERC20Address, mockERC20_2.getAddress(), mockERC20_3.getAddress(), mockERC20_4.getAddress()],
               tokenAmounts,
               shares,
               recipients,
             ),
         ).to.emit(splitter, 'Deposit')
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('50'))
-        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(ethers.parseEther('30'))
-        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('20'))
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(ethers.parseEther('50'))
+        expect(await splitter.balances(mockERC20Address, recipient2.address)).to.equal(ethers.parseEther('30'))
+        expect(await splitter.balances(mockERC20Address, recipient3.address)).to.equal(ethers.parseEther('20'))
 
         expect(await splitter.balances(mockERC20_2.getAddress(), recipient1.address)).to.equal(ethers.parseEther('30'))
         expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(ethers.parseEther('18'))
@@ -220,14 +223,14 @@ describe('ERC20Splitter', () => {
           splitter
             .connect(owner)
             .deposit(
-              [mockERC20.getAddress(), mockERC20_2.getAddress(), mockERC20_3.getAddress()],
+              [mockERC20Address, mockERC20_2.getAddress(), mockERC20_3.getAddress()],
               tokenAmounts,
               shares,
               recipients,
             ),
         ).to.emit(splitter, 'Deposit')
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('50'))
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(ethers.parseEther('50'))
         expect(await splitter.balances(mockERC20_2.getAddress(), recipient2.address)).to.equal(ethers.parseEther('30'))
         expect(await splitter.balances(mockERC20_3.getAddress(), recipient3.address)).to.equal(ethers.parseEther('20'))
       })
@@ -236,13 +239,14 @@ describe('ERC20Splitter', () => {
         const shares = [[5000, 3000, 2000]] // 50%, 30%, 20%
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
-        await expect(
-          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], shares, recipients),
-        ).to.emit(splitter, 'Deposit')
+        await expect(splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], shares, recipients)).to.emit(
+          splitter,
+          'Deposit',
+        )
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(ethers.parseEther('50'))
-        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(ethers.parseEther('30'))
-        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('20'))
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(ethers.parseEther('50'))
+        expect(await splitter.balances(mockERC20Address, recipient2.address)).to.equal(ethers.parseEther('30'))
+        expect(await splitter.balances(mockERC20Address, recipient3.address)).to.equal(ethers.parseEther('20'))
       })
       it('Should deposit native tokens (ETH) and split them between recipients', async () => {
         const shares = [[5000, 3000, 2000]]
@@ -264,7 +268,7 @@ describe('ERC20Splitter', () => {
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
         await expect(
-          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], invalidShares, recipients),
+          splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], invalidShares, recipients),
         ).to.be.revertedWith('ERC20Splitter: Shares must sum to 100%')
       })
 
@@ -273,7 +277,7 @@ describe('ERC20Splitter', () => {
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]] // 3 recipients
 
         await expect(
-          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], invalidShares, recipients),
+          splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], invalidShares, recipients),
         ).to.be.revertedWith('ERC20Splitter: Shares and recipients length mismatch')
       })
 
@@ -284,7 +288,7 @@ describe('ERC20Splitter', () => {
         await mockERC20.transferReverts(true, 0)
 
         await expect(
-          splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], invalidShares, recipients),
+          splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], invalidShares, recipients),
         ).to.be.revertedWith('ERC20Splitter: Transfer failed')
       })
 
@@ -303,7 +307,7 @@ describe('ERC20Splitter', () => {
         ).to.be.revertedWith('ERC20Splitter: Incorrect native token amount sent')
       })
       it('Should revert when tokenAddresses and amounts lengths mismatch', async () => {
-        const tokenAddresses = [mockERC20.getAddress(), ethers.ZeroAddress]
+        const tokenAddresses = [mockERC20Address, ethers.ZeroAddress]
         const amounts = [ethers.parseEther('100')] // Length 1, intentional mismatch
         const shares = [[5000, 3000, 2000]] // Correct length
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
@@ -316,7 +320,7 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should revert when tokenAddresses, shares, and recipients lengths mismatch', async () => {
-        const tokenAddresses = [mockERC20.getAddress(), ethers.ZeroAddress]
+        const tokenAddresses = [mockERC20Address, ethers.ZeroAddress]
         const amounts = [ethers.parseEther('100'), ethers.parseEther('2')]
         const shares = [
           [5000, 3000, 2000], // Length 1
@@ -334,7 +338,7 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should revert when shares and recipients lengths mismatch within sub-arrays', async () => {
-        const tokenAddresses = [mockERC20.getAddress()] // Length 1
+        const tokenAddresses = [mockERC20Address] // Length 1
         const amounts = [ethers.parseEther('100')] // Length 1
         const shares = [[5000, 3000, 2000]] // Length 1, sub-array length 3
         const recipients = [
@@ -395,7 +399,7 @@ describe('ERC20Splitter', () => {
           splitter
             .connect(owner)
             .deposit(
-              [AddressZero, mockERC20.getAddress()],
+              [AddressZero, mockERC20Address],
               [ethAmount, erc20Amount],
               [ethShares[0], erc20Shares[0]],
               [ethRecipients, erc20Recipients],
@@ -408,10 +412,10 @@ describe('ERC20Splitter', () => {
 
         // Check balances for recipient2 (50% of 1 ETH + 60% of 100 ERC-20 tokens = 0.5 ETH + 60 tokens)
         expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(ethers.parseEther('0.5'))
-        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(ethers.parseEther('60'))
+        expect(await splitter.balances(mockERC20Address, recipient2.address)).to.equal(ethers.parseEther('60'))
 
         // Check balances for recipient3 (40% of 100 ERC-20 tokens = 40 tokens)
-        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(ethers.parseEther('40'))
+        expect(await splitter.balances(mockERC20Address, recipient3.address)).to.equal(ethers.parseEther('40'))
       })
     })
 
@@ -421,20 +425,20 @@ describe('ERC20Splitter', () => {
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
         await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
-        await splitter.connect(owner).deposit([mockERC20.getAddress()], [tokenAmount], shares, recipients)
+        await splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], shares, recipients)
       })
 
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
-        const tokens = [await mockERC20.getAddress()]
+        const tokens = [await mockERC20Address]
         await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [await mockERC20Address], [ethers.parseEther('50')])
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(0)
       })
 
       it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
-        const tokens = [await mockERC20.getAddress(), AddressZero]
+        const tokens = [await mockERC20Address, AddressZero]
         const shares = [[5000, 3000, 2000]]
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
@@ -446,22 +450,22 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
-            [await mockERC20.getAddress(), AddressZero],
+            [await mockERC20Address, AddressZero],
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
-        expect(await splitter.balances(await mockERC20.getAddress(), recipient1.address)).to.equal(0)
+        expect(await splitter.balances(await mockERC20Address, recipient1.address)).to.equal(0)
         expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
       })
 
       it('Should handle withdraw() when user has no tokens', async () => {
-        const tokens = [await mockERC20.getAddress(), AddressZero]
+        const tokens = [await mockERC20Address, AddressZero]
         await splitter.connect(anotherUser).withdraw(tokens)
       })
 
       it('Should revert when ERC20 transferFrom fails during withdraw', async () => {
-        const tokens = [await mockERC20.getAddress()]
-        const mockERC20false = await mockERC20.getAddress()
+        const tokens = [await mockERC20Address]
+        const mockERC20false = await mockERC20Address
 
         await network.provider.request({
           method: 'hardhat_impersonateAccount',
@@ -515,20 +519,20 @@ describe('ERC20Splitter', () => {
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
         await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
-        await splitter.connect(owner).deposit([await mockERC20.getAddress()], [tokenAmount], shares, recipients)
+        await splitter.connect(owner).deposit([await mockERC20Address], [tokenAmount], shares, recipients)
       })
 
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
-        const tokens = [await mockERC20.getAddress()]
+        const tokens = [await mockERC20Address]
         await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [await mockERC20.getAddress()], [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [await mockERC20Address], [ethers.parseEther('50')])
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(0)
       })
 
       it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
-        const tokens = [await mockERC20.getAddress(), AddressZero]
+        const tokens = [await mockERC20Address, AddressZero]
         const shares = [[5000, 3000, 2000]]
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
@@ -540,12 +544,12 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address, // Expect both ERC-20 and native token
-            [await mockERC20.getAddress(), AddressZero],
+            [await mockERC20Address, AddressZero],
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
         expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
-        expect(await splitter.balances(mockERC20.getAddress(), recipient1.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(0)
       })
     })
 
@@ -587,14 +591,12 @@ describe('ERC20Splitter', () => {
 
         // Then, deposit ERC-20 tokens for recipient2 and recipient3
         await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
-        await splitter
-          .connect(owner)
-          .deposit([await mockERC20.getAddress()], [tokenAmount], erc20Shares, erc20Recipients)
+        await splitter.connect(owner).deposit([await mockERC20Address], [tokenAmount], erc20Shares, erc20Recipients)
       })
 
       it('Should allow recipient1 to withdraw only their ETH and other recipients to withdraw their ERC-20 tokens', async () => {
         const tokenEth = [AddressZero]
-        const tokenErc20 = [await mockERC20.getAddress()]
+        const tokenErc20 = [await mockERC20Address]
         await expect(splitter.connect(recipient1).withdraw(tokenEth))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
@@ -609,21 +611,21 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
-            [await mockERC20.getAddress()],
+            [await mockERC20Address],
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient2.address)).to.equal(0)
 
         await expect(splitter.connect(recipient3).withdraw(tokenErc20))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient3.address,
-            [await mockERC20.getAddress()],
+            [await mockERC20Address],
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient3.address)).to.equal(0)
       })
     })
     describe('Withdraw for both native tokens (ETH) and ERC-20 tokens multiples addresses 0', () => {
@@ -647,7 +649,7 @@ describe('ERC20Splitter', () => {
         await splitter
           .connect(owner)
           .deposit(
-            [AddressZero, mockERC20.getAddress()],
+            [AddressZero, mockERC20Address],
             [ethAmount, erc20Amount],
             [ethShares[0], erc20Shares[0]],
             [ethRecipients, erc20Recipients],
@@ -669,26 +671,26 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow recipient2 to withdraw both ETH and ERC-20 tokens', async () => {
-        const tokens = [AddressZero, await mockERC20.getAddress()]
+        const tokens = [AddressZero, await mockERC20Address]
         await expect(splitter.connect(recipient2).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
-            [AddressZero, await mockERC20.getAddress()],
+            [AddressZero, await mockERC20Address],
             [ethers.parseEther('0.5'), ethers.parseEther('60')], // 50% of 1 ETH and 60 ERC-20 tokens
           )
 
         expect(await splitter.balances(AddressZero, recipient2.address)).to.equal(0)
-        expect(await splitter.balances(mockERC20.getAddress(), recipient2.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient2.address)).to.equal(0)
       })
 
       it('Should allow recipient3 to withdraw only ERC-20 tokens', async () => {
-        const tokens = [await mockERC20.getAddress()]
+        const tokens = [await mockERC20Address]
         await expect(splitter.connect(recipient3).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient3.address, [await mockERC20.getAddress()], [ethers.parseEther('40')])
+          .withArgs(recipient3.address, [await mockERC20Address], [ethers.parseEther('40')])
 
-        expect(await splitter.balances(mockERC20.getAddress(), recipient3.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient3.address)).to.equal(0)
       })
     })
   })

From 5ce93e482ae52038d6d37dbe6aa4470b77d0a05d Mon Sep 17 00:00:00 2001
From: Eduardo Melo <eduardo.melo00@gmail.com>
Date: Fri, 20 Sep 2024 13:45:13 -0300
Subject: [PATCH 22/22] fix: comments

---
 contracts/ERC20Splitter.sol   |  3 +--
 test/SplitterContract.test.ts | 42 +++++++++++++++++------------------
 2 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
index b245664..f06cdfd 100644
--- a/contracts/ERC20Splitter.sol
+++ b/contracts/ERC20Splitter.sol
@@ -56,7 +56,6 @@ contract ERC20Splitter is ReentrancyGuard {
     }
 
     /// @notice Withdraw all tokens that the caller is entitled to.
-    /// Tokens are automatically determined based on previous deposits.
     /// @param tokenAddresses Array of token addresses (use address(0) for native tokens).
     function withdraw(address[] calldata tokenAddresses) external nonReentrant {
         uint256 tokenCount = tokenAddresses.length;
@@ -68,7 +67,7 @@ contract ERC20Splitter is ReentrancyGuard {
             address tokenAddress = tokenAddresses[i];
             uint256 amount = balances[tokenAddress][msg.sender];
             withdrawnAmounts[i] = amount;
-            
+
             if (amount == 0) {
                 continue;
             }
diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
index bddaed3..f1f7778 100644
--- a/test/SplitterContract.test.ts
+++ b/test/SplitterContract.test.ts
@@ -429,16 +429,16 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
-        const tokens = [await mockERC20Address]
+        const tokens = [mockERC20Address]
         await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [await mockERC20Address], [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [mockERC20Address], [ethers.parseEther('50')])
 
         expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(0)
       })
 
       it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
-        const tokens = [await mockERC20Address, AddressZero]
+        const tokens = [mockERC20Address, AddressZero]
         const shares = [[5000, 3000, 2000]]
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
@@ -450,22 +450,22 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address,
-            [await mockERC20Address, AddressZero],
+            [mockERC20Address, AddressZero],
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
-        expect(await splitter.balances(await mockERC20Address, recipient1.address)).to.equal(0)
+        expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(0)
         expect(await splitter.balances(AddressZero, recipient1.address)).to.equal(0)
       })
 
       it('Should handle withdraw() when user has no tokens', async () => {
-        const tokens = [await mockERC20Address, AddressZero]
+        const tokens = [mockERC20Address, AddressZero]
         await splitter.connect(anotherUser).withdraw(tokens)
       })
 
       it('Should revert when ERC20 transferFrom fails during withdraw', async () => {
-        const tokens = [await mockERC20Address]
-        const mockERC20false = await mockERC20Address
+        const tokens = [mockERC20Address]
+        const mockERC20false = mockERC20Address
 
         await network.provider.request({
           method: 'hardhat_impersonateAccount',
@@ -519,20 +519,20 @@ describe('ERC20Splitter', () => {
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
         await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
-        await splitter.connect(owner).deposit([await mockERC20Address], [tokenAmount], shares, recipients)
+        await splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], shares, recipients)
       })
 
       it('Should allow a recipient to withdraw their split ERC20 tokens without specifying token addresses', async () => {
-        const tokens = [await mockERC20Address]
+        const tokens = [mockERC20Address]
         await expect(splitter.connect(recipient1).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient1.address, [await mockERC20Address], [ethers.parseEther('50')])
+          .withArgs(recipient1.address, [mockERC20Address], [ethers.parseEther('50')])
 
         expect(await splitter.balances(mockERC20Address, recipient1.address)).to.equal(0)
       })
 
       it('Should allow a recipient to withdraw their split native tokens (ETH) and ERC20 tokens', async () => {
-        const tokens = [await mockERC20Address, AddressZero]
+        const tokens = [mockERC20Address, AddressZero]
         const shares = [[5000, 3000, 2000]]
         const recipients = [[recipient1.address, recipient2.address, recipient3.address]]
 
@@ -544,7 +544,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient1.address, // Expect both ERC-20 and native token
-            [await mockERC20Address, AddressZero],
+            [mockERC20Address, AddressZero],
             [ethers.parseEther('50'), ethers.parseEther('0.5')], // 50 ERC20 tokens and 0.5 ETH
           )
 
@@ -591,12 +591,12 @@ describe('ERC20Splitter', () => {
 
         // Then, deposit ERC-20 tokens for recipient2 and recipient3
         await mockERC20.connect(owner).approve(splitter.getAddress(), tokenAmount)
-        await splitter.connect(owner).deposit([await mockERC20Address], [tokenAmount], erc20Shares, erc20Recipients)
+        await splitter.connect(owner).deposit([mockERC20Address], [tokenAmount], erc20Shares, erc20Recipients)
       })
 
       it('Should allow recipient1 to withdraw only their ETH and other recipients to withdraw their ERC-20 tokens', async () => {
         const tokenEth = [AddressZero]
-        const tokenErc20 = [await mockERC20Address]
+        const tokenErc20 = [mockERC20Address]
         await expect(splitter.connect(recipient1).withdraw(tokenEth))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
@@ -611,7 +611,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
-            [await mockERC20Address],
+            [mockERC20Address],
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
@@ -621,7 +621,7 @@ describe('ERC20Splitter', () => {
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient3.address,
-            [await mockERC20Address],
+            [mockERC20Address],
             [ethers.parseEther('50')], // 50% of ERC-20 tokens
           )
 
@@ -671,12 +671,12 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow recipient2 to withdraw both ETH and ERC-20 tokens', async () => {
-        const tokens = [AddressZero, await mockERC20Address]
+        const tokens = [AddressZero, mockERC20Address]
         await expect(splitter.connect(recipient2).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
           .withArgs(
             recipient2.address,
-            [AddressZero, await mockERC20Address],
+            [AddressZero, mockERC20Address],
             [ethers.parseEther('0.5'), ethers.parseEther('60')], // 50% of 1 ETH and 60 ERC-20 tokens
           )
 
@@ -685,10 +685,10 @@ describe('ERC20Splitter', () => {
       })
 
       it('Should allow recipient3 to withdraw only ERC-20 tokens', async () => {
-        const tokens = [await mockERC20Address]
+        const tokens = [mockERC20Address]
         await expect(splitter.connect(recipient3).withdraw(tokens))
           .to.emit(splitter, 'Withdraw')
-          .withArgs(recipient3.address, [await mockERC20Address], [ethers.parseEther('40')])
+          .withArgs(recipient3.address, [mockERC20Address], [ethers.parseEther('40')])
 
         expect(await splitter.balances(mockERC20Address, recipient3.address)).to.equal(0)
       })