events for stats and fix for getUbiDue issue #3

onchain/script/Deploy.sol

@@ -20,7 +20,8 @@ contract SepoliaScript is Script {
         uint256 privateKey = vm.deriveKey(seedPhrase, 0);
         vm.startBroadcast(privateKey);
 
-        TwabController tc = TwabController(TWABC);
+        //TwabController tc = TwabController(TWABC);
+        TwabController tc = new TwabController(60 * 60, uint32(block.timestamp));
         Harb harb = new Harb("Harberger Tax", "HARB", V3_FACTORY, WETH, tc);
         Stake stake = new Stake(address(harb));
         harb.setStakingPool(address(stake));

onchain/src/Harb.sol

@@ -28,6 +28,8 @@ contract Harb is ERC20, ERC20Permit {
 
     /// @notice Address of the TwabController used to keep track of balances.
     TwabController public immutable twabController;
+    uint256 immutable PERIOD_OFFSET;
+    uint256 immutable PERIOD_LENGTH;
     IUniswapV3Pool immutable pool;
 
     /// @notice Address of the LiquidityManager contract that mints and burns supply
@@ -70,10 +72,13 @@ contract Harb is ERC20, ERC20Permit {
     {
         if (address(0) == address(twabController_)) revert ZeroAddressInConstructor();
         twabController = twabController_;
+        PERIOD_OFFSET = twabController.PERIOD_OFFSET();
+        PERIOD_LENGTH = twabController.PERIOD_LENGTH();
         PoolKey memory poolKey = PoolAddress.getPoolKey(_WETH9, address(this), FEE);
         pool = IUniswapV3Pool(PoolAddress.computeAddress(_factory, poolKey));
     }
 
+
     function setLiquidityManager(address liquidityManager_) external {
         // TODO: add trapdoor
         if (address(0) == liquidityManager_) revert ZeroAddressInConstructor();
@@ -126,11 +131,10 @@ contract Harb is ERC20, ERC20Permit {
     function _mint(address receiver, uint256 amount) internal override {
         // make sure staking pool grows proportional to economy
         uint256 stakingPoolBalance = balanceOf(stakingPool);
-        uint256 activeSupply = totalSupply() - stakingPoolBalance;
-        uint256 dormantStake = IStake(stakingPool).dormantSupply();
         if (stakingPoolBalance > 0) {
-            uint256 newStake = stakingPoolBalance * amount / (activeSupply + dormantStake);
-            _mint(stakingPool, newStake);
+            uint256 newStake = stakingPoolBalance *  amount / (totalSupply() - stakingPoolBalance);
+            twabController.mint(stakingPool, SafeCast.toUint96(newStake));
+            emit Transfer(address(0), stakingPool, newStake);
         }
         twabController.mint(receiver, SafeCast.toUint96(amount));
         emit Transfer(address(0), receiver, amount);
@@ -145,7 +149,13 @@ contract Harb is ERC20, ERC20Permit {
      * @param _amount The amount of tokens to burn
      */
     function _burn(address _owner, uint256 _amount) internal override {
-        // TODO
+        // shrink staking pool proportional to economy
+        uint256 stakingPoolBalance = balanceOf(stakingPool);
+        if (stakingPoolBalance > 0) {
+            uint256 excessStake = stakingPoolBalance * _amount / (totalSupply() - stakingPoolBalance);
+            twabController.burn(stakingPool, SafeCast.toUint96(excessStake));
+            emit Transfer(stakingPool, address(0), excessStake);
+        }
         twabController.burn(_owner, SafeCast.toUint96(_amount));
         emit Transfer(_owner, address(0), _amount);
     }
@@ -174,33 +184,38 @@ contract Harb is ERC20, ERC20Permit {
 
     /* ============ UBI stuff ============ */
 
-    function getUbiDue(address _account) public view returns (uint256) {
+    function getUbiDue(address _account) public view returns (uint256 amountDue, uint256 lastPeriodEndAt) {
         UbiTitle storage lastUbiTitle = ubiTitles[_account];
         return ubiDue(_account, lastUbiTitle.time, lastUbiTitle.sumTaxCollected);
     }
 
-    function ubiDue(address _account, uint256 lastTaxClaimed, uint256 _sumTaxCollected) internal view returns (uint256) {
-        if (lastTaxClaimed == 0 || block.timestamp - lastTaxClaimed < 60) {
-            return 0;
+    function ubiDue(address _account, uint256 lastTaxClaimed, uint256 _sumTaxCollected) internal view returns (uint256 amountDue, uint256 lastPeriodEndAt) {
+
+        lastPeriodEndAt = ((block.timestamp - PERIOD_OFFSET) / uint256(PERIOD_LENGTH)) * PERIOD_LENGTH + PERIOD_OFFSET - 1;
+
+        if (lastTaxClaimed == 0 || lastTaxClaimed > lastPeriodEndAt || lastPeriodEndAt - lastTaxClaimed < PERIOD_LENGTH) {
+            return (0, lastPeriodEndAt);
         }
-        uint256 accountTwab = twabController.getTwabBetween(address(this), _account, lastTaxClaimed, block.timestamp);
-        uint256 stakeTwab = twabController.getTwabBetween(address(this), stakingPool, lastTaxClaimed, block.timestamp);
+        uint256 accountTwab = twabController.getTwabBetween(address(this), _account, lastTaxClaimed, lastPeriodEndAt);
+        uint256 stakeTwab = twabController.getTwabBetween(address(this), stakingPool, lastTaxClaimed, lastPeriodEndAt);
 
 
-        uint256 poolTwab = twabController.getTwabBetween(address(this), address(pool), lastTaxClaimed, block.timestamp);
-        uint256 totalSupplyTwab = twabController.getTotalSupplyTwabBetween(address(this), lastTaxClaimed, block.timestamp);
+        uint256 poolTwab = twabController.getTwabBetween(address(this), address(pool), lastTaxClaimed, lastPeriodEndAt);
+        uint256 totalSupplyTwab = twabController.getTotalSupplyTwabBetween(address(this), lastTaxClaimed, lastPeriodEndAt);
 
         uint256 taxCollectedSinceLastClaim = sumTaxCollected - _sumTaxCollected;
-        return taxCollectedSinceLastClaim.mulDiv(accountTwab, (totalSupplyTwab - stakeTwab - poolTwab), Math.Rounding.Down);
+        amountDue = taxCollectedSinceLastClaim.mulDiv(accountTwab, (totalSupplyTwab - stakeTwab - poolTwab), Math.Rounding.Down);
     }
 
     function claimUbi(address _account) external {
         UbiTitle storage lastUbiTitle = ubiTitles[_account];
-        uint256 ubiAmountDue = ubiDue(_account, lastUbiTitle.time, lastUbiTitle.sumTaxCollected);
+        uint256 ubiAmountDue;
+        uint256 lastPeriodEndAt;
+        (ubiAmountDue, lastPeriodEndAt) = ubiDue(_account, lastUbiTitle.time, lastUbiTitle.sumTaxCollected);
 
         if (ubiAmountDue > 0) {
             ubiTitles[_account].sumTaxCollected = sumTaxCollected;
-            ubiTitles[_account].time = block.timestamp;
+            ubiTitles[_account].time = lastPeriodEndAt;
             twabController.transfer(TAX_POOL, _account, SafeCast.toUint96(ubiAmountDue));
             emit UbiClaimed(_account, ubiAmountDue);
         }

onchain/test/Harb.t.sol

@@ -135,7 +135,7 @@ contract HarbTest is Test {
         // claim tax
         {
             balanceBefore = harb.balanceOf(account);
-            uint256 ubiDue = harb.getUbiDue(account);
+            (uint256 ubiDue, ) = harb.getUbiDue(account);
             vm.prank(account);
             harb.claimUbi(account);
             assertFalse(ubiDue == 0, "No UBI paid");