From b9c3cfe1d87748fd1631d3be4d1f1f21a8337bac Mon Sep 17 00:00:00 2001
From: gretzke <daniel@gretzke.de>
Date: Fri, 1 Sep 2023 18:06:28 +0200
Subject: [PATCH] Make unmigration lock a boolean for simplicity and future
 maintenance

---
 src/PolygonMigration.sol             | 14 ++++++--------
 src/interfaces/IPolygonMigration.sol |  2 +-
 test/PolygonMigration.t.sol          | 10 +++-------
 3 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/src/PolygonMigration.sol b/src/PolygonMigration.sol
index 9a3b190..4ab42c2 100644
--- a/src/PolygonMigration.sol
+++ b/src/PolygonMigration.sol
@@ -19,10 +19,10 @@ contract PolygonMigration is Ownable2Step, IPolygonMigration {
     IERC20 public polygon;
     IERC20 public immutable matic;
     uint256 public releaseTimestamp;
-    uint256 public unmigrationLock;
+    bool public unmigrationLocked;
 
     modifier ifUnmigrationUnlocked() {
-        if (unmigrationLock != 0) revert UnmigrationLocked();
+        if (unmigrationLocked) revert UnmigrationLocked();
         _;
     }
 
@@ -109,12 +109,10 @@ contract PolygonMigration is Ownable2Step, IPolygonMigration {
 
     /// @notice Allows governance to lock or unlock the unmigration process
     /// @dev The function does not do any validation since governance can update the unmigration process if required
-    /// @param unmigrationLock_ New unmigration lock status
-    function updateUnmigrationLock(
-        uint256 unmigrationLock_
-    ) external onlyOwner {
-        unmigrationLock = unmigrationLock_;
-        emit UnmigrationLockUpdated(unmigrationLock_);
+    /// @param unmigrationLocked_ New unmigration lock status
+    function updateUnmigrationLock(bool unmigrationLocked_) external onlyOwner {
+        unmigrationLocked = unmigrationLocked_;
+        emit UnmigrationLockUpdated(unmigrationLocked_);
     }
 
     /// @notice Allows governance to release the remaining POL tokens after the migration period has elapsed
diff --git a/src/interfaces/IPolygonMigration.sol b/src/interfaces/IPolygonMigration.sol
index 0b4ee88..546a381 100644
--- a/src/interfaces/IPolygonMigration.sol
+++ b/src/interfaces/IPolygonMigration.sol
@@ -10,7 +10,7 @@ interface IPolygonMigration {
 
     event Migrated(address indexed account, uint256 amount);
     event Unmigrated(address indexed account, uint256 amount);
-    event UnmigrationLockUpdated(uint256 lock);
+    event UnmigrationLockUpdated(bool lock);
     event ReleaseTimestampUpdated(uint256 timestamp);
     event Released(uint256 polAmount, uint256 maticAmount);
 
diff --git a/test/PolygonMigration.t.sol b/test/PolygonMigration.t.sol
index 2897674..177aac6 100644
--- a/test/PolygonMigration.t.sol
+++ b/test/PolygonMigration.t.sol
@@ -113,16 +113,12 @@ contract PolygonMigrationTest is Test {
         assertEq(matic.balanceOf(user), amount2);
     }
 
-    function testRevert_Unmigrate(
-        address user,
-        uint256 amount,
-        uint256 unmigrationLock
-    ) external {
+    function testRevert_Unmigrate(address user, uint256 amount) external {
+        bool unmigrationLock = true;
         vm.assume(
             amount <= 10000000000 * 10 ** 18 &&
                 user != address(0) &&
-                user != address(migration) &&
-                unmigrationLock != 0
+                user != address(migration)
         );
         matic.mint(user, amount);
         vm.startPrank(user);