Opening a new shard fix

contracts/StorageContract.sol

@@ -147,39 +147,53 @@ abstract contract StorageContract is DecentralizedKV {
 
     /// @notice Checks the payment using the last mine time.
     function _prepareAppendWithTimestamp(uint256 _timestamp, uint256 _batchSize) internal {
-        uint256 totalEntries = kvEntryCount + 1; // include the one to be put
-        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS; // shard id of the new KV
-        if ((totalEntries % (1 << SHARD_ENTRY_BITS)) == 1) {
-            // Open a new shard if the KV is the first one of the shard
-            // and mark the shard is ready to mine.
+        uint256 kvEntryCountOld = kvEntryCount - _batchSize; // kvEntryCount already increased
+        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?
-            if (shardId != 0) {
-                // shard0 is already opened in constructor
-                infos[shardId].lastMineTime = _timestamp;
-            }
+            infos[shardId].lastMineTime = _timestamp;
         }
-
-        require(
-            msg.value >= _upfrontPayment(infos[shardId].lastMineTime) * _batchSize,
-            "StorageContract: not enough batch payment"
-        );
     }
 
     /// @notice Upfront payment for the next insertion
     function upfrontPayment() public view virtual override returns (uint256) {
         uint256 totalEntries = kvEntryCount + 1; // include the one to be put
-        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS; // shard id of the new KV
+        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS;
         // shard0 is already opened in constructor
         if ((totalEntries % (1 << SHARD_ENTRY_BITS)) == 1 && shardId != 0) {
             // Open a new shard if the KV is the first one of the shard
-            // and mark the shard is ready to mine.
             // (TODO): Setup shard difficulty as current difficulty / factor?
             return _upfrontPayment(block.timestamp);
         } else {
             return _upfrontPayment(infos[shardId].lastMineTime);
         }
     }
 
+    /// @notice Upfront payment for a batch insertion
+    /// @param _batchSize The blob count for a batch insertion.
+    /// @return The total payment for a batch insertion.
+    function upfrontPaymentInBatch(uint256 _batchSize) public view returns (uint256) {
+        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 ((totalEntries >> SHARD_ENTRY_BITS) > shardId) {
+            uint256 kvCountNew = totalEntries % (1 << SHARD_ENTRY_BITS);
+            totalPayment += _upfrontPayment(block.timestamp) * kvCountNew;
+            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * (_batchSize - kvCountNew);
+        } else {
+            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * _batchSize;
+        }
+        return totalPayment;
+    }
+
     /// @inheritdoc DecentralizedKV
     function _prepareAppend(uint256 _batchSize) internal virtual override {
         return _prepareAppendWithTimestamp(block.timestamp, _batchSize);

contracts/test/EthStorageContractTest.t.sol

@@ -75,7 +75,7 @@ contract EthStorageContractTest is Test {
         storageContract.putBlobs{value: insufficientCost}(keys, blobIdxs, lengths);
 
         // Enough storage cost
-        uint256 sufficientCost = 2 * storageContract.upfrontPayment();
+        uint256 sufficientCost = storageContract.upfrontPaymentInBatch(2);
         storageContract.putBlobs{value: sufficientCost}(keys, blobIdxs, lengths);
 
         assertEq(storageContract.kvEntryCount(), 2);
@@ -99,4 +99,27 @@ contract EthStorageContractTest is Test {
         assertEq(storageContract.size(bytes32(uint256(1))), 20);
         assertEq(storageContract.size(bytes32(uint256(2))), 30);
     }
+
+    function testPutBlobsXshard() public {
+        uint256 size = 6;
+        bytes32[] memory keys = new bytes32[](size);
+        uint256[] memory blobIdxs = new uint256[](size);
+        uint256[] memory lengths = new uint256[](size);
+        for (uint256 i = 0; i < size; i++) {
+            keys[i] = bytes32(uint256(i));
+            blobIdxs[i] = i;
+            lengths[i] = 10 + i * 10;
+        }
+
+        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
+        storageContract.putBlobs{value: sufficientCost}(keys, blobIdxs, lengths);
+        assertEq(storageContract.kvEntryCount(), size);
+    }
 }