From 0c27c8b0956f76e80b2be328eab9c52b94d310ae Mon Sep 17 00:00:00 2001
From: Sergey <2901744+evercoinx@users.noreply.github.com.>
Date: Fri, 30 Aug 2024 20:15:20 +0200
Subject: [PATCH] Implement method to submit POL

---
 contracts/MaticX.sol                     | 38 ++++++++++++++++++++----
 contracts/interfaces/IMaticX.sol         |  2 ++
 contracts/interfaces/IValidatorShare.sol |  4 +++
 contracts/mocks/ValidatorShareMock.sol   | 36 ++++++++++++++--------
 4 files changed, 62 insertions(+), 18 deletions(-)

diff --git a/contracts/MaticX.sol b/contracts/MaticX.sol
index cedcf7d4..d4b17903 100644
--- a/contracts/MaticX.sol
+++ b/contracts/MaticX.sol
@@ -181,7 +181,8 @@ contract MaticX is
 
 		uint256 maticxMinted = delegateToMint(
 			address(this),
-			instantPoolMatic
+			instantPoolMatic,
+			false
 		);
 		instantPoolMaticX += maticxMinted;
 		instantPoolMatic = 0;
@@ -219,7 +220,7 @@ contract MaticX is
 	////////////////////////////////////////////////////////////
 
 	/**
-	 * @dev Send funds to MaticX contract and mints MaticX to msg.sender
+	 * @dev Send MATIC token to MaticX contract and mints MaticX to msg.sender
 	 * @notice Requires that msg.sender has approved _amount of MATIC to this contract
 	 * @param _amount - Amount of MATIC sent from msg.sender to this contract
 	 * @return Amount of MaticX shares generated
@@ -238,7 +239,30 @@ contract MaticX is
 			_amount
 		);
 
-		return delegateToMint(msg.sender, _amount);
+		return delegateToMint(msg.sender, _amount, false);
+	}
+
+	/**
+	 * @dev Send POL token to MaticX contract and mints MaticX to msg.sender
+	 * @notice Requires that msg.sender has approved _amount of POL to this contract
+	 * @param _amount - Amount of POL sent from msg.sender to this contract
+	 * @return Amount of MaticX shares generated
+	 */
+	function submitPOL(uint256 _amount)
+		external
+		override
+		whenNotPaused
+		returns (uint256)
+	{
+		require(_amount > 0, "Invalid amount");
+
+		IERC20Upgradeable(polToken).safeTransferFrom(
+			msg.sender,
+			address(this),
+			_amount
+		);
+
+		return delegateToMint(msg.sender, _amount, true);
 	}
 
 	/**
@@ -483,9 +507,8 @@ contract MaticX is
 	 * @param _amount - Amount of MATIC sent from msg.sender to this contract
 	 * @return Amount of MaticX shares generated
 	 */
-	function delegateToMint(address depositSender, uint256 _amount)
+	function delegateToMint(address depositSender, uint256 _amount, bool pol)
 		internal
-		whenNotPaused
 		returns (uint256)
 	{
 		(
@@ -501,7 +524,10 @@ contract MaticX is
 			.preferredDepositValidatorId();
 		address validatorShare = IStakeManager(stakeManager)
 			.getValidatorContract(preferredValidatorId);
-		IValidatorShare(validatorShare).buyVoucher(_amount, 0);
+
+		pol
+			? IValidatorShare(validatorShare).buyVoucherPOL(_amount, 0)
+			: IValidatorShare(validatorShare).buyVoucher(_amount, 0);
 
 		IFxStateRootTunnel(fxStateRootTunnel).sendMessageToChild(
 			abi.encode(totalShares + amountToMint, totalPooledMatic + _amount)
diff --git a/contracts/interfaces/IMaticX.sol b/contracts/interfaces/IMaticX.sol
index 660db398..5339cded 100644
--- a/contracts/interfaces/IMaticX.sol
+++ b/contracts/interfaces/IMaticX.sol
@@ -53,6 +53,8 @@ interface IMaticX is IERC20Upgradeable {
 
 	function submit(uint256 _amount) external returns (uint256);
 
+	function submitPOL(uint256 _amount) external returns (uint256);
+
 	function requestWithdraw(uint256 _amount) external;
 
 	function claimWithdrawal(uint256 _idx) external;
diff --git a/contracts/interfaces/IValidatorShare.sol b/contracts/interfaces/IValidatorShare.sol
index 8828d99f..cbe80f01 100644
--- a/contracts/interfaces/IValidatorShare.sol
+++ b/contracts/interfaces/IValidatorShare.sol
@@ -19,6 +19,10 @@ interface IValidatorShare {
 		external
 		returns (uint256);
 
+	function buyVoucherPOL(uint256 _amount, uint256 _minSharesToMint)
+		external
+		returns (uint256);
+
 	// solhint-disable-next-line func-name-mixedcase
 	function sellVoucher_new(uint256 claimAmount, uint256 maximumSharesToBurn)
 		external;
diff --git a/contracts/mocks/ValidatorShareMock.sol b/contracts/mocks/ValidatorShareMock.sol
index 03d0edf0..6a22f515 100644
--- a/contracts/mocks/ValidatorShareMock.sol
+++ b/contracts/mocks/ValidatorShareMock.sol
@@ -44,19 +44,16 @@ contract ValidatorShareMock is IValidatorShare {
 		override
 		returns (uint256)
 	{
-		uint256 totalAmount = IERC20(token).balanceOf(address(this));
-
-		uint256 shares = totalAmount != 0
-			? (_amount * totalShares) / totalAmount
-			: _amount;
-
-		totalShares += shares;
-		totalStaked += _amount;
-		require(
-			stakeManager.delegationDeposit(validatorId, _amount, msg.sender),
-			"deposit failed"
-		);
+		_buyVoucher(_amount);
+		return 1;
+	}
 
+	function buyVoucherPOL(uint256 _amount, uint256)
+		external
+		override
+		returns (uint256)
+	{
+		_buyVoucher(_amount);
 		return 1;
 	}
 
@@ -139,4 +136,19 @@ contract ValidatorShareMock is IValidatorShare {
 		);
 		return unbond;
 	}
+
+	function _buyVoucher(uint256 _amount) private {
+		uint256 totalAmount = IERC20(token).balanceOf(address(this));
+
+		uint256 shares = totalAmount != 0
+			? (_amount * totalShares) / totalAmount
+			: _amount;
+
+		totalShares += shares;
+		totalStaked += _amount;
+		require(
+			stakeManager.delegationDeposit(validatorId, _amount, msg.sender),
+			"deposit failed"
+		);
+	}
 }