From 8f387e52c2c235ed49dcfc9ef45c0132ff7915f0 Mon Sep 17 00:00:00 2001
From: Otto Cheung <otto@orca.so>
Date: Mon, 13 Jun 2022 10:11:33 -0700
Subject: [PATCH] revisions

---
 sdk/src/utils/public/pool-utils.ts | 24 +++++++++++++++---------
 sdk/src/utils/public/types.ts      |  6 +++---
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/sdk/src/utils/public/pool-utils.ts b/sdk/src/utils/public/pool-utils.ts
index 6184817..a511555 100644
--- a/sdk/src/utils/public/pool-utils.ts
+++ b/sdk/src/utils/public/pool-utils.ts
@@ -44,17 +44,23 @@ export class PoolUtil {
   /**
    * Given the intended token mint to swap, return the swap direction of a swap for a Whirlpool
    * @param pool The Whirlpool to evaluate the mint against
-   * @param inputTokenMint The token mint PublicKey the user intends to swap
-   * @returns The direction of the swap given the inputTokenMint. undefined if the token mint is not part of the trade pair of the pool.
+   * @param swapTokenMint The token mint PublicKey the user bases their swap against
+   * @param swapTokenIsInput Whether the swap token is the input token. (similar to amountSpecifiedIsInput from swap Ix)
+   * @returns The direction of the swap given the swapTokenMint. undefined if the token mint is not part of the trade pair of the pool.
    */
-  public static getSwapDirection(pool: WhirlpoolData, inputTokenMint: PublicKey) {
-    const tokenType = PoolUtil.getTokenType(pool, inputTokenMint);
-    if (tokenType === TokenType.TokenA) {
-      return SwapDirection.AtoB;
-    } else if (tokenType === TokenType.TokenB) {
-      return SwapDirection.BtoA;
+  public static getSwapDirection(
+    pool: WhirlpoolData,
+    swapTokenMint: PublicKey,
+    swapTokenIsInput: boolean
+  ) {
+    const tokenType = PoolUtil.getTokenType(pool, swapTokenMint);
+    if (!tokenType) {
+      return undefined;
     }
-    return undefined;
+
+    return (tokenType === TokenType.TokenA) === swapTokenIsInput
+      ? SwapDirection.AtoB
+      : SwapDirection.BtoA;
   }
 
   public static getFeeRate(feeRate: number): Percentage {
diff --git a/sdk/src/utils/public/types.ts b/sdk/src/utils/public/types.ts
index 97d1dee..cf60977 100644
--- a/sdk/src/utils/public/types.ts
+++ b/sdk/src/utils/public/types.ts
@@ -1,9 +1,9 @@
 export enum SwapDirection {
-  AtoB,
-  BtoA,
+  AtoB = "aToB",
+  BtoA = "bToA",
 }
 
 export enum TokenType {
-  TokenA,
+  TokenA = 1,
   TokenB,
 }