dedupe some code (Uniswap#129)

* dedupe some code * Apply suggestions from code review Co-authored-by: Noah Zinsmeister <[email protected]> * fix tests Co-authored-by: Noah Zinsmeister <[email protected]>
thaixuandang · Nov 18, 2020 · 22eec79 · 22eec79
1 parent 3f4972f
commit 22eec79
Show file tree

Hide file tree

Showing 9 changed files with 76 additions and 83 deletions.
diff --git a/contracts/UniswapV3Pair.sol b/contracts/UniswapV3Pair.sol
@@ -252,16 +252,20 @@ contract UniswapV3Pair is IUniswapV3Pair {
         return (price0CumulativeLast, price1CumulativeLast);
     }
 
-    function getValueAtPrice(FixedPoint.uq112x112 memory price, int256 liquidity) public pure returns (int256, int256) {
+    function getVirtualReservesDeltaAtPrice(FixedPoint.uq112x112 memory price, int256 liquidity)
+        public
+        pure
+        returns (int256, int256)
+    {
         if (liquidity == 0) return (0, 0);
 
         // we want to round up when liquidity is >0, i.e. being added
         if (liquidity > 0) {
-            (uint112 amount0, uint112 amount1) = PriceMath.getValueAtPriceRoundingUp(price, uint256(liquidity));
+            (uint112 amount0, uint112 amount1) = PriceMath.getVirtualReservesAtPrice(price, uint256(liquidity), true);
             return (amount0, amount1);
         } else {
             // we want to round down when liquidity is <0, i.e. being removed
-            (uint112 amount0, uint112 amount1) = PriceMath.getValueAtPriceRoundingDown(price, uint256(-liquidity));
+            (uint112 amount0, uint112 amount1) = PriceMath.getVirtualReservesAtPrice(price, uint256(-liquidity), false);
             return (-int256(amount0), -int256(amount1));
         }
     }
@@ -318,7 +322,7 @@ contract UniswapV3Pair is IUniswapV3Pair {
         FixedPoint.uq112x112 memory price = TickMath.getRatioAtTick(tick);
 
         // take the tokens
-        (int256 amount0, int256 amount1) = PriceMath.getValueAtPriceRoundingUp(price, liquidity);
+        (int256 amount0, int256 amount1) = PriceMath.getVirtualReservesAtPrice(price, liquidity, true);
         TransferHelper.safeTransferFrom(token0, msg.sender, address(this), uint256(amount0));
         TransferHelper.safeTransferFrom(token1, msg.sender, address(this), uint256(amount1));
 
@@ -407,10 +411,14 @@ contract UniswapV3Pair is IUniswapV3Pair {
 
             // if necessary, initialize both ticks and increment the position counter
             if (position.liquidity == 0 && params.liquidityDelta > 0) {
-                if (tickInfoLower.numPositions == 0) _initializeTick(params.tickLower, tickInfoLower);
-                tickInfoLower.numPositions++;
-                if (tickInfoUpper.numPositions == 0) _initializeTick(params.tickUpper, tickInfoUpper);
-                tickInfoUpper.numPositions++;
+                if (tickInfoLower.numPositions == 0) {
+                    _initializeTick(params.tickLower, tickInfoLower);
+                    tickInfoLower.numPositions = 1;
+                } else tickInfoLower.numPositions++;
+                if (tickInfoUpper.numPositions == 0) {
+                    _initializeTick(params.tickUpper, tickInfoUpper);
+                    tickInfoUpper.numPositions = 1;
+                } else tickInfoUpper.numPositions++;
             }
 
             {
@@ -470,11 +478,11 @@ contract UniswapV3Pair is IUniswapV3Pair {
         // calculate how much the specified liquidity delta is worth at the lower and upper ticks
         // amount0Lower :> amount0Upper
         // amount1Upper :> amount1Lower
-        (int256 amount0Lower, int256 amount1Lower) = getValueAtPrice(
+        (int256 amount0Lower, int256 amount1Lower) = getVirtualReservesDeltaAtPrice(
             TickMath.getRatioAtTick(params.tickLower),
             params.liquidityDelta
         );
-        (int256 amount0Upper, int256 amount1Upper) = getValueAtPrice(
+        (int256 amount0Upper, int256 amount1Upper) = getVirtualReservesDeltaAtPrice(
             TickMath.getRatioAtTick(params.tickUpper),
             params.liquidityDelta
         );
@@ -485,7 +493,10 @@ contract UniswapV3Pair is IUniswapV3Pair {
             amount0 = amount0.add(amount0Lower.sub(amount0Upper));
         } else if (tickCurrent < params.tickUpper) {
             // the current price is inside the passed range
-            (int256 amount0Current, int256 amount1Current) = getValueAtPrice(priceCurrent, params.liquidityDelta);
+            (int256 amount0Current, int256 amount1Current) = getVirtualReservesDeltaAtPrice(
+                priceCurrent,
+                params.liquidityDelta
+            );
             amount0 = amount0.add(amount0Current.sub(amount0Upper));
             amount1 = amount1.add(amount1Current.sub(amount1Lower));
 
@@ -584,9 +595,10 @@ contract UniswapV3Pair is IUniswapV3Pair {
                 step.fee = uint16(Math.max(feeFloor, getFee()));
 
                 // recompute reserves given the current price/liquidity
-                (step.reserve0Virtual, step.reserve1Virtual) = PriceMath.getValueAtPriceRoundingDown(
+                (step.reserve0Virtual, step.reserve1Virtual) = PriceMath.getVirtualReservesAtPrice(
                     state.price,
-                    state.liquidity
+                    state.liquidity,
+                    false
                 );
 
                 // compute the amount of input token required to push the price to the target (and max output token)

diff --git a/contracts/libraries/PriceMath.sol b/contracts/libraries/PriceMath.sol
@@ -33,45 +33,29 @@ library PriceMath {
     }
 
     // given a price and a liquidity amount, return the value of that liquidity at the price, rounded up
-    function getValueAtPriceRoundingUp(FixedPoint.uq112x112 memory price, uint256 liquidity)
-        internal
-        pure
-        returns (uint112 amount0, uint112 amount1)
-    {
+    function getVirtualReservesAtPrice(
+        FixedPoint.uq112x112 memory price,
+        uint256 liquidity,
+        bool roundUp
+    ) internal pure returns (uint112 reserve0, uint112 reserve1) {
         if (liquidity == 0) return (0, 0);
 
         uint8 safeShiftBits = ((255 - BitMath.mostSignificantBit(price._x)) / 2) * 2;
 
         uint256 priceScaled = uint256(price._x) << safeShiftBits; // price * 2**safeShiftBits
         uint256 priceScaledRoot = Babylonian.sqrt(priceScaled); // sqrt(priceScaled)
-        bool roundUp = priceScaledRoot**2 < priceScaled; // flag for whether priceScaledRoot needs to be rounded up
+        bool roundUpRoot = priceScaledRoot**2 < priceScaled; // flag for whether priceScaledRoot needs to be rounded up
 
         uint256 scaleFactor = uint256(1) << (56 + safeShiftBits / 2); // compensate for q112 and shifted bits under root
 
         // calculate amount0 := liquidity / sqrt(price) and amount1 := liquidity * sqrt(price)
-        amount0 = mulDivRoundingUp(liquidity, scaleFactor, priceScaledRoot).toUint112();
-        amount1 = mulDivRoundingUp(liquidity, priceScaledRoot + (roundUp ? 1 : 0), scaleFactor).toUint112();
-    }
-
-    // given a price and a liquidity amount, return the value of that liquidity at the price, rounded down
-    function getValueAtPriceRoundingDown(FixedPoint.uq112x112 memory price, uint256 liquidity)
-        internal
-        pure
-        returns (uint112 amount0, uint112 amount1)
-    {
-        if (liquidity == 0) return (0, 0);
-
-        uint8 safeShiftBits = ((255 - BitMath.mostSignificantBit(price._x)) / 2) * 2;
-
-        uint256 priceScaled = uint256(price._x) << safeShiftBits; // price * 2**safeShiftBits
-        uint256 priceScaledRoot = Babylonian.sqrt(priceScaled); // sqrt(priceScaled)
-        bool roundUp = priceScaledRoot**2 < priceScaled; // flag for whether priceScaledRoot needs to be rounded up
-
-        uint256 scaleFactor = uint256(1) << (56 + safeShiftBits / 2); // compensate for q112 and shifted bits under root
-
-        // calculate amount0 := liquidity / sqrt(price) and amount1 := liquidity * sqrt(price)
-        amount0 = FullMath.mulDiv(liquidity, scaleFactor, priceScaledRoot + (roundUp ? 1 : 0)).toUint112();
-        amount1 = FullMath.mulDiv(liquidity, priceScaledRoot, scaleFactor).toUint112();
+        if (roundUp) {
+            reserve0 = mulDivRoundingUp(liquidity, scaleFactor, priceScaledRoot).toUint112();
+            reserve1 = mulDivRoundingUp(liquidity, priceScaledRoot + (roundUpRoot ? 1 : 0), scaleFactor).toUint112();
+        } else {
+            reserve0 = FullMath.mulDiv(liquidity, scaleFactor, priceScaledRoot + (roundUpRoot ? 1 : 0)).toUint112();
+            reserve1 = FullMath.mulDiv(liquidity, priceScaledRoot, scaleFactor).toUint112();
+        }
     }
 
     function getInputToRatio(
@@ -83,7 +67,7 @@ library PriceMath {
         bool zeroForOne
     ) internal pure returns (uint112 amountIn, uint112 amountOut) {
         // estimate value of reserves at target price, rounding up
-        (uint112 reserve0Target, uint112 reserve1Target) = getValueAtPriceRoundingUp(priceTarget, liquidity);
+        (uint112 reserve0Target, uint112 reserve1Target) = getVirtualReservesAtPrice(priceTarget, liquidity, true);
 
         (amountIn, amountOut) = zeroForOne
             ? (reserve0Target - reserve0, reserve1 - reserve1Target)

diff --git a/contracts/test/PriceMathEchidnaTest.sol b/contracts/test/PriceMathEchidnaTest.sol
@@ -33,13 +33,15 @@ contract PriceMathEchidnaTest {
         require(price >= TickMath.getRatioAtTick(TickMath.MIN_TICK)._x);
         require(price <= TickMath.getRatioAtTick(TickMath.MAX_TICK)._x);
 
-        (uint112 amount0Up, uint112 amount1Up) = PriceMath.getValueAtPriceRoundingUp(
+        (uint112 amount0Up, uint112 amount1Up) = PriceMath.getVirtualReservesAtPrice(
             FixedPoint.uq112x112(price),
-            liquidity
+            liquidity,
+            true
         );
-        (uint112 amount0Down, uint112 amount1Down) = PriceMath.getValueAtPriceRoundingDown(
+        (uint112 amount0Down, uint112 amount1Down) = PriceMath.getVirtualReservesAtPrice(
             FixedPoint.uq112x112(price),
-            liquidity
+            liquidity,
+            false
         );
         assert(amount0Up >= amount0Down);
         assert(amount1Up >= amount1Down);
@@ -58,7 +60,7 @@ contract PriceMathEchidnaTest {
         require(lpFee > 0 && lpFee < PriceMath.LP_FEE_BASE);
 
         FixedPoint.uq112x112 memory price = FixedPoint.uq112x112(priceRaw);
-        (uint112 reserve0, uint112 reserve1) = PriceMath.getValueAtPriceRoundingDown(price, liquidity);
+        (uint112 reserve0, uint112 reserve1) = PriceMath.getVirtualReservesAtPrice(price, liquidity, false);
 
         require(reserve0 > 0 && reserve1 > 0);
 

diff --git a/contracts/test/PriceMathTest.sol b/contracts/test/PriceMathTest.sol
@@ -15,20 +15,12 @@ contract PriceMathTest {
         return PriceMath.getAmountOut(reserveIn, reserveOut, amountIn);
     }
 
-    function getValueAtPriceRoundingUp(FixedPoint.uq112x112 memory price, uint256 liquidity)
-        public
-        pure
-        returns (uint112 amount0, uint112 amount1)
-    {
-        return PriceMath.getValueAtPriceRoundingUp(price, liquidity);
-    }
-
-    function getValueAtPriceRoundingDown(FixedPoint.uq112x112 memory price, uint256 liquidity)
-        public
-        pure
-        returns (uint112 amount0, uint112 amount1)
-    {
-        return PriceMath.getValueAtPriceRoundingDown(price, liquidity);
+    function getVirtualReservesAtPrice(
+        FixedPoint.uq112x112 memory price,
+        uint256 liquidity,
+        bool roundUp
+    ) public pure returns (uint112 amount0, uint112 amount1) {
+        return PriceMath.getVirtualReservesAtPrice(price, liquidity, roundUp);
     }
 
     function getInputToRatio(

diff --git a/test/PriceMath.spec.ts b/test/PriceMath.spec.ts
@@ -33,8 +33,8 @@ describe('PriceMath', () => {
         const price = BigNumber.from('4294967297')
         const liquidity = BigNumber.from('18446744073709551615')
 
-        const [amount0Up, amount1Up] = await priceMath.getValueAtPriceRoundingUp([price], liquidity)
-        const [amount0Down, amount1Down] = await priceMath.getValueAtPriceRoundingDown([price], liquidity)
+        const [amount0Up, amount1Up] = await priceMath.getVirtualReservesAtPrice([price], liquidity, true)
+        const [amount0Down, amount1Down] = await priceMath.getVirtualReservesAtPrice([price], liquidity, false)
 
         expect(amount0Up).to.be.gte(amount0Down)
         expect(amount1Up).to.be.gte(amount1Down)
@@ -47,7 +47,7 @@ describe('PriceMath', () => {
         const liquidity = expandTo18Decimals(10)
         const price = encodePrice(expandTo18Decimals(100), expandTo18Decimals(1))
 
-        const [reserve0, reserve1] = await priceMath.getValueAtPriceRoundingDown([price], liquidity)
+        const [reserve0, reserve1] = await priceMath.getVirtualReservesAtPrice([price], liquidity, false)
 
         expect(reserve0).to.be.eq(expandTo18Decimals(1))
         expect(reserve1).to.be.eq(expandTo18Decimals(100))
@@ -142,7 +142,7 @@ describe('PriceMath', () => {
           let priceAfterSwap: BigNumber
 
           before('compute swap result', async () => {
-            ;[reserve0, reserve1] = await priceMath.getValueAtPriceRoundingDown([priceStarting], liquidity)
+            ;[reserve0, reserve1] = await priceMath.getVirtualReservesAtPrice([priceStarting], liquidity, false)
             ;[amountIn, amountOutMax] = await priceMath.getInputToRatio(
               reserve0,
               reserve1,

diff --git a/test/UniswapV3Pair.spec.ts b/test/UniswapV3Pair.spec.ts
@@ -211,17 +211,20 @@ describe('UniswapV3Pair', () => {
           })
 
           it('increments numPositions', async () => {
-            await pair.setPosition(-231, 0, 0, 100)
+            await pair.setPosition(-231, 5, 0, 100)
             expect((await pair.tickInfos(-231))[0]).to.eq(1)
-            await pair.setPosition(-231, 0, 1, 100)
+            expect((await pair.tickInfos(5))[0]).to.eq(1)
+            await pair.setPosition(-231, 5, 1, 100)
             expect((await pair.tickInfos(-231))[0]).to.eq(2)
+            expect((await pair.tickInfos(5))[0]).to.eq(2)
           })
 
           it('decrements numPositions', async () => {
             await pair.setPosition(-231, 0, 0, 100)
             await pair.setPosition(-231, 0, 1, 100)
             await pair.setPosition(-231, 0, 1, -100)
             expect((await pair.tickInfos(-231))[0]).to.eq(1)
+            expect((await pair.tickInfos(0))[0]).to.eq(1)
           })
 
           it('clears tick if last position is removed', async () => {

diff --git a/test/__snapshots__/PriceMath.spec.ts.snap b/test/__snapshots__/PriceMath.spec.ts.snap
@@ -1,8 +1,8 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`PriceMath #getInputToRatio edge cases gas: returns 0 if price is equal 1`] = `18231`;
+exports[`PriceMath #getInputToRatio edge cases gas: returns 0 if price is equal 1`] = `18267`;
 
-exports[`PriceMath #getInputToRatio invariants 1:49 to 1:75 at 60bps gas 1`] = `18191`;
+exports[`PriceMath #getInputToRatio invariants 1:49 to 1:75 at 60bps gas 1`] = `18227`;
 
 exports[`PriceMath #getInputToRatio invariants 1:49 to 1:75 at 60bps snapshot 1`] = `
 Object {
@@ -16,7 +16,7 @@ Object {
 }
 `;
 
-exports[`PriceMath #getInputToRatio invariants 1:49 to 1:100 at 200bps gas 1`] = `18221`;
+exports[`PriceMath #getInputToRatio invariants 1:49 to 1:100 at 200bps gas 1`] = `18257`;
 
 exports[`PriceMath #getInputToRatio invariants 1:49 to 1:100 at 200bps snapshot 1`] = `
 Object {
@@ -30,7 +30,7 @@ Object {
 }
 `;
 
-exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 30bps gas 1`] = `18278`;
+exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 30bps gas 1`] = `18314`;
 
 exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 30bps snapshot 1`] = `
 Object {
@@ -44,7 +44,7 @@ Object {
 }
 `;
 
-exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 60bps gas 1`] = `18278`;
+exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 60bps gas 1`] = `18314`;
 
 exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 60bps snapshot 1`] = `
 Object {
@@ -58,7 +58,7 @@ Object {
 }
 `;
 
-exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 60bps with small reserves gas 1`] = `18278`;
+exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 60bps with small reserves gas 1`] = `18314`;
 
 exports[`PriceMath #getInputToRatio invariants 1:100 to 1:50 at 60bps with small reserves snapshot 1`] = `
 Object {
@@ -72,7 +72,7 @@ Object {
 }
 `;
 
-exports[`PriceMath #getInputToRatio invariants 1:100 to 1:75 at 45bps gas 1`] = `18181`;
+exports[`PriceMath #getInputToRatio invariants 1:100 to 1:75 at 45bps gas 1`] = `18217`;
 
 exports[`PriceMath #getInputToRatio invariants 1:100 to 1:75 at 45bps snapshot 1`] = `
 Object {

diff --git a/test/__snapshots__/UniswapV3Factory.spec.ts.snap b/test/__snapshots__/UniswapV3Factory.spec.ts.snap
@@ -1,3 +1,3 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`UniswapV3Factory #createPair gas 1`] = `3739446`;
+exports[`UniswapV3Factory #createPair gas 1`] = `3722593`;
diff --git a/test/__snapshots__/UniswapV3Pair.spec.ts.snap b/test/__snapshots__/UniswapV3Pair.spec.ts.snap
@@ -6,16 +6,16 @@ exports[`UniswapV3Pair #getFee gas cost multiple votes median in middle 1`] = `6
 
 exports[`UniswapV3Pair #getFee gas cost uninitialized 1`] = `5887`;
 
-exports[`UniswapV3Pair #setPosition after initialization success cases above current price gas 1`] = `274443`;
+exports[`UniswapV3Pair #setPosition after initialization success cases above current price gas 1`] = `272928`;
 
-exports[`UniswapV3Pair #setPosition after initialization success cases below current price gas 1`] = `219915`;
+exports[`UniswapV3Pair #setPosition after initialization success cases below current price gas 1`] = `218400`;
 
-exports[`UniswapV3Pair #setPosition after initialization success cases including current price gas 1`] = `295471`;
+exports[`UniswapV3Pair #setPosition after initialization success cases including current price gas 1`] = `293993`;
 
-exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap0For1 gas 1`] = `172864`;
+exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap0For1 gas 1`] = `172928`;
 
-exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap0For1 gas large swap 1`] = `3716529`;
+exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap0For1 gas large swap 1`] = `3721777`;
 
-exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap1For0 gas 1`] = `151219`;
+exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap1For0 gas 1`] = `151283`;
 
-exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap1For0 gas large swap 1`] = `3694020`;
+exports[`UniswapV3Pair post-initialize (fee vote 1 - 0.10%) swap1For0 gas large swap 1`] = `3699268`;