vault: change data structure to be recipient oriented

contracts/Vault.sol

@@ -8,27 +8,37 @@ using SafeERC20 for IERC20;
 
 contract Vault {
   IERC20 private immutable _token;
-  mapping(address => mapping(bytes32 => uint256)) private _amounts;
+
+  type Controller is address;
+  type Context is bytes32;
+  type Recipient is address;
+
+  mapping(Controller => mapping(Context => mapping(Recipient => uint256)))
+    private _available;
 
   constructor(IERC20 token) {
     _token = token;
   }
 
-  function amount(bytes32 id) public view returns (uint256) {
-    return _amounts[msg.sender][id];
+  function balance(
+    Context context,
+    Recipient recipient
+  ) public view returns (uint256) {
+    Controller controller = Controller.wrap(msg.sender);
+    return _available[controller][context][recipient];
   }
 
-  function deposit(bytes32 id, address from, uint256 value) public {
-    require(_amounts[msg.sender][id] == 0, DepositAlreadyExists(id));
-    _amounts[msg.sender][id] = value;
-    _token.safeTransferFrom(from, address(this), value);
+  function deposit(Context context, address from, uint256 amount) public {
+    Controller controller = Controller.wrap(msg.sender);
+    Recipient recipient = Recipient.wrap(from);
+    _available[controller][context][recipient] += amount;
+    _token.safeTransferFrom(from, address(this), amount);
   }
 
-  function withdraw(bytes32 id, address recipient) public {
-    uint256 value = _amounts[msg.sender][id];
-    delete _amounts[msg.sender][id];
-    _token.safeTransfer(recipient, value);
+  function withdraw(Context context, Recipient recipient) public {
+    Controller controller = Controller.wrap(msg.sender);
+    uint256 amount = _available[controller][context][recipient];
+    delete _available[controller][context][recipient];
+    _token.safeTransfer(Recipient.unwrap(recipient), amount);
   }
-
-  error DepositAlreadyExists(bytes32 id);
 }

test/Vault.tests.js

@@ -5,72 +5,95 @@ const { randomBytes } = ethers.utils
 describe("Vault", function () {
   let token
   let vault
-  let payer
-  let recipient
+  let account
 
   beforeEach(async function () {
     const TestToken = await ethers.getContractFactory("TestToken")
     token = await TestToken.deploy()
     const Vault = await ethers.getContractFactory("Vault")
     vault = await Vault.deploy(token.address)
-    ;[_, payer, recipient] = await ethers.getSigners()
-    await token.mint(payer.address, 1_000_000)
+    ;[, account] = await ethers.getSigners()
+    await token.mint(account.address, 1_000_000)
   })
 
   describe("depositing", function () {
-    const id = randomBytes(32)
+    const context = randomBytes(32)
     const amount = 42
 
     it("accepts deposits of tokens", async function () {
-      await token.connect(payer).approve(vault.address, amount)
-      await vault.deposit(id, payer.address, amount)
-      expect(await vault.amount(id)).to.equal(amount)
+      await token.connect(account).approve(vault.address, amount)
+      await vault.deposit(context, account.address, amount)
+      expect(await vault.balance(context, account.address)).to.equal(amount)
     })
 
     it("keeps custody of tokens that are deposited", async function () {
-      await token.connect(payer).approve(vault.address, amount)
-      await vault.deposit(id, payer.address, amount)
+      await token.connect(account).approve(vault.address, amount)
+      await vault.deposit(context, account.address, amount)
       expect(await token.balanceOf(vault.address)).to.equal(amount)
     })
 
     it("deposit fails when tokens cannot be transferred", async function () {
-      await token.connect(payer).approve(vault.address, amount - 1)
-      const depositing = vault.deposit(id, payer.address, amount)
+      await token.connect(account).approve(vault.address, amount - 1)
+      const depositing = vault.deposit(context, account.address, amount)
       await expect(depositing).to.be.revertedWith("insufficient allowance")
     })
 
-    it("requires deposit ids to be unique", async function () {
-      await token.connect(payer).approve(vault.address, 2 * amount)
-      await vault.deposit(id, payer.address, amount)
-      const depositing = vault.deposit(id, payer.address, amount)
-      await expect(depositing).to.be.revertedWith("DepositAlreadyExists")
+    it("multiple deposits add to the balance", async function () {
+      await token.connect(account).approve(vault.address, amount)
+      await vault.deposit(context, account.address, amount / 2)
+      await vault.deposit(context, account.address, amount / 2)
+      expect(await vault.balance(context, account.address)).to.equal(amount)
     })
 
-    it("separates deposits from different owners", async function () {
-      let [owner1, owner2] = await ethers.getSigners()
-      await token.connect(payer).approve(vault.address, 3)
-      await vault.connect(owner1).deposit(id, payer.address, 1)
-      await vault.connect(owner2).deposit(id, payer.address, 2)
-      expect(await vault.connect(owner1).amount(id)).to.equal(1)
-      expect(await vault.connect(owner2).amount(id)).to.equal(2)
+    it("separates deposits from different contexts", async function () {
+      const context1 = randomBytes(32)
+      const context2 = randomBytes(32)
+      await token.connect(account).approve(vault.address, 3)
+      await vault.deposit(context1, account.address, 1)
+      await vault.deposit(context2, account.address, 2)
+      expect(await vault.balance(context1, account.address)).to.equal(1)
+      expect(await vault.balance(context2, account.address)).to.equal(2)
+    })
+
+    it("separates deposits from different controllers", async function () {
+      const [, , controller1, controller2] = await ethers.getSigners()
+      const vault1 = vault.connect(controller1)
+      const vault2 = vault.connect(controller2)
+      await token.connect(account).approve(vault.address, 3)
+      await vault1.deposit(context, account.address, 1)
+      await vault2.deposit(context, account.address, 2)
+      expect(await vault1.balance(context, account.address)).to.equal(1)
+      expect(await vault2.balance(context, account.address)).to.equal(2)
     })
   })
 
   describe("withdrawing", function () {
-    const id = randomBytes(32)
+    const context = randomBytes(32)
+    const amount = 42
 
-    it("can withdraw a deposit", async function () {
-      const amount = 42
-      await token.connect(payer).approve(vault.address, amount)
-      await vault.deposit(id, payer.address, amount)
-      await vault.withdraw(id, recipient.address)
-      expect(await vault.amount(id)).to.equal(0)
-      expect(await token.balanceOf(recipient.address)).to.equal(amount)
+    beforeEach(async function () {
+      await token.connect(account).approve(vault.address, amount)
+      await vault.deposit(context, account.address, amount)
     })
 
-    it("ignores withdrawal of an empty deposit", async function () {
-      await vault.withdraw(id, recipient.address)
-      expect(await token.balanceOf(recipient.address)).to.equal(0)
+    it("can withdraw a deposit", async function () {
+      const before = await token.balanceOf(account.address)
+      await vault.withdraw(context, account.address)
+      const after = await token.balanceOf(account.address)
+      expect(after - before).to.equal(amount)
+    })
+
+    it("empties the balance when withdrawing", async function () {
+      await vault.withdraw(context, account.address)
+      expect(await vault.balance(context, account.address)).to.equal(0)
+    })
+
+    it("does not withdraw more than once", async function () {
+      await vault.withdraw(context, account.address)
+      const before = await token.balanceOf(account.address)
+      await vault.withdraw(context, account.address)
+      const after = await token.balanceOf(account.address)
+      expect(after).to.equal(before)
     })
   })
 })