test(invariant): include delay in each period

README.md

@@ -184,4 +184,4 @@ Currently, it's not possible to address this precision problem entirely.
 
 14. if $isVoided = true \implies isPaused = true$, $ra = 0$ and $ud = 0$
 
-15. if $isVoided = false \implies \text{amount streamed} - (td + \text{amount withdrawn}) \le 1$.
+15. if $isVoided = false \implies \text{amount streamed with delay} = td + \text{amount withdrawn}$.

test/invariant/Flow.t.sol

@@ -285,37 +285,44 @@ contract Flow_Invariant_Test is Base_Test {
         }
     }
 
-    /// @dev For non-voided streams, the difference between the total amount streamed and the sum of total debt and
-    /// total withdrawn should never exceed 1. This is indirectly checking that withdrawals do not cause the streamed
-    /// amount to deviate from the theoretical streamed amount by more than 1.
-    function invariant_TotalStreamedApproxEqTotalDebtPlusWithdrawn() external view {
+    /// @dev For non-voided streams, the difference between the total amount streamed adjusted for delay and the sum of
+    /// total debt and total withdrawn should be equal. Also, total streamed amount with delay must never exceed total
+    /// streamed amount without delay.
+    function invariant_TotalStreamedWithDelayEqTotalDebtPlusWithdrawn() external view {
         uint256 lastStreamId = flowStore.lastStreamId();
         for (uint256 i = 0; i < lastStreamId; ++i) {
             uint256 streamId = flowStore.streamIds(i);
 
             // Skip the voided streams.
             if (!flow.isVoided(streamId)) {
-                uint256 totalStreamedAmount =
-                    calculateTotalStreamedAmount(flowStore.streamIds(i), flow.getTokenDecimals(streamId));
+                (uint256 totalStreamedAmount, uint256 totalStreamedAmountWithDelay) =
+                    calculateTotalStreamedAmounts(flowStore.streamIds(i), flow.getTokenDecimals(streamId));
 
-                assertLe(
-                    totalStreamedAmount - flow.totalDebtOf(streamId) - flowStore.withdrawnAmounts(streamId),
-                    1,
-                    "Invariant violation: total debt - streamed amount - withdrawn amount > 1"
+                assertGe(
+                    totalStreamedAmount,
+                    totalStreamedAmountWithDelay,
+                    "Invariant violation: total streamed amount without delay >= total streamed amount with delay"
+                );
+
+                assertEq(
+                    totalStreamedAmountWithDelay,
+                    flow.totalDebtOf(streamId) + flowStore.withdrawnAmounts(streamId),
+                    "Invariant violation: total streamed amount with delay = total debt + withdrawn amount"
                 );
             }
         }
     }
 
-    /// @dev Calculates the total streamed amount iterating over each period.
-    function calculateTotalStreamedAmount(
+    /// @dev Calculates the total streamed amounts by iterating over each period.
+    function calculateTotalStreamedAmounts(
         uint256 streamId,
         uint8 decimals
     )
         internal
         view
-        returns (uint256 totalStreamedAmount)
+        returns (uint256 totalStreamedAmount, uint256 totalStreamedAmountWithDelay)
     {
+        uint256 totalDelayedAmount;
         uint256 periodsCount = flowStore.getPeriods(streamId).length;
 
         for (uint256 i = 0; i < periodsCount; ++i) {
@@ -324,9 +331,14 @@ contract Flow_Invariant_Test is Base_Test {
             // If end time is 0, it means the current period is still active.
             uint40 elapsed = period.end > 0 ? period.end - period.start : uint40(block.timestamp) - period.start;
 
-            totalStreamedAmount += period.ratePerSecond * elapsed;
+            // Calculate the total streamed amount for the current period.
+            totalStreamedAmount += getDescaledAmount(period.ratePerSecond * elapsed, decimals);
+
+            // Calculate the total delayed amount for the current period.
+            totalDelayedAmount += getDescaledAmount(period.delay * period.ratePerSecond, decimals);
         }
 
-        return totalStreamedAmount / 10 ** (18 - decimals);
+        // Calculate the total streamed amount with delay.
+        totalStreamedAmountWithDelay = totalStreamedAmount - totalDelayedAmount;
     }
 }

test/invariant/handlers/FlowHandler.sol

@@ -95,13 +95,16 @@ contract FlowHandler is BaseHandler {
             vm.assume(newRatePerSecond.unwrap() > mvt / 100 && newRatePerSecond.unwrap() <= 1e18);
         }
 
+        uint128 previousRatePerSecond = flow.getRatePerSecond(currentStreamId).unwrap();
+
         // The rate per second must be different from the current rate per second.
-        vm.assume(newRatePerSecond.unwrap() != flow.getRatePerSecond(currentStreamId).unwrap());
+        vm.assume(newRatePerSecond.unwrap() != previousRatePerSecond);
 
         // Adjust the rate per second.
         flow.adjustRatePerSecond(currentStreamId, newRatePerSecond);
 
-        flowStore.updatePeriods(currentStreamId, newRatePerSecond.unwrap(), "adjustRatePerSecond");
+        flowStore.updateDelay(currentStreamId, previousRatePerSecond, decimals);
+        flowStore.pushPeriod(currentStreamId, newRatePerSecond.unwrap(), "adjustRatePerSecond");
     }
 
     function deposit(
@@ -157,10 +160,14 @@ contract FlowHandler is BaseHandler {
         // Paused streams cannot be paused again.
         vm.assume(!flow.isPaused(currentStreamId));
 
+        flowStore.updateDelay(
+            currentStreamId, flow.getRatePerSecond(currentStreamId).unwrap(), flow.getTokenDecimals(currentStreamId)
+        );
+
         // Pause the stream.
         flow.pause(currentStreamId);
 
-        flowStore.updatePeriods(currentStreamId, 0, "pause");
+        flowStore.pushPeriod(currentStreamId, 0, "pause");
     }
 
     function refund(
@@ -226,7 +233,7 @@ contract FlowHandler is BaseHandler {
         // Restart the stream.
         flow.restart(currentStreamId, ratePerSecond);
 
-        flowStore.updatePeriods(currentStreamId, ratePerSecond.unwrap(), "restart");
+        flowStore.pushPeriod(currentStreamId, ratePerSecond.unwrap(), "restart");
     }
 
     function void(
@@ -249,7 +256,7 @@ contract FlowHandler is BaseHandler {
         // Void the stream.
         flow.void(currentStreamId);
 
-        flowStore.updatePeriods(currentStreamId, 0, "void");
+        flowStore.pushPeriod(currentStreamId, 0, "void");
     }
 
     function withdraw(
@@ -285,5 +292,12 @@ contract FlowHandler is BaseHandler {
 
         // Update the withdrawn amount.
         flowStore.updateStreamWithdrawnAmountsSum(currentStreamId, flow.getToken(currentStreamId), amount);
+
+        // If the stream isn't paused, update the delay:
+        uint128 ratePerSecond = flow.getRatePerSecond(currentStreamId).unwrap();
+        if (ratePerSecond > 0) {
+            flowStore.updateDelay(currentStreamId, ratePerSecond, flow.getTokenDecimals(currentStreamId));
+            flowStore.pushPeriod(currentStreamId, ratePerSecond, "withdraw");
+        }
     }
 }

test/invariant/stores/FlowStore.sol

@@ -31,11 +31,13 @@ contract FlowStore {
     /// @param ratePerSecond The rate per second for this period.
     /// @param start The start time of the period.
     /// @param end The end time of the period.
+    /// @param delay The delay for the period.
     struct Period {
         string typeOfPeriod;
         uint128 ratePerSecond;
         uint40 start;
         uint40 end;
+        uint40 delay;
     }
 
     /// @dev Each stream is mapped to an array of periods. This is used to calculate the total streamed amount.
@@ -57,23 +59,59 @@ contract FlowStore {
         // Store the stream id and the period during which provided ratePerSecond applies.
         streamIds.push(streamId);
         periods[streamId].push(
-            Period({ typeOfPeriod: "create", ratePerSecond: ratePerSecond, start: uint40(block.timestamp), end: 0 })
+            Period({
+                typeOfPeriod: "create",
+                ratePerSecond: ratePerSecond,
+                start: uint40(block.timestamp),
+                end: 0,
+                delay: 0
+            })
         );
 
         // Update the last stream id.
         lastStreamId = streamId;
     }
 
-    function updatePeriods(uint256 streamId, uint128 ratePerSecond, string memory typeOfPeriod) external {
+    function pushPeriod(uint256 streamId, uint128 ratePerSecond, string memory typeOfPeriod) external {
         // Update the end time of the previous period.
         periods[streamId][periods[streamId].length - 1].end = uint40(block.timestamp);
 
         // Push the new period with the provided rate per second.
         periods[streamId].push(
-            Period({ typeOfPeriod: typeOfPeriod, ratePerSecond: ratePerSecond, start: uint40(block.timestamp), end: 0 })
+            Period({
+                ratePerSecond: ratePerSecond,
+                start: uint40(block.timestamp),
+                end: 0,
+                delay: 0,
+                typeOfPeriod: typeOfPeriod
+            })
         );
     }
 
+    function updateDelay(uint256 streamId, uint128 ratePerSecond, uint8 decimals) external {
+        // Skip the delay update if the decimals are 18.
+        if (decimals == 18) {
+            return;
+        }
+
+        uint256 periodCount = periods[streamId].length - 1;
+        uint128 factor = uint128(10 ** (18 - decimals));
+        uint40 blockTimestamp = uint40(block.timestamp);
+        uint40 start = periods[streamId][periodCount].start;
+
+        uint128 rescaledStreamedAmount = ratePerSecond * (blockTimestamp - start) / factor * factor;
+
+        uint40 delay;
+        if (rescaledStreamedAmount > ratePerSecond) {
+            delay = blockTimestamp - start - uint40(rescaledStreamedAmount / ratePerSecond);
+            // Since we are reverse engineering the delay, we need to subtract 1 from the delay, which would normally be
+            // added in the constant interval calculation
+            delay = delay > 1 ? delay - 1 : 0;
+        }
+
+        periods[streamId][periodCount].delay += delay;
+    }
+
     function updatePreviousValues(
         uint256 streamId,
         uint40 snapshotTime,