feat(protocol): ProverPool gas optimizations (#14062)

Co-authored-by: adaki2004 <adaki2004@users.noreply.github.com>
Co-authored-by: Daniel Wang <99078276+dantaik@users.noreply.github.com>

packages/protocol/contracts/L1/ProverPool.sol

@@ -41,6 +41,8 @@ contract ProverPool is EssentialContract, IProverPool {
     // reserve more slots than necessary
     uint256[10_000] private proverData;
     mapping(uint256 id => address prover) public idToProver;
+    // Save the weights only when: stake / unstaked / slashed
+    mapping(uint256 id => uint256 weights) public idToWeights;
     mapping(address staker => Staker) public stakers;
 
     uint256[48] private __gap;
@@ -153,6 +155,13 @@ contract ProverPool is EssentialContract, IProverPool {
             _saveProver(staker.proverId, prover);
         }
 
+        uint256 proverWeight = _calcWeight2(
+            staker.maxCapacity,
+            prover.stakedAmount * ONE_TKO,
+            prover.rewardPerGas
+        );
+        idToWeights[staker.proverId] = proverWeight;
+
         emit Slashed(addr, amountToSlash);
     }
 
@@ -224,16 +233,20 @@ contract ProverPool is EssentialContract, IProverPool {
             _stakers[i] = idToProver[i + 1];
         }
     }
-    //Returns each prover's weight dynamically based on feePerGas.
 
+    //Returns each prover's weight dynamically based on feePerGas.
     function getWeights(uint32 feePerGas)
         public
         view
         returns (uint256[MAX_NUM_PROVERS] memory weights, uint256 totalWeight)
     {
         for (uint8 i; i < MAX_NUM_PROVERS; ++i) {
             Prover memory prover = _loadProver(i + 1);
-            weights[i] = _calcWeight(prover, feePerGas);
+            if (prover.currentCapacity == 0) {
+                weights[i] = 0;
+            } else {
+                weights[i] = idToWeights[i + 1];
+            }
             totalWeight += weights[i];
         }
     }
@@ -296,6 +309,10 @@ contract ProverPool is EssentialContract, IProverPool {
         _exit(replaced);
         // }
         idToProver[proverId] = addr;
+        // Keep track of weights when changes ()
+        uint256 proverWeight =
+            _calcWeight2(maxCapacity, amount * ONE_TKO, rewardPerGas);
+        idToWeights[proverId] = proverWeight;
 
         // Assign the staker this proverId
         staker.proverId = proverId;
@@ -319,6 +336,7 @@ contract ProverPool is EssentialContract, IProverPool {
         if (staker.proverId == 0) return;
 
         delete idToProver[staker.proverId];
+        delete idToWeights[staker.proverId];
 
         // Delete the prover but make it non-zero for cheaper rewrites
         // by keep rewardPerGas = 1
@@ -373,6 +391,23 @@ contract ProverPool is EssentialContract, IProverPool {
         }
     }
 
+    // Calculates the user weight's when it stakes/unstakes/slashed
+    function _calcWeight2(
+        uint16 currentCapacity,
+        uint64 stakedAmount,
+        uint16 rewardPerGas
+    )
+        private
+        pure
+        returns (uint256)
+    {
+        if (currentCapacity == 0 || stakedAmount == 0 || rewardPerGas == 0) {
+            return 0;
+        } else {
+            return uint256(stakedAmount) / rewardPerGas / rewardPerGas;
+        }
+    }
+
     function _loadProver(uint256 proverId)
         private
         view

packages/protocol/contracts/L1/ProverPool2.sol

@@ -15,24 +15,32 @@ import { Proxied } from "../common/Proxied.sol";
 contract ProverPool2 is EssentialContract {
     uint256 public constant NUM_SLOTS = 128;
     uint256 public constant EXIT_PERIOD = 1 weeks;
-    uint32 public constant SLASH_POINTS = 500; // basis points
+    uint32 public constant SLASH_POINTS = 9500; // basis points
 
     uint256 public totalStaked;
     uint256 public totalWeight;
 
-    error CAPACITY_TOO_HIGH();
+    error CAPACITY_INCORRECT();
     error NOT_ENOUGH_BALANCE();
+    error CANNOT_BE_PREFERRED();
 
     struct Staker {
         uint256 amount;
         uint256 numSlots;
-        uint256 maxNumSlots; //Max capacity if someone else's unstake would
+        // If type(uint256).max = signals prover can prove all of the blocks
+        // Then gets into the preferredProver (if he is also the max prover)
+        uint256 maxNumSlots; // Max capacity if someone else's unstake would
             // increase a prover's slot count
         uint256 unstakedAt;
-        uint256 unstakedAmount;
         uint16 rewardPerGas;
     }
 
+    // Temporary staker who could jump in as a new prover
+    // when someone unstakes and we need to fill their slots
+    // - until the 'weight-based-owner' claims them (!)
+    // So we basically don't increase anyone's slots unintentionally
+    address preferredProver;
+
     mapping(uint256 slot => address) slots;
     mapping(address staker => Staker) stakers;
 
@@ -70,8 +78,8 @@ contract ProverPool2 is EssentialContract {
     )
         external
     {
-        if (maxCapacity > NUM_SLOTS) {
-            revert CAPACITY_TOO_HIGH();
+        if (maxCapacity > NUM_SLOTS && (maxCapacity != type(uint256).max)) {
+            revert CAPACITY_INCORRECT();
         }
         address staker = msg.sender;
         // If the staker was unstaking, first revert the unstaking
@@ -96,18 +104,26 @@ contract ProverPool2 is EssentialContract {
         }
     }
 
-    function unstake(uint256 unstakedAmount) external {
-        if (stakers[msg.sender].amount < unstakedAmount) {
-            revert NOT_ENOUGH_BALANCE();
-        }
+    function unstake() external {
         address staker = msg.sender;
 
         totalWeight -= getWeight(staker);
         stakers[staker].unstakedAt = block.timestamp;
-        stakers[staker].unstakedAmount += unstakedAmount;
-        stakers[staker].amount -= unstakedAmount;
-        totalStaked -= unstakedAmount;
-        totalWeight += getWeight(staker);
+        totalStaked -= stakers[staker].amount;
+
+        // Exchange unstaked slots with the preferredProver
+        // Auto-claim adjustment
+        uint256 replacedSlots;
+        for (uint256 slotIdx = 0; slotIdx < NUM_SLOTS; slotIdx++) {
+            address current = slots[slotIdx];
+            if (current == staker) {
+                slots[slotIdx] = preferredProver;
+                replacedSlots++;
+            }
+        }
+        // Someone (later) who's weight allows to actually claim
+        // the slots will do that later from preferredProver.
+        stakers[preferredProver].numSlots += replacedSlots;
     }
 
     function setRewardPerGas(uint16 rewardPerGas) external {
@@ -145,43 +161,28 @@ contract ProverPool2 is EssentialContract {
         }
     }
 
-    function slashProver(address slashed) external {
-        Staker memory staker = stakers[slashed];
-
-        uint256 slashableAmount = staker.unstakedAt > 0
-            && block.timestamp <= staker.unstakedAt + EXIT_PERIOD
-            ? staker.amount + staker.unstakedAmount
-            : staker.amount;
-
-        uint256 amountToSlash;
-
-        if (slashableAmount > 0) {
-            amountToSlash = slashableAmount * SLASH_POINTS / 10_000;
-            // make sure we can slash even if  totalAmount is as small as 1
-            if (amountToSlash == 0) amountToSlash = 1;
+    // preferredProver is the one who can (theoretically) prove all
+    // the blocks and also the most staked TKO. He will be assigned
+    // with the slots which will have no 'owner' (until claimed)
+    // when someone unstakes
+    function claimPreferredProverStatus(address staker) external {
+        if (
+            stakers[staker].maxNumSlots != type(uint256).max
+                || stakers[preferredProver].amount >= stakers[staker].amount
+                || stakers[staker].unstakedAt != 0
+        ) {
+            revert CANNOT_BE_PREFERRED();
         }
+        preferredProver = staker;
+    }
 
-        if (amountToSlash == 0) {
-            // do nothing
-        } else if (amountToSlash <= staker.unstakedAmount) {
-            staker.unstakedAmount -= amountToSlash;
-        } else {
-            uint256 _additional = amountToSlash - staker.unstakedAmount;
-            staker.unstakedAmount = 0;
-
-            if (staker.amount > _additional) {
-                staker.amount -= _additional;
-            } else {
-                staker.amount = 0;
-            }
-        }
-        //Write back memory var to storage
-        stakers[slashed] = staker;
+    function slashProver(address slashed) external {
+        stakers[slashed].amount =
+            stakers[slashed].amount * SLASH_POINTS / 10_000;
     }
 
     function withdraw(address staker) public {
         require(stakers[staker].unstakedAt + EXIT_PERIOD >= block.timestamp);
-        stakers[staker].unstakedAmount = 0;
         stakers[staker].unstakedAt = 0;
     }
 

packages/protocol/contracts/L1/TaikoData.sol

@@ -147,6 +147,7 @@ library TaikoData {
         mapping(bytes32 txListHash => TxListInfo) txListInfo;
         mapping(uint256 depositId_mode_ethDepositRingBufferSize => uint256)
             ethDeposits;
+        mapping(address account => uint256 balance) taikoTokenBalances;
         // Never or rarely changed
         // Slot 7: never or rarely changed
         uint64 genesisHeight;
@@ -165,6 +166,6 @@ library TaikoData {
         uint32 feePerGas;
         uint16 avgProofDelay;
         // Reserved
-        uint256[43] __gap; // TODO: update this
+        uint256[42] __gap; // TODO: update this
     }
 }

packages/protocol/contracts/L1/TaikoEvents.sol

@@ -11,7 +11,9 @@ import { TaikoData } from "./TaikoData.sol";
 abstract contract TaikoEvents {
     // The following events must match the definitions in corresponding L1
     // libraries.
-    event BlockProposed(uint256 indexed id, TaikoData.BlockMetadata meta);
+    event BlockProposed(
+        uint256 indexed id, TaikoData.BlockMetadata meta, uint64 blockFee
+    );
 
     event BlockProven(
         uint256 indexed id,

packages/protocol/contracts/L1/TaikoL1.sol

@@ -13,6 +13,7 @@ import { Proxied } from "../common/Proxied.sol";
 import { LibEthDepositing } from "./libs/LibEthDepositing.sol";
 import { LibProposing } from "./libs/LibProposing.sol";
 import { LibProving } from "./libs/LibProving.sol";
+import { LibTkoDistribution } from "./libs/LibTkoDistribution.sol";
 import { LibUtils } from "./libs/LibUtils.sol";
 import { LibVerifying } from "./libs/LibVerifying.sol";
 import { TaikoConfig } from "./TaikoConfig.sol";
@@ -155,6 +156,19 @@ contract TaikoL1 is
         });
     }
 
+    // From proposer side - same way paying the fees - and saving gas.
+    function depositTaikoToken(uint256 amount) external nonReentrant {
+        LibTkoDistribution.depositTaikoToken(
+            state, AddressResolver(this), amount
+        );
+    }
+
+    function withdrawTaikoToken(uint256 amount) external nonReentrant {
+        LibTkoDistribution.withdrawTaikoToken(
+            state, AddressResolver(this), amount
+        );
+    }
+
     function canDepositEthToL2(uint256 amount) public view returns (bool) {
         return LibEthDepositing.canDepositEthToL2({
             state: state,

packages/protocol/contracts/L1/libs/LibProposing.sol

@@ -25,7 +25,9 @@ library LibProposing {
     using LibUtils for TaikoData.State;
     using SafeCastUpgradeable for uint256;
 
-    event BlockProposed(uint256 indexed id, TaikoData.BlockMetadata meta);
+    event BlockProposed(
+        uint256 indexed id, TaikoData.BlockMetadata meta, uint64 blockFee
+    );
 
     error L1_BLOCK_ID();
     error L1_INSUFFICIENT_TOKEN();
@@ -137,15 +139,17 @@ library LibProposing {
             );
         }
 
-        IMintableERC20(resolver.resolve("taiko_token", false)).burn({
-            from: msg.sender,
-            amount: getBlockFee(state, config, meta.gasLimit)
-        });
+        uint64 blockFee = getBlockFee(state, config, meta.gasLimit);
 
-        emit BlockProposed(state.numBlocks, meta);
+        if (state.taikoTokenBalances[msg.sender] < blockFee) {
+            revert L1_INSUFFICIENT_TOKEN();
+        }
+
+        emit BlockProposed(state.numBlocks, meta, blockFee);
 
         unchecked {
             ++state.numBlocks;
+            state.taikoTokenBalances[msg.sender] -= blockFee;
         }
     }
 

packages/protocol/contracts/L1/libs/LibTkoDistribution.sol

@@ -0,0 +1,54 @@
+// SPDX-License-Identifier: MIT
+//  _____     _ _         _         _
+// |_   _|_ _(_) |_____  | |   __ _| |__ ___
+//   | |/ _` | | / / _ \ | |__/ _` | '_ (_-<
+//   |_|\__,_|_|_\_\___/ |____\__,_|_.__/__/
+
+pragma solidity ^0.8.20;
+
+import { AddressResolver } from "../../common/AddressResolver.sol";
+import { LibMath } from "../../libs/LibMath.sol";
+import { SafeCastUpgradeable } from
+    "@openzeppelin/contracts-upgradeable/utils/math/SafeCastUpgradeable.sol";
+import { TaikoData } from "../TaikoData.sol";
+import { TaikoToken } from "../TaikoToken.sol";
+import { LibFixedPointMath as Math } from
+    "../../thirdparty/LibFixedPointMath.sol";
+
+library LibTkoDistribution {
+    error L1_INSUFFICIENT_TOKEN();
+
+    function withdrawTaikoToken(
+        TaikoData.State storage state,
+        AddressResolver resolver,
+        uint256 amount
+    )
+        internal
+    {
+        uint256 balance = state.taikoTokenBalances[msg.sender];
+        if (balance < amount) revert L1_INSUFFICIENT_TOKEN();
+
+        unchecked {
+            state.taikoTokenBalances[msg.sender] -= amount;
+        }
+
+        TaikoToken(resolver.resolve("taiko_token", false)).mint(
+            msg.sender, amount
+        );
+    }
+
+    function depositTaikoToken(
+        TaikoData.State storage state,
+        AddressResolver resolver,
+        uint256 amount
+    )
+        internal
+    {
+        if (amount > 0) {
+            TaikoToken(resolver.resolve("taiko_token", false)).burn(
+                msg.sender, amount
+            );
+            state.taikoTokenBalances[msg.sender] += amount;
+        }
+    }
+}