Unstake & Withdraw Feat on OmnichainStaking Token contract on Base

contracts/governance/locker/staking/OmnichainStakingBase.sol

@@ -156,25 +156,23 @@ abstract contract OmnichainStakingBase is
 
   /// @inheritdoc IOmnichainStaking
   function unstakeToken(uint256 tokenId) external {
-    _updateRewardsAll(msg.sender);
-    require(lockedByToken[tokenId] != address(0), "!tokenId");
-    address lockedBy_ = lockedByToken[tokenId];
-    if (_msgSender() != lockedBy_) {
-      revert InvalidUnstaker(_msgSender(), lockedBy_);
-    }
-
-    delete lockedByToken[tokenId];
-    lockedTokenIdNfts[_msgSender()] = _deleteAnElement(lockedTokenIdNfts[_msgSender()], tokenId);
-
-    // reset and burn voting power
-    _burn(msg.sender, power[tokenId]);
-    power[tokenId] = 0;
-
+    _unstakeToken(tokenId);
     locker.safeTransferFrom(address(this), msg.sender, tokenId);
   }
 
   /// @inheritdoc IOmnichainStaking
-  function increaseLockDuration(uint256 tokenId, uint256 newLockDuration) external {
+  function unstakeAndWithdraw(uint256 tokenId) external nonReentrant {
+    _unstakeToken(tokenId);
+    uint256 lockedAmount = locker.locked(tokenId).amount;
+    locker.withdraw(tokenId);
+    assert(locker.underlying().transfer(msg.sender, lockedAmount));
+  }
+
+  /// @inheritdoc IOmnichainStaking
+  function increaseLockDuration(
+    uint256 tokenId,
+    uint256 newLockDuration
+  ) external {
     require(newLockDuration > 0, "!newLockAmount");
 
     require(msg.sender == lockedByToken[tokenId], "!tokenId");
@@ -379,6 +377,24 @@ abstract contract OmnichainStakingBase is
     return updatedArray;
   }
 
+  function _unstakeToken(uint256 tokenId) internal {
+    _updateRewardsAll(msg.sender);
+    require(lockedByToken[tokenId] != address(0), "!tokenId");
+    address lockedBy_ = lockedByToken[tokenId];
+    if (_msgSender() != lockedBy_) {
+      revert InvalidUnstaker(_msgSender(), lockedBy_);
+    }
+
+    delete lockedByToken[tokenId];
+    lockedTokenIdNfts[_msgSender()] = _deleteAnElement(
+      lockedTokenIdNfts[_msgSender()],
+      tokenId
+    );
+    // reset and burn voting power
+    _burn(msg.sender, power[tokenId]);
+    power[tokenId] = 0;
+  }
+
   /**
    * @notice Called frequently to update the staking parameters associated to an address
    * @param token The token for which the rewards are updated

contracts/interfaces/governance/IOmnichainStaking.sol

@@ -84,6 +84,13 @@ interface IOmnichainStaking is IMultiTokenRewards, IVotes {
    */
   function unstakeToken(uint256 tokenId) external;
 
+  /**
+   * @notice A single withdraw function to unstake NFT, transfer it back to
+   * the user, and also withdraw all tokens for `tokenId` from the locker.
+   * @param tokenId The ID of the regular token NFT to unstake and withdraw.
+   */
+  function unstakeAndWithdraw(uint256 tokenId) external;
+
   /**
    * @notice Updates the lock duration for a specific NFT.
    * @param tokenId The ID of the NFT for which to update the lock duration.

deploy/imp-omnichain-staking-token.ts

@@ -0,0 +1,26 @@
+import { HardhatRuntimeEnvironment } from "hardhat/types";
+import { network } from "hardhat";
+async function main(hre: HardhatRuntimeEnvironment) {
+  const { deployments, getNamedAccounts } = hre;
+  const { deploy } = deployments;
+  const { deployer } = await getNamedAccounts();
+
+  const contract = await deploy("OmnichainStakingTokenImpl", {
+    from: deployer,
+    contract: "OmnichainStakingToken",
+    autoMine: true,
+    skipIfAlreadyDeployed: true,
+    log: true,
+  });
+
+  console.log("OmnichainStakingTokenImpl deployed to:", contract.address);
+
+   if (network.name !== "hardhat") {
+      await hre.run("verify:verify", {
+        address: contract.address,
+      });
+    }
+}
+
+main.tags = ["OmnichainStakingTokenImpl"];
+export default main;