Improved UpdateHistories library

oasysgames · Dec 21, 2022 · 3c04363 · 3c04363
1 parent 2548a79
commit 3c04363
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 30 deletions.
diff --git a/contracts/lib/UpdateHistories.sol b/contracts/lib/UpdateHistories.sol
@@ -4,44 +4,41 @@ pragma solidity 0.8.12;
 
 import { IEnvironment } from "../IEnvironment.sol";
 
+error PastEpoch();
+
 /**
  * @title UpdateHistories
  */
 library UpdateHistories {
-    function set(
-        uint256[] storage epochs,
-        uint256[] storage values,
-        uint256 nextEpoch,
-        uint256 value
-    ) internal {
-        extend(epochs, values, nextEpoch);
-        values[epochs.length - 1] = value;
-    }
-
     function add(
         uint256[] storage epochs,
         uint256[] storage values,
-        uint256 nextEpoch,
+        uint256 epoch,
         uint256 value
     ) internal {
-        extend(epochs, values, nextEpoch);
-        values[epochs.length - 1] += value;
+        uint256 length = epochs.length;
+        if (length > 1 && epoch < epochs[length - 1] && epoch < epochs[length - 2]) revert PastEpoch();
+
+        uint256 pos = extend(epochs, values, epoch);
+        length = epochs.length;
+        for (; pos < length; pos++) {
+            values[pos] += value;
+        }
     }
 
     function sub(
         uint256[] storage epochs,
         uint256[] storage values,
-        uint256 nextEpoch,
+        uint256 epoch,
         uint256 value
     ) internal returns (uint256) {
-        extend(epochs, values, nextEpoch);
-
         uint256 length = epochs.length;
-        uint256 balance = values[length - 1];
+        if (length == 0 || epoch < epochs[length - 1]) revert PastEpoch();
+
+        uint256 pos = extend(epochs, values, epoch);
+        uint256 balance = values[pos];
         value = value <= balance ? value : balance;
-        if (value > 0) {
-            values[length - 1] -= value;
-        }
+        if (value > 0) values[pos] -= value;
         return value;
     }
 
@@ -53,6 +50,7 @@ library UpdateHistories {
         uint256 length = epochs.length;
         if (length == 0 || epochs[0] > epoch) return 0;
         if (epochs[length - 1] <= epoch) return values[length - 1];
+        if (length > 1 && epochs[length - 2] <= epoch) return values[length - 2];
         uint256 idx = sBinarySearch(epochs, epoch, 0, length);
         return values[idx];
     }
@@ -64,26 +62,47 @@ library UpdateHistories {
     ) internal pure returns (IEnvironment.EnvironmentValue memory) {
         uint256 length = epochs.length;
         if (epochs[length - 1] <= epoch) return values[length - 1];
+        if (length > 1 && epochs[length - 2] <= epoch) return values[length - 2];
         uint256 idx = mBinarySearch(epochs, epoch, 0, length);
         return values[idx];
     }
 
     function extend(
         uint256[] storage epochs,
         uint256[] storage values,
-        uint256 nextEpoch
-    ) internal {
+        uint256 epoch
+    ) internal returns (uint256 pos) {
         uint256 length = epochs.length;
+
+        // first time
         if (length == 0) {
-            epochs.push(nextEpoch);
-            values.push();
-            return;
+            epochs.push(epoch);
+            values.push(0);
+            return 0;
+        }
+
+        uint256 lastPos = length - 1;
+        uint256 lastEpoch = epochs[lastPos];
+
+        // same as last epoch
+        if (epoch == lastEpoch) return lastPos;
+
+        // future epoch
+        if (epoch > lastEpoch) {
+            epochs.push(epoch);
+            values.push(values[lastPos]);
+            return lastPos + 1;
         }
 
-        uint256 lastEpoch = epochs[length - 1];
-        if (lastEpoch != nextEpoch) {
-            epochs.push(nextEpoch);
-            values.push(values[length - 1]);
+        // previous epoch
+        if (lastPos > 0 && epoch == epochs[lastPos - 1]) {
+            return lastPos - 1;
+        } else {
+            epochs.push(epochs[lastPos]);
+            values.push(values[lastPos]);
+            epochs[lastPos] = epoch;
+            values[lastPos] = lastPos == 0 ? 0 : values[lastPos - 1];
+            return lastPos;
         }
     }
 

diff --git a/test/StakeManager.spec.ts b/test/StakeManager.spec.ts
@@ -936,7 +936,6 @@ describe('StakeManager', () => {
       const tx2 = await staker1.unstakeV2(Token.wOAS, validator1, '2.5')
       await expect(tx1).to.emit(stakeManager, 'UnstakedV2').withArgs(staker1.address, validator1.owner.address, 0)
       await expect(tx2).to.emit(stakeManager, 'UnstakedV2').withArgs(staker1.address, validator1.owner.address, 1)
-      await expect(staker1.unstakeV2(Token.sOAS, validator1, '1')).to.revertedWith('NoAmount')
 
       await expectBalance(stakeManager, '515', '5', '10')
       await expectBalance(staker1.signer, '7985', '995', '990')
@@ -1005,6 +1004,7 @@ describe('StakeManager', () => {
       await staker1.unstakeV2(Token.wOAS, validator1, '9999')
       await staker1.unstakeV2(Token.OAS, validator2, '5')
       await staker1.unstakeV2(Token.sOAS, validator2, '5')
+      await expect(staker1.unstakeV2(Token.wOAS, validator1, '1')).to.revertedWith('NoAmount')
 
       expect(await staker1.getLockedUnstakeCount()).to.equal(10)