add tests

pkg/interfaces/contracts/pool-aclamm/IAclAmmPool.sol

@@ -28,5 +28,7 @@ interface IAclAmmPool is IBasePool {
 
     function getLastTimestamp() external view returns (uint256);
 
+    function getCurrentSqrtQ0() external view returns (uint256);
+
     function setSqrtQ0(uint256 newSqrtQ0, uint256 startTime, uint256 endTime) external;
 }

pkg/pool-aclamm/contracts/AclAmmPool.sol

@@ -191,6 +191,11 @@ contract AclAmmPool is
         return _lastTimestamp;
     }
 
+    /// @inheritdoc IAclAmmPool
+    function getCurrentSqrtQ0() external view override returns (uint256) {
+        return _calculateCurrentSqrtQ0();
+    }
+
     /// @inheritdoc IAclAmmPool
     function setSqrtQ0(
         uint256 newSqrtQ0,

pkg/pool-aclamm/contracts/lib/AclAmmMath.sol

@@ -159,7 +159,7 @@ library AclAmmMath {
         uint256 numerator = ((endTime - currentTime) * startSqrtQ0) + ((currentTime - startTime) * endSqrtQ0);
         uint256 denominator = endTime - startTime;
 
-        return numerator.divDown(denominator);
+        return numerator / denominator;
     }
 
     function isAboveCenter(

pkg/pool-aclamm/contracts/test/AclAmmMathMock.sol

@@ -0,0 +1,125 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pragma solidity ^0.8.24;
+
+import { Rounding } from "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
+import { AclAmmMath } from "../lib/AclAmmMath.sol";
+
+contract AclAmmMathMock {
+    function computeInvariant(
+        uint256[] memory balancesScaled18,
+        uint256[] memory lastVirtualBalances,
+        uint256 c,
+        uint256 sqrtQ0,
+        uint256 lastTimestamp,
+        uint256 centernessMargin,
+        Rounding rounding
+    ) external view returns (uint256) {
+        return
+            AclAmmMath.computeInvariant(
+                balancesScaled18,
+                lastVirtualBalances,
+                c,
+                sqrtQ0,
+                lastTimestamp,
+                centernessMargin,
+                rounding
+            );
+    }
+
+    function calculateOutGivenIn(
+        uint256[] memory balancesScaled18,
+        uint256[] memory virtualBalances,
+        uint256 tokenInIndex,
+        uint256 tokenOutIndex,
+        uint256 amountGivenScaled18
+    ) external pure returns (uint256) {
+        return
+            AclAmmMath.calculateOutGivenIn(
+                balancesScaled18,
+                virtualBalances,
+                tokenInIndex,
+                tokenOutIndex,
+                amountGivenScaled18
+            );
+    }
+
+    function calculateInGivenOut(
+        uint256[] memory balancesScaled18,
+        uint256[] memory virtualBalances,
+        uint256 tokenInIndex,
+        uint256 tokenOutIndex,
+        uint256 amountGivenScaled18
+    ) external pure returns (uint256) {
+        return
+            AclAmmMath.calculateInGivenOut(
+                balancesScaled18,
+                virtualBalances,
+                tokenInIndex,
+                tokenOutIndex,
+                amountGivenScaled18
+            );
+    }
+
+    function initializeVirtualBalances(
+        uint256[] memory balancesScaled18,
+        uint256 sqrtQ0
+    ) external pure returns (uint256[] memory virtualBalances) {
+        return AclAmmMath.initializeVirtualBalances(balancesScaled18, sqrtQ0);
+    }
+
+    function getVirtualBalances(
+        uint256[] memory balancesScaled18,
+        uint256[] memory lastVirtualBalances,
+        uint256 c,
+        uint256 sqrtQ0,
+        uint256 lastTimestamp,
+        uint256 centernessMargin
+    ) external view returns (uint256[] memory virtualBalances, bool changed) {
+        return
+            AclAmmMath.getVirtualBalances(
+                balancesScaled18,
+                lastVirtualBalances,
+                c,
+                sqrtQ0,
+                lastTimestamp,
+                centernessMargin
+            );
+    }
+
+    function isPoolInRange(
+        uint256[] memory balancesScaled18,
+        uint256[] memory virtualBalances,
+        uint256 centernessMargin
+    ) external pure returns (bool) {
+        return AclAmmMath.isPoolInRange(balancesScaled18, virtualBalances, centernessMargin);
+    }
+
+    function calculateCenterness(
+        uint256[] memory balancesScaled18,
+        uint256[] memory virtualBalances
+    ) external pure returns (uint256) {
+        return AclAmmMath.calculateCenterness(balancesScaled18, virtualBalances);
+    }
+
+    function calculateSqrtQ0(
+        uint256 currentTime,
+        uint256 startSqrtQ0,
+        uint256 endSqrtQ0,
+        uint256 startTime,
+        uint256 endTime
+    ) external pure returns (uint256) {
+        return AclAmmMath.calculateSqrtQ0(currentTime, startSqrtQ0, endSqrtQ0, startTime, endTime);
+    }
+
+    function isAboveCenter(
+        uint256[] memory balancesScaled18,
+        uint256[] memory virtualBalances
+    ) external pure returns (bool) {
+        return AclAmmMath.isAboveCenter(balancesScaled18, virtualBalances);
+    }
+
+    function parseIncreaseDayRate(uint256 increaseDayRate) external pure returns (uint256) {
+        return AclAmmMath.parseIncreaseDayRate(increaseDayRate);
+    }
+}

pkg/pool-aclamm/test/AclAmmMath.test.ts

@@ -0,0 +1,49 @@
+import { Contract, BigNumberish } from 'ethers';
+import { deploy } from '@balancer-labs/v3-helpers/src/contract';
+import { expect } from 'chai';
+import { bn, fp } from '@balancer-labs/v3-helpers/src/numbers';
+import { expectEqualWithError } from '@balancer-labs/v3-helpers/src/test/relativeError';
+import { random } from 'lodash';
+import { calculateSqrtQ0 } from '@balancer-labs/v3-helpers/src/math/aclAmm';
+
+const MAX_RELATIVE_ERROR = 0.0001; // Max relative error
+
+// TODO: Test this math by checking extremes values for the amplification field (0 and infinite)
+// to verify that it equals constant sum and constant product (weighted) invariants.
+
+describe('AclAmmMath', function () {
+  let mock: Contract;
+
+  before(async function () {
+    mock = await deploy('AclAmmMathMock');
+  });
+
+  context('calculateSqrtQ0', () => {
+    it('should return endSqrtQ0Fp when currentTime > endTime', async () => {
+      const currentTime = 100;
+      const startSqrtQ0Fp = bn(100e18);
+      const endSqrtQ0Fp = bn(300e18);
+      const startTime = 1;
+      const endTime = 50;
+
+      const contractResult = await mock.calculateSqrtQ0(currentTime, startSqrtQ0Fp, endSqrtQ0Fp, startTime, endTime);
+      const mathResult = calculateSqrtQ0(currentTime, startSqrtQ0Fp, endSqrtQ0Fp, startTime, endTime);
+
+      expect(contractResult).to.equal(mathResult);
+      expect(contractResult).to.equal(endSqrtQ0Fp);
+    });
+
+    it('should return the correct value when currentTime < endTime', async () => {
+      const currentTime = 25;
+      const startSqrtQ0Fp = bn(100e18);
+      const endSqrtQ0Fp = bn(300e18);
+      const startTime = 1;
+      const endTime = 50;
+
+      const contractResult = await mock.calculateSqrtQ0(currentTime, startSqrtQ0Fp, endSqrtQ0Fp, startTime, endTime);
+      const mathResult = calculateSqrtQ0(currentTime, startSqrtQ0Fp, endSqrtQ0Fp, startTime, endTime);
+
+      expect(contractResult).to.equal(mathResult);
+    });
+  });
+});

pkg/pool-aclamm/test/foundry/AclAmmMath.t.sol

@@ -103,4 +103,48 @@ contract AclAmmMathTest is Test {
             assertEq(isAboveCenter, balance0.divDown(balance1) > virtualBalance0.divDown(virtualBalance1));
         }
     }
+
+    function testCalculateSqrtQ0__Fuzz(
+        uint256 currentTime,
+        uint256 startSqrtQ0,
+        uint256 endSqrtQ0,
+        uint256 startTime,
+        uint256 endTime
+    ) public view {
+        endTime = bound(endTime, 2, type(uint64).max);
+        startTime = bound(startTime, 1, endTime - 1);
+        currentTime = bound(currentTime, startTime, endTime);
+
+        console.log("currentTime: %s", currentTime);
+        console.log("startSqrtQ0: %s", startSqrtQ0);
+        console.log("endSqrtQ0: %s", endSqrtQ0);
+        console.log("startTime: %s", startTime);
+        console.log("endTime: %s", endTime);
+
+        endSqrtQ0 = bound(endSqrtQ0, 1, type(uint128).max);
+        startSqrtQ0 = bound(endSqrtQ0, 1, type(uint128).max);
+
+        uint256 sqrtQ0 = AclAmmMath.calculateSqrtQ0(currentTime, startSqrtQ0, endSqrtQ0, startTime, endTime);
+
+        currentTime++;
+        uint256 nextSqrtQ0 = AclAmmMath.calculateSqrtQ0(currentTime, startSqrtQ0, endSqrtQ0, startTime, endTime);
+
+        if (startSqrtQ0 >= endSqrtQ0) {
+            assertLe(nextSqrtQ0, sqrtQ0, "Next sqrtQ0 should be less than current sqrtQ0");
+        } else {
+            assertGe(nextSqrtQ0, sqrtQ0, "Next sqrtQ0 should be greater than current sqrtQ0");
+        }
+    }
+
+    function testCalculateSqrtQ0WhenCurrentTimeIsAfterEndTime() public pure {
+        uint256 startSqrtQ0 = 100;
+        uint256 endSqrtQ0 = 200;
+        uint256 startTime = 0;
+        uint256 endTime = 50;
+        uint256 currentTime = 100;
+
+        uint256 sqrtQ0 = AclAmmMath.calculateSqrtQ0(currentTime, startSqrtQ0, endSqrtQ0, startTime, endTime);
+
+        assertEq(sqrtQ0, endSqrtQ0, "SqrtQ0 should be equal to endSqrtQ0");
+    }
 }

pkg/pool-aclamm/test/foundry/AclAmmPool.t.sol

@@ -11,6 +11,7 @@ import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/Fixe
 
 import { BaseAclAmmTest } from "./utils/BaseAclAmmTest.sol";
 import { AclAmmPool } from "../../contracts/AclAmmPool.sol";
+import { AclAmmMath } from "../../contracts/lib/AclAmmMath.sol";
 
 contract AclAmmPoolTest is BaseAclAmmTest {
     using FixedPoint for uint256;
@@ -138,4 +139,30 @@ contract AclAmmPoolTest is BaseAclAmmTest {
         counter++;
         return uint256(keccak256(abi.encodePacked(block.prevrandao, block.timestamp, counter)));
     }
+
+    function testGetCurrentSqrtQ0() public {
+        uint256 sqrtQ0 = AclAmmPool(pool).getCurrentSqrtQ0();
+        assertEq(sqrtQ0, _DEFAULT_SQRT_Q0, "Invalid default sqrtQ0");
+    }
+
+    function testSetSqrtQ0() public {
+        uint256 newSqrtQ0 = 2e18;
+        uint256 startTime = block.timestamp;
+        uint256 duration = 1 hours;
+        uint256 endTime = block.timestamp + duration;
+
+        uint256 startSqrtQ0 = AclAmmPool(pool).getCurrentSqrtQ0();
+        vm.prank(admin);
+        AclAmmPool(pool).setSqrtQ0(newSqrtQ0, startTime, endTime);
+
+        skip(duration / 2);
+        uint256 sqrtQ0 = AclAmmPool(pool).getCurrentSqrtQ0();
+        uint256 mathSqrtQ0 = AclAmmMath.calculateSqrtQ0(block.timestamp, startSqrtQ0, newSqrtQ0, startTime, endTime);
+
+        assertEq(sqrtQ0, mathSqrtQ0, "SqrtQ0 not updated correctly");
+
+        skip(duration / 2 + 1);
+        sqrtQ0 = AclAmmPool(pool).getCurrentSqrtQ0();
+        assertEq(sqrtQ0, newSqrtQ0, "SqrtQ0 does not match new value");
+    }
 }

pkg/pool-aclamm/test/foundry/utils/BaseAclAmmTest.sol

@@ -62,6 +62,8 @@ contract BaseAclAmmTest is AclAmmPoolContractsDeployer, BaseVaultTest {
 
         PoolRoleAccounts memory roleAccounts;
 
+        roleAccounts = PoolRoleAccounts({ pauseManager: address(0), swapFeeManager: admin, poolCreator: address(0) });
+
         newPool = AclAmmPoolFactory(poolFactory).create(
             name,
             symbol,

pvt/helpers/src/math/aclAmm.ts

@@ -0,0 +1,19 @@
+import { BigNumberish } from 'ethers';
+import { bn } from '../numbers';
+
+export function calculateSqrtQ0(
+  currentTime: number,
+  startSqrtQ0Fp: BigNumberish,
+  endSqrtQ0Fp: BigNumberish,
+  startTime: number,
+  endTime: number
+): bigint {
+  if (currentTime > endTime) {
+    return bn(endSqrtQ0Fp);
+  }
+
+  const numerator = bn(endTime - currentTime) * bn(startSqrtQ0Fp) + bn(currentTime - startTime) * bn(endSqrtQ0Fp);
+  const denominator = bn(endTime - startTime);
+
+  return numerator / denominator;
+}