chore: make storage vars immutable

0xPolygon · Oct 2, 2023 · b3b4d97 · b3b4d97
1 parent d9fd897
commit b3b4d97
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 83 deletions.
diff --git a/script/Deploy.s.sol b/script/Deploy.s.sol
@@ -27,17 +27,19 @@ contract Deploy is Script {
         ProxyAdmin admin = new ProxyAdmin();
         admin.transferOwnership(governance);
 
-        address migrationImplementation = address(new PolygonMigration());
+        address migrationImplementation = address(new PolygonMigration(matic));
 
         address migrationProxy = address(
             new TransparentUpgradeableProxy(
                 migrationImplementation,
                 address(admin),
-                abi.encodeCall(PolygonMigration.initialize, matic)
+                abi.encodeWithSelector(PolygonMigration.initialize.selector)
             )
         );
 
-        address emissionManagerImplementation = address(new DefaultEmissionManager());
+        address emissionManagerImplementation = address(
+            new DefaultEmissionManager(migrationProxy, stakeManager, treasury)
+        );
         address emissionManagerProxy = address(
             new TransparentUpgradeableProxy(address(emissionManagerImplementation), address(admin), "")
         );
@@ -49,13 +51,7 @@ contract Deploy is Script {
             permit2revoker
         );
 
-        DefaultEmissionManager(emissionManagerProxy).initialize(
-            address(polygonToken),
-            migrationProxy,
-            stakeManager,
-            treasury,
-            governance
-        );
+        DefaultEmissionManager(emissionManagerProxy).initialize(address(polygonToken), governance);
 
         PolygonMigration(migrationProxy).setPolygonToken(address(polygonToken));
 

diff --git a/src/DefaultEmissionManager.sol b/src/DefaultEmissionManager.sol
@@ -21,45 +21,34 @@ contract DefaultEmissionManager is Ownable2StepUpgradeable, IDefaultEmissionMana
     uint256 public constant START_SUPPLY = 10_000_000_000e18;
     address private immutable DEPLOYER;
 
-    IPolygonEcosystemToken public token;
-    IPolygonMigration public migration;
-    address public stakeManager;
-    address public treasury;
+    IPolygonMigration public immutable migration;
+    address public immutable stakeManager;
+    address public immutable treasury;
 
+    IPolygonEcosystemToken public token;
     uint256 public startTimestamp;
 
-    constructor() {
+    constructor(address migration_, address stakeManager_, address treasury_) {
         DEPLOYER = msg.sender;
+        migration = IPolygonMigration(migration_);
+        stakeManager = stakeManager_;
+        treasury = treasury_;
+
         // so that the implementation contract cannot be initialized
         _disableInitializers();
     }
 
-    function initialize(
-        address token_,
-        address migration_,
-        address stakeManager_,
-        address treasury_,
-        address owner_
-    ) external initializer {
+    function initialize(address token_, address owner_) external initializer {
         // prevent front-running since we can't initialize on proxy deployment
         if (DEPLOYER != msg.sender) revert();
-        if (
-            token_ == address(0) ||
-            migration_ == address(0) ||
-            stakeManager_ == address(0) ||
-            treasury_ == address(0) ||
-            owner_ == address(0)
-        ) revert InvalidAddress();
+        if (token_ == address(0) || owner_ == address(0)) revert InvalidAddress();
 
         token = IPolygonEcosystemToken(token_);
-        migration = IPolygonMigration(migration_);
-        stakeManager = stakeManager_;
-        treasury = treasury_;
         startTimestamp = block.timestamp;
 
         assert(START_SUPPLY == token.totalSupply());
 
-        token.safeApprove(migration_, type(uint256).max);
+        token.safeApprove(address(migration), type(uint256).max);
         // initial ownership setup bypassing 2 step ownership transfer process
         _transferOwnership(owner_);
     }

diff --git a/src/PolygonMigration.sol b/src/PolygonMigration.sol
@@ -16,24 +16,23 @@ contract PolygonMigration is Ownable2StepUpgradeable, IPolygonMigration {
     using SafeERC20 for IERC20;
     using SafeERC20 for IERC20Permit;
 
+    IERC20 public immutable matic;
     IERC20 public polygon;
-    IERC20 public matic;
     bool public unmigrationLocked;
 
     modifier onlyUnmigrationUnlocked() {
         if (unmigrationLocked) revert UnmigrationLocked();
         _;
     }
 
-    constructor() {
+    constructor(address matic_) {
+        matic = IERC20(matic_);
         // so that the implementation contract cannot be initialized
         _disableInitializers();
     }
 
-    function initialize(address matic_) external initializer {
+    function initialize() external initializer {
         __Ownable_init();
-        if (matic_ == address(0)) revert InvalidAddress();
-        matic = IERC20(matic_);
     }
 
     /// @notice This function allows owner/governance to set POL token address *only once*

diff --git a/test/DefaultEmissionManager.t.sol b/test/DefaultEmissionManager.t.sol
@@ -29,16 +29,21 @@ contract DefaultEmissionManagerTest is Test {
         treasury = makeAddr("treasury");
         stakeManager = makeAddr("stakeManager");
         governance = makeAddr("governance");
-        emissionManagerImplementation = new DefaultEmissionManager();
         ProxyAdmin admin = new ProxyAdmin();
-        emissionManager = DefaultEmissionManager(
-            address(new TransparentUpgradeableProxy(address(emissionManagerImplementation), address(admin), ""))
-        );
         matic = new ERC20PresetMinterPauser("Matic Token", "MATIC");
         migration = PolygonMigration(
-            address(new TransparentUpgradeableProxy(address(new PolygonMigration()), address(admin), ""))
+            address(
+                new TransparentUpgradeableProxy(
+                    address(new PolygonMigration(address(matic))),
+                    address(admin),
+                    abi.encodeWithSelector(PolygonMigration.initialize.selector)
+                )
+            )
+        );
+        emissionManagerImplementation = new DefaultEmissionManager(address(migration), stakeManager, treasury);
+        emissionManager = DefaultEmissionManager(
+            address(new TransparentUpgradeableProxy(address(emissionManagerImplementation), address(admin), ""))
         );
-        migration.initialize(address(matic));
         polygon = new PolygonEcosystemToken(
             address(migration),
             address(emissionManager),
@@ -49,7 +54,7 @@ contract DefaultEmissionManagerTest is Test {
         migration.transferOwnership(governance);
         vm.prank(governance);
         migration.acceptOwnership();
-        emissionManager.initialize(address(polygon), address(migration), stakeManager, treasury, governance);
+        emissionManager.initialize(address(polygon), governance);
         // POL being emissionary, while MATIC having a constant supply,
         // the requirement of unmigrating POL to MATIC for StakeManager on each mint
         // is satisfied by a one-time transfer of MATIC to the migration contract
@@ -62,7 +67,7 @@ contract DefaultEmissionManagerTest is Test {
 
     function testRevert_Initialize() external {
         vm.expectRevert("Initializable: contract is already initialized");
-        emissionManager.initialize(address(0), address(0), address(0), address(0), address(0));
+        emissionManager.initialize(address(0), address(0));
     }
 
     function test_Deployment() external {
@@ -75,44 +80,35 @@ contract DefaultEmissionManagerTest is Test {
         assertEq(polygon.totalSupply(), 10_000_000_000e18);
     }
 
-    function test_InvalidDeployment(uint160 seed) external {
-        address[5] memory params = [
-            makeAddr("polygon"),
-            makeAddr("migration"),
-            makeAddr("stakeManager"),
-            makeAddr("treasury"),
-            makeAddr("governance")
-        ];
+    function test_InvalidDeployment() external {
+        address _polygon = makeAddr("polygon");
+        address _governance = makeAddr("governance");
 
-        address proxy = address(new TransparentUpgradeableProxy(address(new DefaultEmissionManager()), msg.sender, ""));
+        address proxy = address(
+            new TransparentUpgradeableProxy(
+                address(new DefaultEmissionManager(address(migration), stakeManager, treasury)),
+                msg.sender,
+                ""
+            )
+        );
 
-        vm.prank(address(seed));
+        vm.prank(address(0x1337));
         vm.expectRevert();
-        DefaultEmissionManager(proxy).initialize(params[0], params[1], params[2], params[3], params[4]);
-
-        params[seed % params.length] = address(0); // any one is zero addr
+        DefaultEmissionManager(proxy).initialize(_polygon, _governance);
 
         vm.expectRevert(InvalidAddress.selector);
-        DefaultEmissionManager(proxy).initialize(params[0], params[1], params[2], params[3], params[4]);
+        DefaultEmissionManager(proxy).initialize(_polygon, address(0));
+        vm.expectRevert(InvalidAddress.selector);
+        DefaultEmissionManager(proxy).initialize(address(0), _governance);
+        vm.expectRevert(InvalidAddress.selector);
+        DefaultEmissionManager(proxy).initialize(address(0), address(0));
     }
 
     function test_ImplementationCannotBeInitialized() external {
         vm.expectRevert("Initializable: contract is already initialized");
-        DefaultEmissionManager(address(emissionManagerImplementation)).initialize(
-            address(0),
-            address(0),
-            address(0),
-            address(0),
-            address(0)
-        );
+        DefaultEmissionManager(address(emissionManagerImplementation)).initialize(address(0), address(0));
         vm.expectRevert("Initializable: contract is already initialized");
-        DefaultEmissionManager(address(emissionManager)).initialize(
-            address(0),
-            address(0),
-            address(0),
-            address(0),
-            address(0)
-        );
+        DefaultEmissionManager(address(emissionManager)).initialize(address(0), address(0));
     }
 
     function test_Mint() external {

diff --git a/test/PolygonEcosystemToken.t.sol b/test/PolygonEcosystemToken.t.sol
@@ -29,10 +29,16 @@ contract PolygonTest is Test {
         permit2revoker = makeAddr("permit2revoker");
         ProxyAdmin admin = new ProxyAdmin();
         emissionManager = DefaultEmissionManager(
-            address(new TransparentUpgradeableProxy(address(new DefaultEmissionManager()), address(admin), ""))
+            address(
+                new TransparentUpgradeableProxy(
+                    address(new DefaultEmissionManager(migration, stakeManager, treasury)),
+                    address(admin),
+                    ""
+                )
+            )
         );
         polygon = new PolygonEcosystemToken(migration, address(emissionManager), governance, permit2revoker);
-        emissionManager.initialize(address(polygon), migration, stakeManager, treasury, msg.sender);
+        emissionManager.initialize(address(polygon), msg.sender);
     }
 
     function test_Deployment(address owner) external {

diff --git a/test/PolygonMigration.t.sol b/test/PolygonMigration.t.sol
@@ -30,9 +30,9 @@ contract PolygonMigrationTest is Test {
         migration = PolygonMigration(
             address(
                 new TransparentUpgradeableProxy(
-                    address(new PolygonMigration()),
+                    address(new PolygonMigration(address(matic))),
                     address(admin),
-                    abi.encodeCall(PolygonMigration.initialize, address(matic))
+                    abi.encodeWithSelector(PolygonMigration.initialize.selector)
                 )
             )
         );
@@ -58,13 +58,11 @@ contract PolygonMigrationTest is Test {
 
     function test_InvalidDeployment() external {
         PolygonMigration temp = PolygonMigration(
-            address(new TransparentUpgradeableProxy(address(new PolygonMigration()), msg.sender, ""))
+            address(new TransparentUpgradeableProxy(address(new PolygonMigration(address(matic))), msg.sender, ""))
         );
-        vm.expectRevert(IPolygonMigration.InvalidAddress.selector);
-        temp.initialize(address(0));
-        temp.initialize(address(matic));
+        temp.initialize();
         vm.expectRevert("Initializable: contract is already initialized");
-        temp.initialize(address(1));
+        temp.initialize();
     }
 
     function test_Migrate(address user, uint256 amount) external {