add reward cap

contracts/StorageContract.sol

@@ -100,12 +100,9 @@ abstract contract StorageContract is DecentralizedKV {
     );
 
     /// @notice Constructs the StorageContract contract. Initializes the storage config.
-    constructor(
-        Config memory _config,
-        uint256 _startTime,
-        uint256 _storageCost,
-        uint256 _dcfFactor
-    ) DecentralizedKV(1 << _config.maxKvSizeBits, _startTime, _storageCost, _dcfFactor) {
+    constructor(Config memory _config, uint256 _startTime, uint256 _storageCost, uint256 _dcfFactor)
+        DecentralizedKV(1 << _config.maxKvSizeBits, _startTime, _storageCost, _dcfFactor)
+    {
         /* Assumptions */
         require(_config.shardSizeBits >= _config.maxKvSizeBits, "StorageContract: shardSize too small");
         require(_config.maxKvSizeBits >= SAMPLE_SIZE_BITS, "StorageContract: maxKvSize too small");
@@ -211,10 +208,11 @@ abstract contract StorageContract is DecentralizedKV {
     /// @param _shardId  The shard id.
     /// @param _minedTs  The mined timestamp.
     /// @return diff_ The difficulty of the shard.
-    function _calculateDiffAndInitHashSingleShard(
-        uint256 _shardId,
-        uint256 _minedTs
-    ) internal view returns (uint256 diff_) {
+    function _calculateDiffAndInitHashSingleShard(uint256 _shardId, uint256 _minedTs)
+        internal
+        view
+        returns (uint256 diff_)
+    {
         MiningLib.MiningInfo storage info = infos[_shardId];
         require(_minedTs >= info.lastMineTime, "StorageContract: minedTs too small");
         diff_ = MiningLib.expectedDiff(info, _minedTs, CUTOFF, DIFF_ADJ_DIVISOR, minimumDiff);
@@ -260,7 +258,12 @@ abstract contract StorageContract is DecentralizedKV {
             reward = _paymentIn(STORAGE_COST * (kvEntryCount % (1 << SHARD_ENTRY_BITS)), info.lastMineTime, _minedTs);
             // Additional prepaid for the last shard
             if (prepaidLastMineTime < _minedTs) {
-                reward += _paymentIn(prepaidAmount, prepaidLastMineTime, _minedTs);
+                uint256 prepaidAmountCap =
+                    STORAGE_COST * ((1 << SHARD_ENTRY_BITS) - kvEntryCount % (1 << SHARD_ENTRY_BITS));
+                if (prepaidAmountCap > prepaidAmount) {
+                    prepaidAmountCap = prepaidAmount;
+                }
+                reward += _paymentIn(prepaidAmountCap, prepaidLastMineTime, _minedTs);
                 updatePrepaidTime = true;
             }
         }
@@ -276,7 +279,7 @@ abstract contract StorageContract is DecentralizedKV {
     /// @return The mining reward.
     function miningReward(uint256 _shardId, uint256 _blockNumber) public view returns (uint256) {
         uint256 minedTs = block.timestamp - (block.number - _blockNumber) * 12;
-        (, , uint256 minerReward) = _miningReward(_shardId, minedTs);
+        (,, uint256 minerReward) = _miningReward(_shardId, minedTs);
         return minerReward;
     }
 

contracts/test/StorageContractTest.t.sol

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import "./TestStorageContract.sol";
+import "../StorageContract.sol";
+import "forge-std/Test.sol";
+import "forge-std/Vm.sol";
+
+contract StorageContractTest is Test {
+    uint256 constant STORAGE_COST = 1000;
+    uint256 constant SHARD_SIZE_BITS = 19;
+    uint256 constant MAX_KV_SIZE = 17;
+    uint256 constant PREPAID_AMOUNT = 2 * STORAGE_COST;
+    TestStorageContract storageContract;
+
+    function setUp() public {
+        storageContract = new TestStorageContract(
+            StorageContract.Config(MAX_KV_SIZE, SHARD_SIZE_BITS, 2, 0, 0, 0), 0, STORAGE_COST, 0
+        );
+        storageContract.initialize(0, PREPAID_AMOUNT, 0, address(0x1), address(0x1));
+    }
+
+    function testMiningReward() public {
+        // no key-value stored on EthStorage, only use prepaid amount as the reward
+        (,, uint256 reward) = storageContract.miningRewards(0, 1);
+        assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT, 0, 1));
+
+        // 1 key-value stored on EthStorage
+        storageContract.setKvEntryCount(1);
+        (,, reward) = storageContract.miningRewards(0, 1);
+        assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT + STORAGE_COST * 1, 0, 1));
+
+        // 2 key-value stored on EthStorage
+        storageContract.setKvEntryCount(2);
+        (,, reward) = storageContract.miningRewards(0, 1);
+        assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT + STORAGE_COST * 2, 0, 1));
+
+        // 3 key-value stored on EthStorage, but the reward is capped with 4 * STORAGE_COST
+        storageContract.setKvEntryCount(3);
+        (,, reward) = storageContract.miningRewards(0, 1);
+        assertEq(reward, storageContract.paymentIn(PREPAID_AMOUNT + STORAGE_COST * 2, 0, 1));
+    }
+}

contracts/test/TestStorageContract.sol

@@ -0,0 +1,44 @@
+// SPDX License Identifier: MIT
+pragma solidity ^0.8.0;
+
+import "../StorageContract.sol";
+
+contract TestStorageContract is StorageContract {
+    constructor(Config memory _config, uint256 _startTime, uint256 _storageCost, uint256 _dcfFactor)
+        StorageContract(_config, _startTime, _storageCost, _dcfFactor)
+    {}
+
+    function initialize(
+        uint256 _minimumDiff,
+        uint256 _prepaidAmount,
+        uint256 _nonceLimit,
+        address _treasury,
+        address _owner
+    ) public payable initializer {
+        __init_storage(_minimumDiff, _prepaidAmount, _nonceLimit, _treasury, _owner);
+    }
+
+    function verifySamples(
+        uint256 _startShardId,
+        bytes32 _hash0,
+        address _miner,
+        bytes32[] memory _encodedSamples,
+        uint256[] memory _masks,
+        bytes[] calldata _inclusiveProofs,
+        bytes[] calldata _decodeProof
+    ) public pure override returns (bytes32) {
+        return bytes32(0);
+    }
+
+    function setKvEntryCount(uint40 _kvEntryCount) public {
+        kvEntryCount = _kvEntryCount;
+    }
+
+    function paymentIn(uint256 _x, uint256 _fromTs, uint256 _toTs) public view returns (uint256) {
+        return _paymentIn(_x, _fromTs, _toTs);
+    }
+
+    function miningRewards(uint256 _shardId, uint256 _minedTs) public view returns (bool, uint256, uint256) {
+        return _miningReward(_shardId, _minedTs);
+    }
+}

foundry.toml

@@ -3,5 +3,5 @@ src = 'contracts'
 out = 'out'
 libs = ['node_modules', 'lib']
 test = 'test'
-cache_path  = 'cache_forge'
-evm_version = 'cancun'
+cache_path = 'cache_forge'
+evm_version = 'cancun'

package.json

@@ -31,7 +31,7 @@
   },
   "scripts": {
     "compile": "hardhat compile",
-    "test": "hardhat test",
+    "test": "hardhat test && forge test",
     "prettier:check": "prettier-check contracts/**/*.sol",
     "prettier:fix": "prettier --write contracts/**/*.sol test/**/*.js scripts/**/*.js",
     "deploy": "npx hardhat run scripts/deploy.js --network sepolia",
@@ -42,4 +42,4 @@
       "packages/arb-shared-dependencies"
     ]
   }
-}
+}