From a7d7578cb8dea32daef1f96fde9b8cae4b4ade54 Mon Sep 17 00:00:00 2001
From: Daniel Von Fange <daniel@leancoder.com>
Date: Wed, 30 Oct 2024 15:34:38 -0400
Subject: [PATCH] Create delegations between accounts

---
 src/token/OUSD.sol | 26 +++++++++++++++++++++--
 test/OUSD.t.sol    | 52 ++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/src/token/OUSD.sol b/src/token/OUSD.sol
index c4e3e1e..a5717e8 100644
--- a/src/token/OUSD.sol
+++ b/src/token/OUSD.sol
@@ -46,6 +46,9 @@ contract OUSD is Governable {
     mapping(address => uint256) public nonRebasingCreditsPerToken;
     mapping(address => RebaseOptions) public rebaseState;
     mapping(address => uint256) public isUpgraded;
+    mapping(address => address) public yieldDelegate;
+    mapping(address => uint256) public yieldDelegateeCount;
+    mapping(address => uint256) public aintMoney;
 
     uint256 private constant RESOLUTION_INCREASE = 1e9;
 
@@ -127,7 +130,6 @@ contract OUSD is Governable {
         view
         returns (uint256)
     {
-        if (_creditBalances[_account] == 0) return 0;
         return
             _creditBalances[_account] * 1e18 / _creditsPerToken(_account);
     }
@@ -582,7 +584,7 @@ contract OUSD is Governable {
             : _newTotalSupply;
 
         _rebasingCreditsPerToken = _rebasingCredits 
-            * 1e18 / _totalSupply - nonRebasingSupply;
+            * 1e18 / (_totalSupply - nonRebasingSupply);
 
         require(_rebasingCreditsPerToken > 0, "Invalid change in supply");
 
@@ -595,4 +597,24 @@ contract OUSD is Governable {
             _rebasingCreditsPerToken
         );
     }
+
+    function delegateYield(address from, address to) external onlyGovernor() {
+        require(from != to, "Cannot delegate to self");
+        require(yieldDelegateeCount[from] == 0, "Cannot delegate from delegatee");
+        require(yieldDelegate[to] == address(0), "Cannot delegate to delegator");
+        require(yieldDelegate[from] == address(0), "Already delegated");
+        require(!_isNonRebasingAccount(to), "Must delegate to a rebasing account");
+        require(_isNonRebasingAccount(from), "Must delegate from a non-rebasing account");
+        
+        yieldDelegate[from] = to;
+        rebaseState[from] = RebaseOptions.OptIn;
+        rebaseState[to] = RebaseOptions.OptIn;
+        yieldDelegateeCount[to] += 1;
+    }
+
+    function undelegateYield(address from) external onlyGovernor() {
+        require(yieldDelegate[from] != address(0), "");
+        yieldDelegateeCount[yieldDelegate[from]] -= 1;
+        yieldDelegate[from] = address(0);
+    }
 }
\ No newline at end of file
diff --git a/test/OUSD.t.sol b/test/OUSD.t.sol
index 55ab048..b4f9e80 100644
--- a/test/OUSD.t.sol
+++ b/test/OUSD.t.sol
@@ -20,19 +20,59 @@ contract CounterTest is Test {
         ousd.initialize("", "", address(this), 314159265358979323846264333);
         
         ousd.mint(matt, 1000 ether);
+        assertEq(ousd.totalSupply(), 1000 ether);
         ousd.mint(labs, 1000 ether);
         ousd.mint(pool, 1000 ether);
         ousd.mint(collector, 1000 ether);
         ousd.mint(attacker, 1000 ether);
+        assertEq(ousd.totalSupply(), 5000 ether);
+
+        vm.prank(pool);
+        ousd.rebaseOptOut();
+
+        vm.prank(collector);
+        ousd.rebaseOptOut();
+        vm.prank(collector);
+        ousd.rebaseOptIn();
+        
+        ousd.delegateYield(pool, collector);
+        assertEq(ousd.totalSupply(), 5000 ether);
+        assertEq(ousd.nonRebasingSupply(), 1000 ether);
     }
 
     function test_ChangeSupply() public {
-        ousd.changeSupply(1000 ether);
-        assertEq(ousd.totalSupply(), 1000 ether);
+        assertEq(ousd.totalSupply(), 5000 ether);
+        assertEq(ousd.nonRebasingSupply(), 1000 ether);
+        ousd.changeSupply(6000 ether);
+        assertEq(ousd.totalSupply(), 6000 ether);
+        assertEq(ousd.nonRebasingSupply(), 1000 ether);
+    }
+
+
+    function test_CanDelegateYield() public {
+        vm.prank(matt);
+        ousd.rebaseOptOut();
+        ousd.delegateYield(matt, attacker);
+    }
+
+    function test_NoDelegateYieldToSelf() public {
+        vm.expectRevert("Cannot delegate to self");
+        ousd.delegateYield(matt, matt);
     }
 
-    // function testFuzz_SetNumber(uint256 x) public {
-    //     counter.setNumber(x);
-    //     assertEq(counter.number(), x);
-    // }
+    function test_NoDelegateYieldToDelegator() public {
+        vm.expectRevert("Cannot delegate to delegator");
+        ousd.delegateYield(matt, pool);
+    }
+
+    function test_NoDelegateYieldFromReceiver() public {
+        assertEq(ousd.yieldDelegateeCount(collector),1);
+        vm.expectRevert("Cannot delegate from delegatee");
+        ousd.delegateYield(collector, matt);
+    }
+    
+    function test_CanUndelegeteYield() public {
+        ousd.undelegateYield(pool);
+        assertEq(ousd.yieldDelegate(pool), address(0));
+    }
 }