Skip to content

Commit 1edc622

Browse files
joaobrunoahjoaobrunoah
authored
Adjust Virtual Balances when Adding/Removing liquidity (#5)
1 parent 202e373 commit 1edc622

File tree

2 files changed

+201
-30
lines changed

2 files changed

+201
-30
lines changed

‎contracts/AclAmmPool.sol

+73-30
Original file line numberDiff line numberDiff line change
@@ -4,24 +4,16 @@
44
pragma solidity ^0.8.24;
55

66
import { ISwapFeePercentageBounds } from "@balancer-labs/v3-interfaces/contracts/vault/ISwapFeePercentageBounds.sol";
7-
import {
8-
IUnbalancedLiquidityInvariantRatioBounds
9-
} from "@balancer-labs/v3-interfaces/contracts/vault/IUnbalancedLiquidityInvariantRatioBounds.sol";
7+
import "@balancer-labs/v3-interfaces/contracts/vault/IUnbalancedLiquidityInvariantRatioBounds.sol";
8+
import { IBasePool } from "@balancer-labs/v3-interfaces/contracts/vault/IBasePool.sol";
109
import { IVault } from "@balancer-labs/v3-interfaces/contracts/vault/IVault.sol";
1110
import { IHooks } from "@balancer-labs/v3-interfaces/contracts/vault/IHooks.sol";
12-
import { IBasePool } from "@balancer-labs/v3-interfaces/contracts/vault/IBasePool.sol";
13-
import {
14-
Rounding,
15-
PoolSwapParams,
16-
SwapKind,
17-
HookFlags,
18-
TokenConfig,
19-
LiquidityManagement
20-
} from "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
11+
import "@balancer-labs/v3-interfaces/contracts/vault/VaultTypes.sol";
2112

13+
import { BasePoolAuthentication } from "@balancer-labs/v3-pool-utils/contracts/BasePoolAuthentication.sol";
14+
import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
2215
import { BalancerPoolToken } from "@balancer-labs/v3-vault/contracts/BalancerPoolToken.sol";
2316
import { Version } from "@balancer-labs/v3-solidity-utils/contracts/helpers/Version.sol";
24-
import { BasePoolAuthentication } from "@balancer-labs/v3-pool-utils/contracts/BasePoolAuthentication.sol";
2517
import { PoolInfo } from "@balancer-labs/v3-pool-utils/contracts/PoolInfo.sol";
2618
import { BaseHooks } from "@balancer-labs/v3-vault/contracts/BaseHooks.sol";
2719

@@ -37,6 +29,8 @@ contract AclAmmPool is
3729
Version,
3830
BaseHooks
3931
{
32+
using FixedPoint for uint256;
33+
4034
// uint256 private constant _MIN_SWAP_FEE_PERCENTAGE = 0.001e16; // 0.001%
4135
uint256 private constant _MIN_SWAP_FEE_PERCENTAGE = 0;
4236
uint256 private constant _MAX_SWAP_FEE_PERCENTAGE = 10e16; // 10%
@@ -103,7 +97,7 @@ contract AclAmmPool is
10397
);
10498
_lastTimestamp = block.timestamp;
10599
if (changed) {
106-
_virtualBalances = virtualBalances;
100+
_setVirtualBalances(virtualBalances);
107101

108102
if (_sqrtQ0State.startTime != 0) {
109103
_sqrtQ0State.startTime = 0;
@@ -135,16 +129,18 @@ contract AclAmmPool is
135129
/// @inheritdoc IHooks
136130
function getHookFlags() public pure override returns (HookFlags memory hookFlags) {
137131
hookFlags.shouldCallBeforeInitialize = true;
132+
hookFlags.shouldCallBeforeAddLiquidity = true;
133+
hookFlags.shouldCallBeforeRemoveLiquidity = true;
138134
}
139135

140136
/// @inheritdoc IHooks
141137
function onRegister(
142138
address,
143139
address,
144140
TokenConfig[] memory tokenConfig,
145-
LiquidityManagement calldata
141+
LiquidityManagement calldata liquidityManagement
146142
) public view override onlyVault returns (bool) {
147-
return tokenConfig.length == 2;
143+
return tokenConfig.length == 2 && liquidityManagement.disableUnbalancedLiquidity;
148144
}
149145

150146
/// @inheritdoc IHooks
@@ -153,7 +149,45 @@ contract AclAmmPool is
153149
bytes memory
154150
) public override onlyVault returns (bool) {
155151
_lastTimestamp = block.timestamp;
156-
_virtualBalances = AclAmmMath.initializeVirtualBalances(balancesScaled18, _calculateCurrentSqrtQ0());
152+
_setVirtualBalances(AclAmmMath.initializeVirtualBalances(balancesScaled18, _calculateCurrentSqrtQ0()));
153+
return true;
154+
}
155+
156+
/// @inheritdoc IHooks
157+
function onBeforeAddLiquidity(
158+
address,
159+
address pool,
160+
AddLiquidityKind,
161+
uint256[] memory,
162+
uint256 minBptAmountOut,
163+
uint256[] memory,
164+
bytes memory
165+
) public override onlyVault returns (bool) {
166+
uint256 totalSupply = _vault.totalSupply(pool);
167+
uint256 proportion = minBptAmountOut.divUp(totalSupply);
168+
uint256[] memory virtualBalances = _getLastVirtualBalances();
169+
virtualBalances[0] = virtualBalances[0].mulUp(FixedPoint.ONE + proportion);
170+
virtualBalances[1] = virtualBalances[1].mulUp(FixedPoint.ONE + proportion);
171+
_setVirtualBalances(virtualBalances);
172+
return true;
173+
}
174+
175+
/// @inheritdoc IHooks
176+
function onBeforeRemoveLiquidity(
177+
address,
178+
address pool,
179+
RemoveLiquidityKind,
180+
uint256 maxBptAmountIn,
181+
uint256[] memory,
182+
uint256[] memory,
183+
bytes memory
184+
) public override onlyVault returns (bool) {
185+
uint256 totalSupply = _vault.totalSupply(pool);
186+
uint256 proportion = maxBptAmountIn.divUp(totalSupply);
187+
uint256[] memory virtualBalances = _getLastVirtualBalances();
188+
virtualBalances[0] = virtualBalances[0].mulDown(FixedPoint.ONE - proportion);
189+
virtualBalances[1] = virtualBalances[1].mulDown(FixedPoint.ONE - proportion);
190+
_setVirtualBalances(virtualBalances);
157191
return true;
158192
}
159193

@@ -178,19 +212,8 @@ contract AclAmmPool is
178212
}
179213

180214
/// @inheritdoc IAclAmmPool
181-
function getLastVirtualBalances() external view returns (uint256[] memory virtualBalances) {
182-
(, , uint256[] memory balancesScaled18, ) = _vault.getPoolTokenInfo(address(this));
183-
184-
// Calculate virtual balances
185-
(virtualBalances, ) = AclAmmMath.getVirtualBalances(
186-
balancesScaled18,
187-
_virtualBalances,
188-
_c,
189-
_lastTimestamp,
190-
block.timestamp,
191-
_centerednessMargin,
192-
_sqrtQ0State
193-
);
215+
function getLastVirtualBalances() external view returns (uint256[] memory) {
216+
return _getLastVirtualBalances();
194217
}
195218

196219
/// @inheritdoc IAclAmmPool
@@ -212,6 +235,11 @@ contract AclAmmPool is
212235
_setSqrtQ0(newSqrtQ0, startTime, endTime);
213236
}
214237

238+
function _setVirtualBalances(uint256[] memory virtualBalances) internal {
239+
_virtualBalances = virtualBalances;
240+
_lastTimestamp = block.timestamp;
241+
}
242+
215243
function _setSqrtQ0(uint256 endSqrtQ0, uint256 startTime, uint256 endTime) internal {
216244
if (startTime > endTime) {
217245
revert GradualUpdateTimeTravel(startTime, endTime);
@@ -246,4 +274,19 @@ contract AclAmmPool is
246274
function _setCenterednessMargin(uint256 centerednessMargin) internal {
247275
_centerednessMargin = centerednessMargin;
248276
}
277+
278+
function _getLastVirtualBalances() internal view returns (uint256[] memory virtualBalances) {
279+
(, , uint256[] memory balancesScaled18, ) = _vault.getPoolTokenInfo(address(this));
280+
281+
// Calculate virtual balances
282+
(virtualBalances, ) = AclAmmMath.getVirtualBalances(
283+
balancesScaled18,
284+
_virtualBalances,
285+
_c,
286+
_lastTimestamp,
287+
block.timestamp,
288+
_centerednessMargin,
289+
_sqrtQ0State
290+
);
291+
}
249292
}

‎test/foundry/AclAmmLiquidity.t.sol

+128
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
// SPDX-License-Identifier: GPL-3.0-or-later
2+
3+
pragma solidity ^0.8.24;
4+
5+
import { console2 } from "forge-std/console2.sol";
6+
7+
import { FixedPoint } from "@balancer-labs/v3-solidity-utils/contracts/math/FixedPoint.sol";
8+
9+
import { BaseAclAmmTest } from "./utils/BaseAclAmmTest.sol";
10+
import { AclAmmPool } from "../../contracts/AclAmmPool.sol";
11+
import { AclAmmMath } from "../../contracts/lib/AclAmmMath.sol";
12+
13+
contract AclAmmLiquidityTest is BaseAclAmmTest {
14+
using FixedPoint for uint256;
15+
16+
uint256 constant _MAX_PRICE_ERROR_ABS = 5;
17+
uint256 constant _MAX_CENTEREDNESS_ERROR_ABS = 1e9;
18+
19+
function testAddLiquidity_Fuzz(
20+
uint256 exactBptAmountOut,
21+
uint256 initialDaiBalance,
22+
uint256 initialUsdcBalance
23+
) public {
24+
_setPoolBalances(initialDaiBalance, initialUsdcBalance);
25+
26+
uint256 totalSupply = vault.totalSupply(pool);
27+
exactBptAmountOut = bound(exactBptAmountOut, 1e6, 100 * totalSupply);
28+
29+
uint256[] memory maxAmountsIn = new uint256[](2);
30+
maxAmountsIn[daiIdx] = dai.balanceOf(alice);
31+
maxAmountsIn[usdcIdx] = usdc.balanceOf(alice);
32+
33+
uint256[] memory virtualBalancesBefore = AclAmmPool(pool).getLastVirtualBalances();
34+
(, , uint256[] memory balancesBefore, ) = vault.getPoolTokenInfo(pool);
35+
36+
vm.prank(alice);
37+
router.addLiquidityProportional(pool, maxAmountsIn, exactBptAmountOut, false, "");
38+
39+
uint256[] memory virtualBalancesAfter = AclAmmPool(pool).getLastVirtualBalances();
40+
(, , uint256[] memory balancesAfter, ) = vault.getPoolTokenInfo(pool);
41+
42+
// Check if virtual balances were correctly updated.
43+
uint256 proportion = exactBptAmountOut.divUp(totalSupply);
44+
assertEq(
45+
virtualBalancesAfter[daiIdx],
46+
virtualBalancesBefore[daiIdx].mulUp(FixedPoint.ONE + proportion),
47+
"DAI virtual balance does not match"
48+
);
49+
assertEq(
50+
virtualBalancesAfter[usdcIdx],
51+
virtualBalancesBefore[usdcIdx].mulUp(FixedPoint.ONE + proportion),
52+
"USDC virtual balance does not match"
53+
);
54+
55+
_checkPriceAndCenteredness(balancesBefore, balancesAfter, virtualBalancesBefore, virtualBalancesAfter);
56+
}
57+
58+
function testRemoveLiquidity_Fuzz(
59+
uint256 exactBptAmountIn,
60+
uint256 initialDaiBalance,
61+
uint256 initialUsdcBalance
62+
) public {
63+
_setPoolBalances(initialDaiBalance, initialUsdcBalance);
64+
65+
uint256 totalSupply = vault.totalSupply(pool);
66+
exactBptAmountIn = bound(exactBptAmountIn, 1e6, (9 * totalSupply) / 10);
67+
68+
uint256[] memory minAmountsOut = new uint256[](2);
69+
minAmountsOut[daiIdx] = 0;
70+
minAmountsOut[usdcIdx] = 0;
71+
72+
uint256[] memory virtualBalancesBefore = AclAmmPool(pool).getLastVirtualBalances();
73+
(, , uint256[] memory balancesBefore, ) = vault.getPoolTokenInfo(pool);
74+
75+
vm.prank(lp);
76+
router.removeLiquidityProportional(pool, exactBptAmountIn, minAmountsOut, false, "");
77+
78+
uint256[] memory virtualBalancesAfter = AclAmmPool(pool).getLastVirtualBalances();
79+
(, , uint256[] memory balancesAfter, ) = vault.getPoolTokenInfo(pool);
80+
81+
// Check if virtual balances were correctly updated.
82+
uint256 proportion = exactBptAmountIn.divUp(totalSupply);
83+
assertEq(
84+
virtualBalancesAfter[daiIdx],
85+
virtualBalancesBefore[daiIdx].mulDown(FixedPoint.ONE - proportion),
86+
"DAI virtual balance does not match"
87+
);
88+
assertEq(
89+
virtualBalancesAfter[usdcIdx],
90+
virtualBalancesBefore[usdcIdx].mulDown(FixedPoint.ONE - proportion),
91+
"USDC virtual balance does not match"
92+
);
93+
94+
_checkPriceAndCenteredness(balancesBefore, balancesAfter, virtualBalancesBefore, virtualBalancesAfter);
95+
}
96+
97+
function _checkPriceAndCenteredness(
98+
uint256[] memory balancesBefore,
99+
uint256[] memory balancesAfter,
100+
uint256[] memory virtualBalancesBefore,
101+
uint256[] memory virtualBalancesAfter
102+
) internal view {
103+
// Check if price is constant.
104+
uint256 daiPriceBefore = (balancesBefore[usdcIdx] + virtualBalancesBefore[usdcIdx]).divDown(
105+
balancesBefore[daiIdx] + virtualBalancesBefore[daiIdx]
106+
);
107+
uint256 daiPriceAfter = (balancesAfter[usdcIdx] + virtualBalancesAfter[usdcIdx]).divDown(
108+
balancesAfter[daiIdx] + virtualBalancesAfter[daiIdx]
109+
);
110+
assertApproxEqAbs(daiPriceAfter, daiPriceBefore, _MAX_PRICE_ERROR_ABS, "Price changed");
111+
112+
// Check if centeredness is constant.
113+
uint256 centerednessBefore = AclAmmMath.calculateCenteredness(balancesBefore, virtualBalancesBefore);
114+
uint256 centerednessAfter = AclAmmMath.calculateCenteredness(balancesAfter, virtualBalancesAfter);
115+
assertApproxEqAbs(centerednessAfter, centerednessBefore, _MAX_CENTEREDNESS_ERROR_ABS, "Centeredness changed");
116+
}
117+
118+
function _setPoolBalances(uint256 initialDaiBalance, uint256 initialUsdcBalance) internal {
119+
initialDaiBalance = bound(initialDaiBalance, 1e10, dai.balanceOf(address(vault)));
120+
initialUsdcBalance = bound(initialUsdcBalance, 1e10, usdc.balanceOf(address(vault)));
121+
122+
uint256[] memory initialBalances = new uint256[](2);
123+
initialBalances[daiIdx] = initialDaiBalance;
124+
initialBalances[usdcIdx] = initialUsdcBalance;
125+
126+
vault.manualSetPoolBalances(pool, initialBalances, initialBalances);
127+
}
128+
}

0 commit comments

Comments
 (0)