Opening a new shard fix

contracts/DecentralizedKV.sol

@@ -114,8 +114,8 @@ contract DecentralizedKV is OwnableUpgradeable {
         return _upfrontPayment(block.timestamp);
     }
 
-    /// @notice Checks before appending the key-value.
-    function _prepareAppend(uint256 _batchSize) internal virtual {
+    /// @notice Checks while appending the key-value.
+    function _checkAppend(uint256 _batchSize) internal virtual {
         require(msg.value >= upfrontPayment() * _batchSize, "DecentralizedKV: not enough batch payment");
     }
 
@@ -152,7 +152,7 @@ contract DecentralizedKV is OwnableUpgradeable {
             res[i] = paddr.kvIdx;
         }
 
-        _prepareAppend(batchPaymentSize);
+        _checkAppend(batchPaymentSize);
 
         return res;
     }

contracts/StorageContract.sol

@@ -145,44 +145,45 @@ abstract contract StorageContract is DecentralizedKV {
     /// @notice People can sent ETH to the contract.
     function sendValue() public payable {}
 
-    /// @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.
-            // (TODO): Setup shard difficulty as current difficulty / factor?
-            if (shardId != 0) {
-                // shard0 is already opened in constructor
-                infos[shardId].lastMineTime = _timestamp;
-            }
-        }
+    /// @notice Upfront payment for the next insertion
+    function upfrontPayment() public view virtual override returns (uint256) {
+        return _upfrontPaymentInBatch(kvEntryCount, 1);
+    }
 
-        require(
-            msg.value >= _upfrontPayment(infos[shardId].lastMineTime) * _batchSize,
-            "StorageContract: not enough batch payment"
-        );
+    /// @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 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
-        // 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);
+    /// @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 {
-            return _upfrontPayment(infos[shardId].lastMineTime);
+            totalPayment += _upfrontPayment(infos[shardId].lastMineTime) * _batchSize;
         }
+        return totalPayment;
     }
 
     /// @inheritdoc DecentralizedKV
-    function _prepareAppend(uint256 _batchSize) internal virtual override {
-        return _prepareAppendWithTimestamp(block.timestamp, _batchSize);
+    function _checkAppend(uint256 _batchSize) internal virtual override {
+        uint256 kvEntryCountPrev = kvEntryCount - _batchSize; // kvEntryCount already increased
+        uint256 totalPayment = _upfrontPaymentInBatch(kvEntryCountPrev, _batchSize);
+        require(msg.value >= totalPayment, "StorageContract: not enough batch payment");
+
+        uint256 shardId = kvEntryCount >> SHARD_ENTRY_BITS; // shard id after the batch
+        if (shardId > (kvEntryCountPrev >> SHARD_ENTRY_BITS)) {
+            // Open a new shard and mark the shard is ready to mine.
+            // (TODO): Setup shard difficulty as current difficulty / factor?
+            infos[shardId].lastMineTime = block.timestamp;
+        }
     }
 
     /// @notice Verify the samples of the BLOBs by the miner (storage provider) including

contracts/test/EthStorageContractTest.t.sol

@@ -35,7 +35,7 @@ contract EthStorageContractTest is Test {
 
         uint256 insufficientCost = storageContract.upfrontPayment() - 1;
 
-        // Expect the specific revert reason from _prepareAppend due to insufficient msg.value
+        // Expect the specific revert reason from _checkAppend due to insufficient msg.value
         vm.expectRevert("StorageContract: not enough batch payment");
         storageContract.putBlob{value: insufficientCost}(key, blobIdx, length);
 
@@ -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);
+    }
 }