Added overflow handling for TWAB timestamp

contracts/DrawCalculator.sol

@@ -239,16 +239,16 @@ contract DrawCalculator is IDrawCalculator, Ownable {
         IDrawBeacon.Draw[] memory _draws,
         IPrizeDistributionHistory.PrizeDistribution[] memory _prizeDistributions
     ) internal view returns (uint256[] memory) {
-        uint32[] memory _timestampsWithStartCutoffTimes = new uint32[](_draws.length);
-        uint32[] memory _timestampsWithEndCutoffTimes = new uint32[](_draws.length);
+        uint64[] memory _timestampsWithStartCutoffTimes = new uint64[](_draws.length);
+        uint64[] memory _timestampsWithEndCutoffTimes = new uint64[](_draws.length);
 
         // generate timestamps with draw cutoff offsets included
         for (uint32 i = 0; i < _draws.length; i++) {
             unchecked {
-                _timestampsWithStartCutoffTimes[i] = uint32(
+                _timestampsWithStartCutoffTimes[i] = (
                     _draws[i].timestamp - _prizeDistributions[i].startTimestampOffset
                 );
-                _timestampsWithEndCutoffTimes[i] = uint32(
+                _timestampsWithEndCutoffTimes[i] = (
                     _draws[i].timestamp - _prizeDistributions[i].endTimestampOffset
                 );
             }

contracts/Ticket.sol

@@ -75,7 +75,7 @@ contract Ticket is ControlledToken, ITicket {
     }
 
     /// @inheritdoc ITicket
-    function getBalanceAt(address _user, uint256 _target) external view override returns (uint256) {
+    function getBalanceAt(address _user, uint64 _target) external view override returns (uint256) {
         TwabLib.Account storage account = userTwabs[_user];
 
         return
@@ -90,25 +90,25 @@ contract Ticket is ControlledToken, ITicket {
     /// @inheritdoc ITicket
     function getAverageBalancesBetween(
         address _user,
-        uint32[] calldata _startTimes,
-        uint32[] calldata _endTimes
+        uint64[] calldata _startTimes,
+        uint64[] calldata _endTimes
     ) external view override returns (uint256[] memory) {
         return _getAverageBalancesBetween(userTwabs[_user], _startTimes, _endTimes);
     }
 
     /// @inheritdoc ITicket
     function getAverageTotalSuppliesBetween(
-        uint32[] calldata _startTimes,
-        uint32[] calldata _endTimes
+        uint64[] calldata _startTimes,
+        uint64[] calldata _endTimes
     ) external view override returns (uint256[] memory) {
         return _getAverageBalancesBetween(totalSupplyTwab, _startTimes, _endTimes);
     }
 
     /// @inheritdoc ITicket
     function getAverageBalanceBetween(
         address _user,
-        uint256 _startTime,
-        uint256 _endTime
+        uint64 _startTime,
+        uint64 _endTime
     ) external view override returns (uint256) {
         TwabLib.Account storage account = userTwabs[_user];
 
@@ -123,7 +123,7 @@ contract Ticket is ControlledToken, ITicket {
     }
 
     /// @inheritdoc ITicket
-    function getBalancesAt(address _user, uint32[] calldata _targets)
+    function getBalancesAt(address _user, uint64[] calldata _targets)
         external
         view
         override
@@ -139,7 +139,7 @@ contract Ticket is ControlledToken, ITicket {
             _balances[i] = TwabLib.getBalanceAt(
                 twabContext.twabs,
                 details,
-                _targets[i],
+                uint32(_targets[i]),
                 uint32(block.timestamp)
             );
         }
@@ -148,18 +148,18 @@ contract Ticket is ControlledToken, ITicket {
     }
 
     /// @inheritdoc ITicket
-    function getTotalSupplyAt(uint32 _target) external view override returns (uint256) {
+    function getTotalSupplyAt(uint64 _target) external view override returns (uint256) {
         return
             TwabLib.getBalanceAt(
                 totalSupplyTwab.twabs,
                 totalSupplyTwab.details,
-                _target,
+                uint32(_target),
                 uint32(block.timestamp)
             );
     }
 
     /// @inheritdoc ITicket
-    function getTotalSuppliesAt(uint32[] calldata _targets)
+    function getTotalSuppliesAt(uint64[] calldata _targets)
         external
         view
         override
@@ -174,7 +174,7 @@ contract Ticket is ControlledToken, ITicket {
             totalSupplies[i] = TwabLib.getBalanceAt(
                 totalSupplyTwab.twabs,
                 details,
-                _targets[i],
+                uint32(_targets[i]),
                 uint32(block.timestamp)
             );
         }
@@ -247,8 +247,8 @@ contract Ticket is ControlledToken, ITicket {
      */
     function _getAverageBalancesBetween(
         TwabLib.Account storage _account,
-        uint32[] calldata _startTimes,
-        uint32[] calldata _endTimes
+        uint64[] calldata _startTimes,
+        uint64[] calldata _endTimes
     ) internal view returns (uint256[] memory) {
         require(_startTimes.length == _endTimes.length, "Ticket/start-end-times-length-match");
 
@@ -261,8 +261,8 @@ contract Ticket is ControlledToken, ITicket {
             averageBalances[i] = TwabLib.getAverageBalanceBetween(
                 _account.twabs,
                 accountDetails,
-                _startTimes[i],
-                _endTimes[i],
+                uint32(_startTimes[i]),
+                uint32(_endTimes[i]),
                 currentTimestamp
             );
         }

contracts/interfaces/ITicket.sol

@@ -127,15 +127,15 @@ interface ITicket is IControlledToken {
      * @param timestamp Timestamp at which we want to retrieve the TWAB balance.
      * @return The TWAB balance at the given timestamp.
      */
-    function getBalanceAt(address user, uint256 timestamp) external view returns (uint256);
+    function getBalanceAt(address user, uint64 timestamp) external view returns (uint256);
 
     /**
      * @notice Retrieves `user` TWAB balances.
      * @param user Address of the user whose TWABs are being fetched.
      * @param timestamps Timestamps range at which we want to retrieve the TWAB balances.
      * @return `user` TWAB balances.
      */
-    function getBalancesAt(address user, uint32[] calldata timestamps)
+    function getBalancesAt(address user, uint64[] calldata timestamps)
         external
         view
         returns (uint256[] memory);
@@ -149,8 +149,8 @@ interface ITicket is IControlledToken {
      */
     function getAverageBalanceBetween(
         address user,
-        uint256 startTime,
-        uint256 endTime
+        uint64 startTime,
+        uint64 endTime
     ) external view returns (uint256);
 
     /**
@@ -162,23 +162,23 @@ interface ITicket is IControlledToken {
      */
     function getAverageBalancesBetween(
         address user,
-        uint32[] calldata startTimes,
-        uint32[] calldata endTimes
+        uint64[] calldata startTimes,
+        uint64[] calldata endTimes
     ) external view returns (uint256[] memory);
 
     /**
      * @notice Retrieves the total supply TWAB balance at the given timestamp.
      * @param timestamp Timestamp at which we want to retrieve the total supply TWAB balance.
      * @return The total supply TWAB balance at the given timestamp.
      */
-    function getTotalSupplyAt(uint32 timestamp) external view returns (uint256);
+    function getTotalSupplyAt(uint64 timestamp) external view returns (uint256);
 
     /**
      * @notice Retrieves the total supply TWAB balance between the given timestamps range.
      * @param timestamps Timestamps range at which we want to retrieve the total supply TWAB balance.
      * @return Total supply TWAB balances.
      */
-    function getTotalSuppliesAt(uint32[] calldata timestamps)
+    function getTotalSuppliesAt(uint64[] calldata timestamps)
         external
         view
         returns (uint256[] memory);
@@ -190,7 +190,7 @@ interface ITicket is IControlledToken {
      * @return The average total supplies held during the time frame.
      */
     function getAverageTotalSuppliesBetween(
-        uint32[] calldata startTimes,
-        uint32[] calldata endTimes
+        uint64[] calldata startTimes,
+        uint64[] calldata endTimes
     ) external view returns (uint256[] memory);
 }

contracts/libraries/OverflowSafeComparatorLib.sol

@@ -2,6 +2,8 @@
 
 pragma solidity 0.8.6;
 
+import "hardhat/console.sol";
+
 /// @title OverflowSafeComparatorLib library to share comparator functions between contracts
 /// @dev Code taken from Uniswap V3 Oracle.sol: https://github.com/Uniswap/v3-core/blob/3e88af408132fc957e3e406f65a0ce2b1ca06c3d/contracts/libraries/Oracle.sol
 /// @author PoolTogether Inc.
@@ -36,7 +38,8 @@ library OverflowSafeComparatorLib {
         uint32 _a,
         uint32 _b,
         uint32 _timestamp
-    ) internal pure returns (bool) {
+    ) internal view returns (bool) {
+
         // No need to adjust if there hasn't been an overflow
         if (_a <= _timestamp && _b <= _timestamp) return _a <= _b;
 

contracts/libraries/TwabLib.sol

@@ -228,7 +228,7 @@ library TwabLib {
         );
 
         // Difference in amount / time
-        return (endTwab.amount - startTwab.amount) / (endTwab.timestamp - startTwab.timestamp);
+        return (endTwab.amount - startTwab.amount) / OverflowSafeComparatorLib.checkedSub(endTwab.timestamp, startTwab.timestamp, _currentTime);
     }
 
     /** @notice Searches TWAB history and calculate the difference between amount(s)/timestamp(s) to return average balance
@@ -278,7 +278,7 @@ library TwabLib {
         // The time-weighted average balance uses time measured between two epoch timestamps as
         // a constaint on the measurement when calculating the time weighted average balance.
         return
-            (afterOrAt.amount - beforeOrAt.amount) / (afterOrAt.timestamp - beforeOrAt.timestamp);
+            (afterOrAt.amount - beforeOrAt.amount) / OverflowSafeComparatorLib.checkedSub(afterOrAt.timestamp, beforeOrAt.timestamp, _currentTime);
     }
 
     /** @notice Calculates a user TWAB for a target timestamp using the historical TWAB records.
@@ -339,7 +339,7 @@ library TwabLib {
             );
 
         uint224 heldBalance = (afterOrAtStart.amount - beforeOrAtStart.amount) /
-            (afterOrAtStart.timestamp - beforeOrAtStart.timestamp);
+            OverflowSafeComparatorLib.checkedSub(afterOrAtStart.timestamp, beforeOrAtStart.timestamp, _time);
 
         return _computeNextTwab(beforeOrAtStart, heldBalance, _targetTimestamp);
     }
@@ -368,6 +368,7 @@ library TwabLib {
     }
 
     /// @notice Sets a new TWAB Observation at the next available index and returns the new account details.
+    /// @dev Note that if _currentTime is before the last observation timestamp, it appears as an overflow
     /// @param _twabs The twabs array to insert into
     /// @param _accountDetails The current account details
     /// @param _currentTime The current time
@@ -387,7 +388,6 @@ library TwabLib {
         )
     {
         (, ObservationLib.Observation memory _newestTwab) = newestTwab(_twabs, _accountDetails);
-        require(_currentTime >= _newestTwab.timestamp, "TwabLib/twab-time-monotonic");
 
         // if we're in the same block, return
         if (_newestTwab.timestamp == _currentTime) {

contracts/test/libraries/OverflowSafeComparatorLibHarness.sol

@@ -19,7 +19,7 @@ contract OverflowSafeComparatorLibHarness {
         uint32 _a,
         uint32 _b,
         uint32 _timestamp
-    ) external pure returns (bool) {
+    ) external view returns (bool) {
         return _a.lte(_b, _timestamp);
     }
 

test/Ticket.test.ts

@@ -965,4 +965,31 @@ describe('Ticket', () => {
 
         })
     })
+
+    context('when the timestamp overflows', () => {
+
+        let overflowMintTimestamp: number
+
+        beforeEach(async () => {
+            await ticket.mint(wallet1.address, toWei('100'))
+            const timestamp = (await ethers.provider.getBlock('latest')).timestamp
+            const timeUntilOverflow = (2**32 - timestamp)
+            await increaseTime(timeUntilOverflow)
+            await ticket.mint(wallet1.address, toWei('100'))
+            overflowMintTimestamp = (await ethers.provider.getBlock('latest')).timestamp
+            await increaseTime(100)
+        })
+
+        describe('getAverageBalanceBetween()', () => {
+            it('should function across overflow boundary', async () => {
+                expect(await ticket.getAverageBalanceBetween(wallet1.address, overflowMintTimestamp-100, overflowMintTimestamp+100)).to.equal(toWei('150'))
+            })
+        })
+
+        describe('getBalanceAt', () => {
+            it('should function across overflow boundary', async () => {
+                expect(await ticket.getBalanceAt(wallet1.address, overflowMintTimestamp)).to.equal(toWei('200'))
+            })
+        })
+    })
 });

test/TwabLibraryExposed.test.ts

@@ -41,14 +41,6 @@ describe('TwabLib', () => {
                 .withArgs([200, 1, 1], [0, timestamp], false);
         });
 
-        it('should require the timestamp to always increase', async () => {
-            await twabLib.increaseBalance(100, timestamp);
-
-            await expect(twabLib.increaseBalance(100, timestamp - 10)).to.be.revertedWith(
-                'TwabLib/twab-time-monotonic',
-            );
-        });
-
         it('should add second twab', async () => {
             await expect(twabLib.increaseBalance(100, timestamp))
                 .to.emit(twabLib, 'Updated')