From f4846624605cdbe0fa699973bc594112db4c1ec2 Mon Sep 17 00:00:00 2001
From: syntrust <lundeng@quarkchain.org>
Date: Mon, 19 Aug 2024 18:35:58 +0800
Subject: [PATCH] refactor

---
 contracts/StorageContract.sol               | 29 ++++++++++-----------
 contracts/test/EthStorageContractTest.t.sol |  7 +++--
 2 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/contracts/StorageContract.sol b/contracts/StorageContract.sol
index e845722..e80dfea 100644
--- 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
index 059d5bd..0f9466f 100644
--- 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);
     }
 }