refactor

ethstorage · Aug 19, 2024 · f484662 · f484662
1 parent 6437963
commit f484662
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 19 deletions.
diff --git a/contracts/StorageContract.sol b/contracts/StorageContract.sol
@@ -150,21 +150,16 @@ abstract contract StorageContract is DecentralizedKV {
 
     /// @notice Checks the payment using the last mine time.
     function _prepareAppendWithTimestamp(uint256 _timestamp, uint256 _batchSize) internal {
-        uint256 shardId = (kvEntryCount - _batchSize) >> SHARD_ENTRY_BITS;
-        uint256 shardIdNew = kvEntryCount >> SHARD_ENTRY_BITS; // shard id after the batch
-        uint256 totalPayment = 0;
-        if (shardIdNew > shardId) {
-            uint256 kvCountNew = kvEntryCount % (1 << SHARD_ENTRY_BITS);
-            totalPayment += _upfrontPayment(block.timestamp) * kvCountNew;
-            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * (_batchSize - kvCountNew);
+        uint256 kvEntryCountOld = kvEntryCount - _batchSize;
+        uint256 totalPayment = _upfrontPaymentInBatch(kvEntryCountOld, _batchSize);
+        require(msg.value >= totalPayment, "StorageContract: not enough batch payment");
+        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS; // shard id after the batch
+        if (shardId > (kvEntryCountOld >> SHARD_ENTRY_BITS)) {
             // Open a new shard and mark the shard is ready to mine.
             // (TODO): Setup shard difficulty as current difficulty / factor?
-            infos[shardIdNew].lastMineTime = _timestamp;
-        } else {
-            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * _batchSize;
+            infos[shardId].lastMineTime = _timestamp;
         }
 
-        require(msg.value >= totalPayment, "StorageContract: not enough batch payment");
     }
 
     /// @notice Upfront payment for the next insertion
@@ -183,11 +178,15 @@ abstract contract StorageContract is DecentralizedKV {
 
     /// @notice Upfront payment for a batch insertion
     function upfrontPaymentInBatch(uint256 _batchSize) public view virtual override returns (uint256) {
-        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS;
-        uint256 totalEntries = kvEntryCount + _batchSize; // include the batch to be put
-        uint256 shardIdNew = totalEntries >> SHARD_ENTRY_BITS; // shard id after the batch
+        return _upfrontPaymentInBatch(kvEntryCount, _batchSize);
+    }
+
+    /// @notice Upfront payment for a batch insertion
+    function _upfrontPaymentInBatch(uint256 _kvEntryCount, uint256 _batchSize) private view returns (uint256) {
+        uint256 shardId = _kvEntryCount >> SHARD_ENTRY_BITS;
+        uint256 totalEntries = _kvEntryCount + _batchSize; // include the batch to be put
         uint256 totalPayment = 0;
-        if (shardIdNew > shardId) {
+        if ((totalEntries >> SHARD_ENTRY_BITS) > shardId) {
             uint256 kvCountNew = totalEntries % (1 << SHARD_ENTRY_BITS);
             totalPayment += _upfrontPayment(block.timestamp) * kvCountNew;
             totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * (_batchSize - kvCountNew);

diff --git a/contracts/test/EthStorageContractTest.t.sol b/contracts/test/EthStorageContractTest.t.sol
@@ -101,7 +101,7 @@ contract EthStorageContractTest is Test {
     }
 
     function testPutBlobsXshard() public {
-        uint256 size = 5;
+        uint256 size = 6;
         bytes32[] memory keys = new bytes32[](size);
         uint256[] memory blobIdxs = new uint256[](size);
         uint256[] memory lengths = new uint256[](size);
@@ -112,15 +112,14 @@ contract EthStorageContractTest is Test {
         }
         vm.warp(0);
 
-        uint256 insufficientCost = storageContract.upfrontPaymentInBatch(size) - 1;
+        uint256 sufficientCost = storageContract.upfrontPaymentInBatch(size);
+        uint256 insufficientCost = sufficientCost - 1;
         // Expect the specific revert reason from _prepareBatchAppend due to insufficient msg.value
         vm.expectRevert("StorageContract: not enough batch payment");
         storageContract.putBlobs{value: insufficientCost}(keys, blobIdxs, lengths);
 
         // Enough storage cost
-        uint256 sufficientCost = storageContract.upfrontPaymentInBatch(size);
         storageContract.putBlobs{value: sufficientCost}(keys, blobIdxs, lengths);
-
         assertEq(storageContract.kvEntryCount(), size);
     }
 }