From 3f7471c93084652553aa099706dab992035ca51c Mon Sep 17 00:00:00 2001
From: Gary Wang <gary@garywang.net>
Date: Tue, 1 Sep 2020 07:03:05 -0700
Subject: [PATCH] Add option to
 findBaseTokenAccountsForOwner/findQuoteTokenAccountsForOwner to include
 unwrapped SOL for wrapped-SOL markets

---
 src/market.ts | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/src/market.ts b/src/market.ts
index f980764..1babc57 100644
--- a/src/market.ts
+++ b/src/market.ts
@@ -180,7 +180,18 @@ export class Market {
   async findBaseTokenAccountsForOwner(
     connection: Connection,
     ownerAddress: PublicKey,
+    includeUnwrappedSol = false,
   ): Promise<Array<{ pubkey: PublicKey; account: AccountInfo<Buffer> }>> {
+    if (this.baseMintAddress.equals(WRAPPED_SOL_MINT) && includeUnwrappedSol) {
+      const [wrapped, unwrapped] = await Promise.all([
+        this.findBaseTokenAccountsForOwner(connection, ownerAddress, false),
+        connection.getAccountInfo(ownerAddress),
+      ]);
+      if (unwrapped !== null) {
+        return [{ pubkey: ownerAddress, account: unwrapped }, ...wrapped];
+      }
+      return wrapped;
+    }
     return (
       await connection.getTokenAccountsByOwner(ownerAddress, {
         mint: this.baseMintAddress,
@@ -191,7 +202,18 @@ export class Market {
   async findQuoteTokenAccountsForOwner(
     connection: Connection,
     ownerAddress: PublicKey,
+    includeUnwrappedSol = false,
   ): Promise<{ pubkey: PublicKey; account: AccountInfo<Buffer> }[]> {
+    if (this.quoteMintAddress.equals(WRAPPED_SOL_MINT) && includeUnwrappedSol) {
+      const [wrapped, unwrapped] = await Promise.all([
+        this.findQuoteTokenAccountsForOwner(connection, ownerAddress, false),
+        connection.getAccountInfo(ownerAddress),
+      ]);
+      if (unwrapped !== null) {
+        return [{ pubkey: ownerAddress, account: unwrapped }, ...wrapped];
+      }
+      return wrapped;
+    }
     return (
       await connection.getTokenAccountsByOwner(ownerAddress, {
         mint: this.quoteMintAddress,