vault: flow tokens from one recipient to the other

contracts/Timestamps.sol

@@ -4,6 +4,10 @@ pragma solidity 0.8.28;
 type Timestamp is uint64;
 
 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);
   }

contracts/TokenFlows.sol

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.28;
+
+import "./Timestamps.sol";
+
+type TokensPerSecond is int256;
+
+using {_negate as -} for TokensPerSecond global;
+
+function _negate(TokensPerSecond rate) pure returns (TokensPerSecond) {
+  return TokensPerSecond.wrap(-TokensPerSecond.unwrap(rate));
+}
+
+library TokenFlows {
+  function accumulated(
+    TokensPerSecond rate,
+    Timestamp start,
+    Timestamp end
+  ) internal pure returns (int256) {
+    if (TokensPerSecond.unwrap(rate) == 0) {
+      return 0;
+    }
+    uint64 duration = Timestamp.unwrap(end) - Timestamp.unwrap(start);
+    return TokensPerSecond.unwrap(rate) * int256(uint256(duration));
+  }
+}

contracts/Vault.sol

@@ -76,4 +76,14 @@ contract Vault is VaultBase {
     Controller controller = Controller.wrap(msg.sender);
     _extendLock(controller, context, expiry);
   }
+
+  function flow(
+    Context context,
+    Recipient from,
+    Recipient to,
+    TokensPerSecond rate
+  ) public {
+    Controller controller = Controller.wrap(msg.sender);
+    _flow(controller, context, from, to, rate);
+  }
 }

contracts/VaultBase.sol

@@ -4,9 +4,11 @@ pragma solidity 0.8.28;
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "./Timestamps.sol";
+import "./TokenFlows.sol";
 
 using SafeERC20 for IERC20;
 using Timestamps for Timestamp;
+using TokenFlows for TokensPerSecond;
 
 abstract contract VaultBase {
   IERC20 internal immutable _token;
@@ -20,11 +22,18 @@ abstract contract VaultBase {
     Timestamp maximum;
   }
 
+  struct Flow {
+    Timestamp start;
+    TokensPerSecond rate;
+  }
+
   mapping(Controller => mapping(Context => Lock)) private _locks;
   mapping(Controller => mapping(Context => mapping(Recipient => uint256)))
     private _available;
   mapping(Controller => mapping(Context => mapping(Recipient => uint256)))
     private _designated;
+  mapping(Controller => mapping(Context => mapping(Recipient => Flow)))
+    private _flows;
 
   constructor(IERC20 token) {
     _token = token;
@@ -35,9 +44,15 @@ abstract contract VaultBase {
     Context context,
     Recipient recipient
   ) internal view returns (uint256) {
-    return
-      _available[controller][context][recipient] +
-      _designated[controller][context][recipient];
+    Flow memory flow = _flows[controller][context][recipient];
+    int256 flowed = flow.rate.accumulated(flow.start, Timestamps.currentTime());
+    uint256 available = _available[controller][context][recipient];
+    uint256 designated = _designated[controller][context][recipient];
+    if (flowed >= 0) {
+      return available + designated + uint256(flowed);
+    } else {
+      return available + designated - uint256(-flowed);
+    }
   }
 
   function _getDesignated(
@@ -152,6 +167,18 @@ abstract contract VaultBase {
     _locks[controller][context].expiry = expiry;
   }
 
+  function _flow(
+    Controller controller,
+    Context context,
+    Recipient from,
+    Recipient to,
+    TokensPerSecond rate
+  ) internal {
+    Timestamp start = Timestamps.currentTime();
+    _flows[controller][context][to] = Flow({start: start, rate: rate});
+    _flows[controller][context][from] = Flow({start: start, rate: -rate});
+  }
+
   error InsufficientBalance();
   error Locked();
   error AlreadyLocked();

test/Vault.tests.js

@@ -1,7 +1,7 @@
 const { expect } = require("chai")
 const { ethers } = require("hardhat")
 const { randomBytes } = ethers.utils
-const { currentTime, advanceTimeTo } = require("./evm")
+const { currentTime, advanceTimeTo, mine } = require("./evm")
 
 describe("Vault", function () {
   let token
@@ -374,10 +374,31 @@ describe("Vault", function () {
 
     it("deletes lock when funds are withdrawn", async function () {
       await vault.lockup(context, expiry, expiry)
-      await advanceTimeToForNextBlock(expiry)
+      await advanceTimeTo(expiry)
       await vault.withdraw(context, account.address)
       expect((await vault.lock(context))[0]).to.equal(0)
       expect((await vault.lock(context))[1]).to.equal(0)
     })
   })
+
+  describe("flow", function () {
+    const context = randomBytes(32)
+    const amount = 42
+
+    beforeEach(async function () {
+      await token.connect(account).approve(vault.address, amount)
+      await vault.deposit(context, account.address, amount)
+    })
+
+    it("moves tokens over time", async function () {
+      await vault.flow(context, account.address, account2.address, 2)
+      const start = await currentTime()
+      await advanceTimeTo(start + 2)
+      expect(await vault.balance(context, account.address)).to.equal(amount - 4)
+      expect(await vault.balance(context, account2.address)).to.equal(4)
+      await advanceTimeTo(start + 4)
+      expect(await vault.balance(context, account.address)).to.equal(amount - 8)
+      expect(await vault.balance(context, account2.address)).to.equal(8)
+    })
+  })
 })