more UBI tests

onchain/src/Harb.sol

@@ -52,6 +52,8 @@ contract Harb is ERC20, ERC20Permit {
 
     /// @notice Thrown if the some address is unexpectedly the zero address.
     error ZeroAddressInConstructor();
+    error ZeroAddressInSetter();
+    error AddressAlreadySet();
 
     event UbiClaimed(address indexed owner, uint256 ubiAmount);
 
@@ -82,14 +84,14 @@ contract Harb is ERC20, ERC20Permit {
 
 
     function setLiquidityManager(address liquidityManager_) external {
-        // TODO: add trapdoor
-        if (address(0) == liquidityManager_) revert ZeroAddressInConstructor();
+        if (address(0) == liquidityManager_) revert ZeroAddressInSetter();
+        if (liquidityManager != address(0)) revert AddressAlreadySet();
         liquidityManager = liquidityManager_;
     }
 
     function setStakingPool(address stakingPool_) external {
-        // TODO: add trapdoor
-        if (address(0) == stakingPool_) revert ZeroAddressInConstructor();
+        if (address(0) == stakingPool_) revert ZeroAddressInSetter();
+        if (stakingPool != address(0)) revert AddressAlreadySet();
         stakingPool = stakingPool_;
     }
 
@@ -135,6 +137,7 @@ contract Harb is ERC20, ERC20Permit {
      * @param amount Tokens to mint
      */
     function _mint(address receiver, uint256 amount) internal override {
+        // TODO: limit supply to 2^96?
         // make sure staking pool grows proportional to economy
         uint256 stakingPoolBalance = balanceOf(stakingPool);
         if (stakingPoolBalance > 0) {
@@ -216,9 +219,8 @@ contract Harb is ERC20, ERC20Permit {
         amountDue = taxCollectedSinceLastClaim.mulDiv(accountTwab, (totalSupplyTwab - stakeTwab - poolTwab - taxTwab), Math.Rounding.Down);
     }
 
-    function claimUbi(address _account) external {
+    function claimUbi(address _account) external returns (uint256 ubiAmountDue) {
         UbiTitle storage lastUbiTitle = ubiTitles[_account];
-        uint256 ubiAmountDue;
         uint256 lastPeriodEndAt;
         (ubiAmountDue, lastPeriodEndAt) = ubiDue(_account, lastUbiTitle.time, lastUbiTitle.sumTaxCollected);
 

onchain/src/Stake.sol

@@ -205,6 +205,7 @@ contract Stake is IStake {
         if (pos.owner != msg.sender) {
             revert NoPermission(msg.sender, pos.owner);
         }
+        //TODO: implement not found
         // to prevent snatch-and-change grieving attack, pay TAX_FLOOR_DURATION
         require(taxRate > pos.taxRate, "tax too low to snatch");
         _payTax(positionID, pos, TAX_FLOOR_DURATION);
@@ -216,6 +217,7 @@ contract Stake is IStake {
         if (pos.owner != msg.sender) {
             revert NoPermission(msg.sender, pos.owner);
         }
+        //TODO: implement not found
         // to prevent snatch-and-exit grieving attack, pay TAX_FLOOR_DURATION
         _payTax(positionId, pos, TAX_FLOOR_DURATION);
         _exitPosition(positionId, pos);

onchain/test/Harb.t.sol

@@ -22,6 +22,7 @@ contract HarbTest is Test {
     Harb harb;
     IUniswapV3Factory factory;
     address stakingPool;
+    address liqPool;
     BaseLineLP liquidityManager;
     TwabController tc;
 
@@ -39,12 +40,12 @@ contract HarbTest is Test {
         factory = IUniswapV3Factory(factoryAddress);
         weth = IWETH9(address(new WETH()));
 
-        tc = new TwabController(60 * 60 * 24, uint32(block.timestamp));
+        tc = new TwabController(60 * 60, uint32(block.timestamp));
         harb = new Harb("HARB", "HARB", factoryAddress, address(weth), tc);
 
 
         factory = IUniswapV3Factory(factoryAddress);
-        factory.createPool(address(weth), address(harb), FEE);
+        liqPool = factory.createPool(address(weth), address(harb), FEE);
         stakingPool = makeAddr("stakingPool"); // This represents the staking pool
         harb.setStakingPool(stakingPool);
         liquidityManager = new BaseLineLP(factoryAddress, address(weth), address(harb));
@@ -358,4 +359,73 @@ contract HarbTest is Test {
         assertEq(harb.sumTaxCollected(), 0, "No tax should be collected, and sumTaxCollected should remain zero after the claim attempt.");
     }
 
+    function testEdgeCaseWithMaximumTaxCollection() public {
+        uint256 initialSupply = 1e24; // Large number of HARB tokens to simulate realistic large-scale deployment
+        uint256 maxTaxAmount = type(uint96).max - initialSupply; // Setting max tax just below overflow threshold when added to total supply
+
+        address account1 = makeAddr("alice");
+
+        // Setup initial supply and allocate to user
+        vm.startPrank(address(liquidityManager));
+        harb.mint(initialSupply + maxTaxAmount);
+        harb.transfer(account1, initialSupply);
+        harb.transfer(TAX_POOL, maxTaxAmount); // Simulate tax collection at the theoretical maximum
+        vm.stopPrank();
+
+        // Assert that maximum tax was collected
+        assertEq(harb.sumTaxCollected(), maxTaxAmount, "Max tax collected should match the max tax amount transferred.");
+
+        // Simulate time passage and UBI claim
+        vm.warp(block.timestamp + 30 days);
+
+        // Account 1 claims UBI
+        vm.prank(account1);
+        harb.claimUbi(account1);
+
+        // Check if the account's balance increased correctly
+        uint256 expectedBalance = initialSupply + maxTaxAmount; // This assumes the entire tax pool goes to one account, simplify as needed
+        assertEq(harb.balanceOf(account1), expectedBalance, "Account 1's balance after claiming UBI with max tax collection is incorrect.");
+
+        // Verify that no taxes are left unclaimed
+        assertEq(harb.balanceOf(TAX_POOL), 0, "All taxes should be claimed after the UBI claim.");
+    }
+
+
+    // TODO: why is this test passing even though it exceeds MAX_CARDINALITY?
+    function testTwabBeyondBuffer() public {
+        uint256 initialSupply = 1000 * 1e18; // 1000 HARB tokens
+        uint256 taxAmount = 300 * 1e18; // 300 HARB tokens to be collected as tax
+
+        address account1 = makeAddr("alice");
+
+        // Setup initial supply and allocate to user
+        vm.startPrank(address(liquidityManager));
+        harb.mint(initialSupply + taxAmount);
+        harb.transfer(account1, initialSupply / 800);
+        harb.transfer(TAX_POOL, taxAmount); // Simulate tax collection at the theoretical maximum
+        harb.transfer(liqPool, harb.balanceOf(address(liquidityManager)));
+        vm.stopPrank();
+        vm.warp(block.timestamp + 1 hours);
+
+
+        // Simulate updates over a longer period, e.g., enough to potentially wrap the buffer.
+        uint numHours = 399;  // More than 365 to potentially test buffer wrapping (MAX_CARDINALITY)
+        for (uint i = 0; i < numHours; i++) {
+            vm.prank(liqPool);
+            harb.transfer(account1, initialSupply / 800);
+            vm.warp(block.timestamp + 1 hours);  // Fast-forward time by one hour.
+        }
+
+        // Account 1 claims UBI
+        vm.prank(account1);
+        uint256 ubiCollected = harb.claimUbi(account1);
+
+        // Check if the account's balance increased correctly
+        uint256 expectedBalance = (initialSupply / 2) + ubiCollected; // This assumes the entire tax pool goes to one account, simplify as needed
+        assertApproxEqRel(harb.balanceOf(account1), expectedBalance, 1 * 1e18, "Account 1's balance after claiming UBI with max tax collection is incorrect.");
+
+        // Verify that no taxes are left unclaimed
+        assertEq(harb.balanceOf(TAX_POOL), 0, "All taxes should be claimed after the UBI claim.");
+    }
+
 }