From b945c87aeb9042eb031604d735bd62dd26987459 Mon Sep 17 00:00:00 2001
From: jaredvu <jaredvu@gmail.com>
Date: Wed, 15 Jan 2025 10:36:36 -0800
Subject: [PATCH] override size instead of update

---
 src/abacus-ts/calculators/accountActions.ts | 177 +++++++++++++-------
 src/abacus-ts/types/operationTypes.ts       |   2 +-
 2 files changed, 114 insertions(+), 65 deletions(-)

diff --git a/src/abacus-ts/calculators/accountActions.ts b/src/abacus-ts/calculators/accountActions.ts
index ba94c687e..499bb5e96 100644
--- a/src/abacus-ts/calculators/accountActions.ts
+++ b/src/abacus-ts/calculators/accountActions.ts
@@ -1,50 +1,137 @@
 import { produce } from 'immer';
 
-import { IndexerPositionSide } from '@/types/indexer/indexerApiGen';
+import {
+  IndexerAssetPositionResponseObject,
+  IndexerPositionSide,
+} from '@/types/indexer/indexerApiGen';
 
 import { MustBigNumber } from '@/lib/numbers';
 
-import { freshChildSubaccount } from '../lib/subaccountUtils';
+import { freshChildSubaccount, newUsdcAssetPosition } from '../lib/subaccountUtils';
 import {
   ModifyUsdcAssetPositionProps,
   SubaccountBatchedOperations,
   SubaccountOperations,
 } from '../types/operationTypes';
-import { ChildSubaccountData, ParentSubaccountData } from '../types/rawTypes';
-
-export function createUsdcDepositOperations({
-  subaccountNumber,
-  depositAmount,
-}: {
-  subaccountNumber: number;
-  depositAmount: string;
-}): SubaccountBatchedOperations {
+import { ParentSubaccountData } from '../types/rawTypes';
+
+function addUsdcAssetPosition(
+  parentSubaccount: ParentSubaccountData,
+  payload: Pick<IndexerAssetPositionResponseObject, 'side' | 'size' | 'subaccountNumber'>
+): ParentSubaccountData {
+  const { side, size, subaccountNumber } = payload;
+  return produce(parentSubaccount, (draftParentSubaccountData) => {
+    let childSubaccount = draftParentSubaccountData.childSubaccounts[subaccountNumber];
+
+    if (childSubaccount == null) {
+      // Upsert ChildSubaccountData into parentSubaccountData.childSubaccounts
+      const updatedChildSubaccount = freshChildSubaccount({
+        address: draftParentSubaccountData.address,
+        subaccountNumber,
+      });
+
+      childSubaccount = {
+        ...updatedChildSubaccount,
+        assetPositions: {
+          ...updatedChildSubaccount.assetPositions,
+          USDC: newUsdcAssetPosition({
+            side,
+            size,
+            subaccountNumber,
+          }),
+        },
+      };
+    } else {
+      if (childSubaccount.assetPositions.USDC == null) {
+        // Upsert USDC Asset Position
+        childSubaccount.assetPositions.USDC = newUsdcAssetPosition({
+          side,
+          size,
+          subaccountNumber,
+        });
+      } else {
+        if (childSubaccount.assetPositions.USDC.side !== side) {
+          const signedSizeBN = MustBigNumber(childSubaccount.assetPositions.USDC.size).minus(size);
+
+          if (signedSizeBN.lte(0)) {
+            // New size flips the Asset Position Side
+            childSubaccount.assetPositions.USDC.side =
+              side === IndexerPositionSide.LONG
+                ? IndexerPositionSide.SHORT
+                : IndexerPositionSide.LONG;
+            childSubaccount.assetPositions.USDC.size = signedSizeBN.abs().toString();
+          } else {
+            // Set the new size of the Asset Position
+            childSubaccount.assetPositions.USDC.size = signedSizeBN.toString();
+          }
+        } else {
+          // Side is maintained, add the size to the existing position
+          childSubaccount.assetPositions.USDC.size = MustBigNumber(
+            childSubaccount.assetPositions.USDC.size
+          )
+            .plus(size)
+            .toString();
+        }
+      }
+    }
+  });
+}
+
+export function createUsdcDepositOperations(
+  parentSubaccount: ParentSubaccountData,
+  {
+    subaccountNumber,
+    depositAmount,
+  }: {
+    subaccountNumber: number;
+    depositAmount: string;
+  }
+): SubaccountBatchedOperations {
+  const updatedParentSubaccountData = addUsdcAssetPosition(parentSubaccount, {
+    side: IndexerPositionSide.LONG,
+    size: depositAmount,
+    subaccountNumber,
+  });
+
+  if (updatedParentSubaccountData.childSubaccounts[subaccountNumber]?.assetPositions.USDC == null) {
+    throw new Error('USDC Asset Position was improperly modified');
+  }
+
   return {
     operations: [
       SubaccountOperations.ModifyUsdcAssetPosition({
         subaccountNumber,
-        changes: {
-          size: depositAmount,
-        },
+        changes: updatedParentSubaccountData.childSubaccounts[subaccountNumber].assetPositions.USDC,
       }),
     ],
   };
 }
 
-export function createUsdcWithdrawalOperations({
-  subaccountNumber,
-  withdrawAmount,
-}: {
-  subaccountNumber: number;
-  withdrawAmount: string;
-}): SubaccountBatchedOperations {
+export function createUsdcWithdrawalOperations(
+  parentSubaccount: ParentSubaccountData,
+  {
+    subaccountNumber,
+    withdrawAmount,
+  }: {
+    subaccountNumber: number;
+    withdrawAmount: string;
+  }
+): SubaccountBatchedOperations {
+  const updatedParentSubaccountData = addUsdcAssetPosition(parentSubaccount, {
+    side: IndexerPositionSide.SHORT,
+    size: withdrawAmount,
+    subaccountNumber,
+  });
+
+  if (updatedParentSubaccountData.childSubaccounts[subaccountNumber]?.assetPositions.USDC == null) {
+    throw new Error('USDC Asset Position was improperly modified');
+  }
+
   return {
     operations: [
       SubaccountOperations.ModifyUsdcAssetPosition({
         subaccountNumber,
-        changes: {
-          size: MustBigNumber(withdrawAmount).negated().toString(),
-        },
+        changes: updatedParentSubaccountData.childSubaccounts[subaccountNumber].assetPositions.USDC,
       }),
     ],
   };
@@ -55,49 +142,11 @@ function modifyUsdcAssetPosition(
   payload: ModifyUsdcAssetPositionProps
 ): ParentSubaccountData {
   const { subaccountNumber, changes } = payload;
-  if (!changes.size) return parentSubaccountData;
 
   return produce(parentSubaccountData, (draftParentSubaccountData) => {
-    const sizeBN = MustBigNumber(changes.size);
-
-    let childSubaccount: ChildSubaccountData | undefined =
-      draftParentSubaccountData.childSubaccounts[subaccountNumber];
-
-    if (childSubaccount != null) {
-      // Modify childSubaccount
-      if (childSubaccount.assetPositions.USDC != null) {
-        const size = MustBigNumber(childSubaccount.assetPositions.USDC.size)
-          .plus(sizeBN)
-          .toString();
-
-        childSubaccount.assetPositions.USDC.size = size;
-      } else if (sizeBN.gt(0)) {
-        // Upsert USDC Asset Position
-        childSubaccount.assetPositions.USDC = {
-          assetId: '0',
-          symbol: 'USDC',
-          size: sizeBN.toString(),
-          side: IndexerPositionSide.LONG,
-          subaccountNumber,
-        };
-      }
-    } else {
-      // Upsert ChildSubaccountData into parentSubaccountData.childSubaccounts
-      childSubaccount = freshChildSubaccount({
-        address: parentSubaccountData.address,
-        subaccountNumber,
-      });
-
-      childSubaccount.assetPositions.USDC = {
-        assetId: '0',
-        symbol: 'USDC',
-        size: sizeBN.toString(),
-        side: IndexerPositionSide.LONG,
-        subaccountNumber,
-      };
+    if (draftParentSubaccountData.childSubaccounts[subaccountNumber]?.assetPositions.USDC != null) {
+      draftParentSubaccountData.childSubaccounts[subaccountNumber].assetPositions.USDC = changes;
     }
-
-    draftParentSubaccountData.childSubaccounts[subaccountNumber] = childSubaccount;
   });
 }
 
diff --git a/src/abacus-ts/types/operationTypes.ts b/src/abacus-ts/types/operationTypes.ts
index 42698c910..f691ea630 100644
--- a/src/abacus-ts/types/operationTypes.ts
+++ b/src/abacus-ts/types/operationTypes.ts
@@ -17,7 +17,7 @@ export type ModifyPerpetualPositionProps = {
 
 export type ModifyUsdcAssetPositionProps = {
   subaccountNumber: number;
-  changes: Partial<Pick<IndexerAssetPositionResponseObject, 'size'>>;
+  changes: IndexerAssetPositionResponseObject;
 };
 
 export const SubaccountOperations = unionize(