From 8f5e0f14f8cdbcda9241358e275ef45b9888d17b Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Tue, 4 Feb 2025 16:10:40 +0100
Subject: [PATCH] vault: combine Account and Flow structs

---
 contracts/vault/Accounts.sol  | 52 +++++++++++++++++++++++++++++++++++
 contracts/vault/Flows.sol     | 23 ----------------
 contracts/vault/VaultBase.sol | 43 ++++++++---------------------
 3 files changed, 64 insertions(+), 54 deletions(-)
 create mode 100644 contracts/vault/Accounts.sol
 delete mode 100644 contracts/vault/Flows.sol

diff --git a/contracts/vault/Accounts.sol b/contracts/vault/Accounts.sol
new file mode 100644
index 0000000..e4135c2
--- /dev/null
+++ b/contracts/vault/Accounts.sol
@@ -0,0 +1,52 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.28;
+
+import "./TokensPerSecond.sol";
+
+struct Account {
+  uint128 available;
+  uint128 designated;
+  TokensPerSecond flow;
+  Timestamp flowUpdated;
+}
+
+library Accounts {
+  function isValidAt(
+    Account memory account,
+    Timestamp timestamp
+  ) internal pure returns (bool) {
+    if (account.flow < TokensPerSecond.wrap(0)) {
+      return uint128(-accumulateFlow(account, timestamp)) <= account.available;
+    } else {
+      return true;
+    }
+  }
+
+  function at(
+    Account memory account,
+    Timestamp timestamp
+  ) internal pure returns (Account memory) {
+    Account memory result = account;
+    if (result.flow != TokensPerSecond.wrap(0)) {
+      int128 accumulated = accumulateFlow(result, timestamp);
+      if (accumulated >= 0) {
+        result.designated += uint128(accumulated);
+      } else {
+        result.available -= uint128(-accumulated);
+      }
+    }
+    result.flowUpdated = timestamp;
+    return result;
+  }
+
+  function accumulateFlow(
+    Account memory account,
+    Timestamp timestamp
+  ) private pure returns (int128) {
+    int128 rate = TokensPerSecond.unwrap(account.flow);
+    Timestamp start = account.flowUpdated;
+    Timestamp end = timestamp;
+    uint64 duration = Timestamp.unwrap(end) - Timestamp.unwrap(start);
+    return rate * int128(uint128(duration));
+  }
+}
diff --git a/contracts/vault/Flows.sol b/contracts/vault/Flows.sol
deleted file mode 100644
index 46bcfb3..0000000
--- a/contracts/vault/Flows.sol
+++ /dev/null
@@ -1,23 +0,0 @@
-// SPDX-License-Identifier: MIT
-pragma solidity 0.8.28;
-
-import "./Timestamps.sol";
-import "./TokensPerSecond.sol";
-
-struct Flow {
-  TokensPerSecond rate;
-  Timestamp updated;
-}
-
-library Flows {
-  function totalAt(
-    Flow memory flow,
-    Timestamp timestamp
-  ) internal pure returns (int128) {
-    int128 rate = TokensPerSecond.unwrap(flow.rate);
-    Timestamp start = flow.updated;
-    Timestamp end = timestamp;
-    uint64 duration = Timestamp.unwrap(end) - Timestamp.unwrap(start);
-    return rate * int128(uint128(duration));
-  }
-}
diff --git a/contracts/vault/VaultBase.sol b/contracts/vault/VaultBase.sol
index 44862a7..2e2773f 100644
--- a/contracts/vault/VaultBase.sol
+++ b/contracts/vault/VaultBase.sol
@@ -3,14 +3,14 @@ pragma solidity 0.8.28;
 
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+import "./Accounts.sol";
 import "./Timestamps.sol";
 import "./TokensPerSecond.sol";
-import "./Flows.sol";
 import "./Locks.sol";
 
 using SafeERC20 for IERC20;
 using Timestamps for Timestamp;
-using Flows for Flow;
+using Accounts for Account;
 using Locks for Lock;
 
 abstract contract VaultBase {
@@ -20,12 +20,6 @@ abstract contract VaultBase {
   type Fund is bytes32;
   type Recipient is address;
 
-  struct Account {
-    uint128 available;
-    uint128 designated;
-    Flow flow;
-  }
-
   mapping(Controller => mapping(Fund => Lock)) private _locks;
   mapping(Controller => mapping(Fund => mapping(Recipient => Account)))
     private _accounts;
@@ -47,19 +41,12 @@ abstract contract VaultBase {
     Recipient recipient
   ) internal view returns (Account memory) {
     Account memory account = _accounts[controller][fund][recipient];
-    Timestamp timestamp = Timestamps.currentTime();
-    if (account.flow.rate != TokensPerSecond.wrap(0)) {
-      Lock memory lock = _locks[controller][fund];
-      Timestamp end = Timestamps.earliest(timestamp, lock.expiry);
-      int128 accumulated = account.flow.totalAt(end);
-      if (accumulated >= 0) {
-        account.designated += uint128(accumulated);
-      } else {
-        account.available -= uint128(-accumulated);
-      }
-    }
-    account.flow.updated = timestamp;
-    return account;
+    Lock memory lock = _locks[controller][fund];
+    Timestamp timestamp = Timestamps.earliest(
+      Timestamps.currentTime(),
+      lock.expiry
+    );
+    return account.at(timestamp);
   }
 
   function _lock(
@@ -165,8 +152,8 @@ abstract contract VaultBase {
     Account memory senderAccount = _getAccount(controller, fund, from);
     Account memory receiverAccount = _getAccount(controller, fund, to);
 
-    senderAccount.flow.rate = senderAccount.flow.rate - rate;
-    receiverAccount.flow.rate = receiverAccount.flow.rate + rate;
+    senderAccount.flow = senderAccount.flow - rate;
+    receiverAccount.flow = receiverAccount.flow + rate;
 
     _checkAccountInvariant(senderAccount, lock);
 
@@ -183,10 +170,7 @@ abstract contract VaultBase {
     require(lock.isLocked(), LockRequired());
 
     Account memory account = _getAccount(controller, fund, recipient);
-    require(
-      account.flow.rate == TokensPerSecond.wrap(0),
-      CannotBurnFlowingTokens()
-    );
+    require(account.flow == TokensPerSecond.wrap(0), CannotBurnFlowingTokens());
 
     uint128 amount = account.available + account.designated;
 
@@ -229,10 +213,7 @@ abstract contract VaultBase {
     Account memory account,
     Lock memory lock
   ) private pure {
-    if (account.flow.rate < TokensPerSecond.wrap(0)) {
-      uint128 outgoing = uint128(-account.flow.totalAt(lock.maximum));
-      require(outgoing <= account.available, InsufficientBalance());
-    }
+    require(account.isValidAt(lock.maximum), InsufficientBalance());
   }
 
   error InsufficientBalance();