From d9722d55fd5e87eb7dd457d1b87b0e0c1f52ffb5 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Thu, 23 Jan 2025 10:19:41 +0100
Subject: [PATCH] vault: use custom operators to improve readability

---
 contracts/Timestamps.sol      | 19 +++++++++++--------
 contracts/TokensPerSecond.sol |  5 +++++
 contracts/VaultBase.sol       | 10 +++++-----
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/contracts/Timestamps.sol b/contracts/Timestamps.sol
index 4910842..96756b7 100644
--- a/contracts/Timestamps.sol
+++ b/contracts/Timestamps.sol
@@ -3,16 +3,19 @@ pragma solidity 0.8.28;
 
 type Timestamp is uint64;
 
+using {_lessThan as <} for Timestamp global;
+using {_atMost as <=} for Timestamp global;
+
+function _lessThan(Timestamp a, Timestamp b) pure returns (bool) {
+  return Timestamp.unwrap(a) < Timestamp.unwrap(b);
+}
+
+function _atMost(Timestamp a, Timestamp b) pure returns (bool) {
+  return Timestamp.unwrap(a) <= Timestamp.unwrap(b);
+}
+
 library Timestamps {
   function currentTime() internal view returns (Timestamp) {
     return Timestamp.wrap(uint64(block.timestamp));
   }
-
-  function isAfter(Timestamp a, Timestamp b) internal pure returns (bool) {
-    return Timestamp.unwrap(a) > Timestamp.unwrap(b);
-  }
-
-  function isFuture(Timestamp timestamp) internal view returns (bool) {
-    return Timestamp.unwrap(timestamp) > block.timestamp;
-  }
 }
diff --git a/contracts/TokensPerSecond.sol b/contracts/TokensPerSecond.sol
index 9459020..677da3a 100644
--- a/contracts/TokensPerSecond.sol
+++ b/contracts/TokensPerSecond.sol
@@ -6,7 +6,12 @@ import "./Timestamps.sol";
 type TokensPerSecond is int256;
 
 using {_negate as -} for TokensPerSecond global;
+using {_equals as ==} for TokensPerSecond global;
 
 function _negate(TokensPerSecond rate) pure returns (TokensPerSecond) {
   return TokensPerSecond.wrap(-TokensPerSecond.unwrap(rate));
 }
+
+function _equals(TokensPerSecond a, TokensPerSecond b) pure returns (bool) {
+  return TokensPerSecond.unwrap(a) == TokensPerSecond.unwrap(b);
+}
diff --git a/contracts/VaultBase.sol b/contracts/VaultBase.sol
index 9d7f591..3789b2e 100644
--- a/contracts/VaultBase.sol
+++ b/contracts/VaultBase.sol
@@ -99,7 +99,7 @@ abstract contract VaultBase {
     Context context,
     Recipient recipient
   ) internal {
-    require(!_getLock(controller, context).expiry.isFuture(), Locked());
+    require(_getLock(controller, context).expiry <= Timestamps.currentTime(), Locked());
     delete _locks[controller][context];
     Balance memory balance = _getBalance(controller, context, recipient);
     uint256 amount = balance.available + balance.designated;
@@ -155,7 +155,7 @@ abstract contract VaultBase {
       Timestamp.unwrap(_getLock(controller, context).maximum) == 0,
       AlreadyLocked()
     );
-    require(!expiry.isAfter(maximum), ExpiryPastMaximum());
+    require(expiry <= maximum, ExpiryPastMaximum());
     _locks[controller][context] = Lock({expiry: expiry, maximum: maximum});
   }
 
@@ -165,9 +165,9 @@ abstract contract VaultBase {
     Timestamp expiry
   ) internal {
     Lock memory previous = _getLock(controller, context);
-    require(previous.expiry.isFuture(), LockExpired());
-    require(!previous.expiry.isAfter(expiry), InvalidExpiry());
-    require(!expiry.isAfter(previous.maximum), ExpiryPastMaximum());
+    require(Timestamps.currentTime() < previous.expiry, LockExpired());
+    require(previous.expiry <= expiry, InvalidExpiry());
+    require(expiry <= previous.maximum, ExpiryPastMaximum());
     _locks[controller][context].expiry = expiry;
   }