fix(#68): Refactors Remove and Allocate with better var names.

primitivefinance · Jun 20, 2021 · 12d1243 · 12d1243
1 parent b4f3deb
commit 12d1243
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 34 deletions.
diff --git a/contracts/PrimitiveEngine.sol b/contracts/PrimitiveEngine.sol
@@ -189,22 +189,21 @@ contract PrimitiveEngine is IPrimitiveEngine {
             (reserve.liquidity, reserve.reserveRisky, reserve.reserveStable);
 
         require(resLiquidity > 0, "Not initialized");
-        delRisky = (resRisky * delLiquidity) / resLiquidity;
-        delStable = (resStable * delLiquidity) / resLiquidity;
+        delRisky = (resRisky * delLiquidity) / resLiquidity; // amount of risky tokens to provide
+        delStable = (resStable * delLiquidity) / resLiquidity; // amount of stable tokens to provide
         require(delRisky * delStable > 0, "Deltas are 0");
 
         if (fromMargin) {
-            margins.withdraw(delRisky, delStable); // removes tokens from `msg.sender` margin account
+            margins.withdraw(delRisky, delStable); // removes tokens from `msg.sender` margin account, notice the mapping
         } else {
             uint256 balRisky = balanceRisky();
             uint256 balStable = balanceStable();
-            IPrimitiveLiquidityCallback(msg.sender).allocateCallback(delRisky, delStable, data);
+            IPrimitiveLiquidityCallback(msg.sender).allocateCallback(delRisky, delStable, data); // agnostic payment
             require(balanceRisky() >= balRisky + delRisky, "Not enough risky");
             require(balanceStable() >= balStable + delStable, "Not enough stable");
         }
 
-        bytes32 pid_ = pid;
-        Position.Data storage position = positions.fetch(owner, pid_);
+        Position.Data storage position = positions.fetch(owner, pid);
         position.allocate(delLiquidity); // increase position liquidity
         reserve.allocate(delRisky, delStable, delLiquidity, _blockTimestamp()); // increase reserves and liquidity
         emit Allocated(msg.sender, delRisky, delStable);
@@ -214,42 +213,35 @@ contract PrimitiveEngine is IPrimitiveEngine {
     function remove(
         bytes32 pid,
         uint256 delLiquidity,
-        bool isInternal,
+        bool toMargin,
         bytes calldata data
     ) external override lock returns (uint256 delRisky, uint256 delStable) {
-        require(delLiquidity > 0, "Cannot be 0");
+        require(delLiquidity > 0, "Cannot be 0"); // fail early
+
         Reserve.Data storage reserve = reserves[pid];
+        (uint256 resRisky, uint256 resStable, uint256 resLiquidity) =
+            (reserve.reserveRisky, reserve.reserveStable, reserve.liquidity);
 
-        uint256 nextRisky;
-        uint256 nextStable;
+        require(resLiquidity >= delLiquidity, "Above max burn");
+        delRisky = (resRisky * delLiquidity) / resLiquidity; // amount of risky to remove
+        delStable = (resStable * delLiquidity) / resLiquidity; // amount of stable to remove
+        require(delRisky * delStable > 0, "Deltas are 0");
 
-        {
-            (uint256 resRisky, uint256 resStable, uint256 resLiquidity) =
-                (reserve.reserveRisky, reserve.reserveStable, reserve.liquidity);
-            require(resLiquidity >= delLiquidity, "Above max burn");
-            delRisky = (resRisky * delLiquidity) / resLiquidity;
-            delStable = (resStable * delLiquidity) / resLiquidity;
-            require(delRisky * delStable > 0, "Deltas are 0");
-            nextRisky = resRisky - delRisky;
-            nextStable = resStable - delStable;
-        }
+        positions.remove(pid, delLiquidity); // update position liquidity, notice the fn call on the mapping
+        reserve.remove(delRisky, delStable, delLiquidity, _blockTimestamp()); // update global reserves
 
-        // Updated state
-        if (isInternal) {
+        if (toMargin) {
             Margin.Data storage margin = margins[msg.sender];
-            margin.deposit(delRisky, delStable);
+            margin.deposit(delRisky, delStable); // increase margin balance
         } else {
             uint256 balRisky = balanceRisky();
             uint256 balStable = balanceStable();
             IERC20(risky).safeTransfer(msg.sender, delRisky);
             IERC20(stable).safeTransfer(msg.sender, delStable);
-            IPrimitiveLiquidityCallback(msg.sender).removeCallback(delRisky, delStable, data);
+            IPrimitiveLiquidityCallback(msg.sender).removeCallback(delRisky, delStable, data); // agnostic withdrawals
             require(balanceRisky() >= balRisky - delRisky, "Not enough risky");
             require(balanceStable() >= balStable - delStable, "Not enough stable");
         }
-
-        positions.remove(pid, delLiquidity); // Update position liquidity
-        reserve.remove(delRisky, delStable, delLiquidity, _blockTimestamp());
         emit Removed(msg.sender, delRisky, delStable);
     }
 
@@ -412,7 +404,7 @@ contract PrimitiveEngine is IPrimitiveEngine {
         bytes32 pid,
         address owner,
         uint256 delLiquidity,
-        bool isInternal,
+        bool fromMargin,
         bytes calldata data
     ) external override lock returns (uint256 deltaRisky, uint256 deltaStable) {
         Reserve.Data storage res = reserves[pid];
@@ -424,7 +416,7 @@ contract PrimitiveEngine is IPrimitiveEngine {
         deltaRisky = (delLiquidity * res.reserveRisky) / res.liquidity;
         deltaStable = (delLiquidity * res.reserveStable) / res.liquidity;
 
-        if (isInternal) {
+        if (fromMargin) {
             margins.withdraw(delLiquidity - deltaRisky, deltaStable);
 
             res.allocate(deltaRisky, deltaStable, delLiquidity, _blockTimestamp());

diff --git a/package.json b/package.json
@@ -22,6 +22,7 @@
     "test:create": "hardhat test ./test/unit/primitiveEngine/effect/create.ts",
     "test:deposit": "hardhat test ./test/unit/primitiveEngine/effect/deposit.ts",
     "test:allocate": "hardhat test ./test/unit/primitiveEngine/effect/allocate.ts",
+    "test:remove": "hardhat test ./test/unit/primitiveEngine/effect/remove.ts",
     "test:factory": "hardhat test ./test/unit/primitiveFactory/effect/create.ts",
     "test:lib": "hardhat test ./test/unit/libraries",
     "prepare": "husky install"

diff --git a/test/unit/primitiveEngine/effect/allocate.ts b/test/unit/primitiveEngine/effect/allocate.ts
@@ -8,12 +8,12 @@ import { allocateFragment } from '../fragments'
 
 import loadContext from '../../context'
 
-const [strike, sigma, time, _] = [parseWei('1000').raw, 0.85 * PERCENTAGE, 31449600, parseWei('1100').raw]
+const [strike, sigma, time, _] = [parseWei('1000').raw, 0.85 * PERCENTAGE, 1655655140, parseWei('1100').raw]
 const empty: BytesLike = constants.HashZero
 
 describe('allocate', function () {
   before(async function () {
-    loadContext(waffle.provider, ['engineCreate', 'engineDeposit', 'engineAllocate'], allocateFragment)
+    await loadContext(waffle.provider, ['engineCreate', 'engineDeposit', 'engineAllocate'], allocateFragment)
   })
 
   describe('when the parameters are valid', function () {

diff --git a/test/unit/primitiveEngine/effect/remove.ts b/test/unit/primitiveEngine/effect/remove.ts
@@ -8,18 +8,21 @@ import { removeFragment } from '../fragments'
 
 import loadContext from '../../context'
 
-const [strike, sigma, time, _] = [parseWei('1000').raw, 0.85 * PERCENTAGE, 31449600, parseWei('1100').raw]
+const [strike, sigma, time, _] = [parseWei('1000').raw, 0.85 * PERCENTAGE, 1655655140, parseWei('1100').raw]
+
 const empty: BytesLike = constants.HashZero
 
 describe('remove', function () {
   before(async function () {
-    loadContext(waffle.provider, ['engineCreate', 'engineDeposit', 'engineAllocate', 'engineRemove'], removeFragment)
+    await loadContext(waffle.provider, ['engineCreate', 'engineDeposit', 'engineAllocate', 'engineRemove'], removeFragment)
   })
 
   describe('when the parameters are valid', function () {
     it('removes 1 liquidity share and deposits the resultant risky and stable to margin', async function () {
       const pid = await this.contracts.engine.getPoolId(strike, sigma, time)
       const posid = await this.contracts.engineRemove.getPosition(pid)
+      console.log('removing from pid', pid.slice(0, 5))
+      console.log('engine', this.contracts.engine.address.slice(0, 5))
       await this.contracts.engineRemove.removeToMargin(pid, parseWei('1').raw, empty)
 
       expect(await this.contracts.engine.positions(posid)).to.be.deep.eq([

diff --git a/test/unit/primitiveEngine/fragments.ts b/test/unit/primitiveEngine/fragments.ts
@@ -2,7 +2,7 @@ import { Wallet, constants } from 'ethers'
 import { Contracts } from '../../../types'
 import { parseWei, PERCENTAGE } from '../../shared/Units'
 
-const [strike, sigma, time, riskyPrice] = [parseWei('1000').raw, 0.85 * PERCENTAGE, 31449600, parseWei('1100').raw]
+const [strike, sigma, time, riskyPrice] = [parseWei('1000').raw, 0.85 * PERCENTAGE, 1655655140, parseWei('1100').raw]
 const empty = constants.HashZero
 
 export async function createFragment(signers: Wallet[], contracts: Contracts): Promise<void> {