From e72af254730f6e1c01dd69188e95249e6cc5cc18 Mon Sep 17 00:00:00 2001
From: shardul <shardul.aeer@somaiya.edu>
Date: Sat, 28 Aug 2021 19:20:54 +0530
Subject: [PATCH] feat: Save created OpenOrders and token accounts to cache

Allow swaps to be performed after necessary accounts for 'toMint'
and 'quoteMint' are created.
---
 src/components/Swap.tsx | 287 +++++++++++++++++++++++++++-------------
 src/context/Dex.tsx     |  14 +-
 src/context/Swap.tsx    |   1 +
 src/context/Token.tsx   |  91 ++++++++-----
 src/utils/tokens.ts     | 126 +++++++++++++-----
 5 files changed, 359 insertions(+), 160 deletions(-)

diff --git a/src/components/Swap.tsx b/src/components/Swap.tsx
index 4a3630e9..cecccdab 100644
--- a/src/components/Swap.tsx
+++ b/src/components/Swap.tsx
@@ -1,11 +1,10 @@
-import { useState } from "react";
+import { useMemo, useState } from "react";
 import {
   PublicKey,
   Keypair,
   Transaction,
   SystemProgram,
   Signer,
-  Account,
   SYSVAR_RENT_PUBKEY,
 } from "@solana/web3.js";
 import {
@@ -28,13 +27,15 @@ import { ExpandMore, ImportExportRounded } from "@material-ui/icons";
 import { useSwapContext, useSwapFair } from "../context/Swap";
 import {
   useDexContext,
-  useOpenOrders,
   useRouteVerbose,
   useMarket,
   FEE_MULTIPLIER,
+  _DexContext,
 } from "../context/Dex";
 import { useTokenMap } from "../context/TokenList";
 import {
+  addTokensToCache,
+  CachedToken,
   useMint,
   useOwnedTokenAccount,
   useTokenContext,
@@ -44,6 +45,7 @@ import TokenDialog from "./TokenDialog";
 import { SettingsButton } from "./Settings";
 import { InfoLabel } from "./Info";
 import { SOL_MINT, WRAPPED_SOL_MINT, DEX_PID } from "../utils/pubkeys";
+import { getTokenAddrressAndCreateIx } from "../utils/tokens";
 
 const useStyles = makeStyles((theme) => ({
   card: {
@@ -339,38 +341,63 @@ export function SwapButton() {
     isClosingNewAccounts,
     isStrict,
   } = useSwapContext();
-  const { swapClient, isLoaded: isDexLoaded } = useDexContext();
-  const { isLoaded: isTokensLoaded } = useTokenContext();
+  const {
+    swapClient,
+    isLoaded: isDexLoaded,
+    addOpenOrderAccount,
+    openOrders,
+  } = useDexContext();
+  const { isLoaded: isTokensLoaded, refreshTokenState } = useTokenContext();
+
+  // Token to be traded away
   const fromMintInfo = useMint(fromMint);
+  // End destination token
   const toMintInfo = useMint(toMint);
-  const openOrders = useOpenOrders();
+
   const route = useRouteVerbose(fromMint, toMint);
   const fromMarket = useMarket(
     route && route.markets ? route.markets[0] : undefined
   );
+
   const toMarket = useMarket(
     route && route.markets ? route.markets[1] : undefined
   );
-  const canSwap = useCanSwap();
-  const referral = useReferral(fromMarket);
-  const fair = useSwapFair();
-  let fromWallet = useOwnedTokenAccount(fromMint);
-  let toWallet = useOwnedTokenAccount(toMint);
+
+  const toWallet = useOwnedTokenAccount(toMint);
+  const fromWallet = useOwnedTokenAccount(fromMint);
+
+  // Intermediary token for multi-market swaps, eg. USDC in a SRM -> BTC swap
   const quoteMint = fromMarket && fromMarket.quoteMintAddress;
   const quoteMintInfo = useMint(quoteMint);
   const quoteWallet = useOwnedTokenAccount(quoteMint);
+
+  const canSwap = useCanSwap();
+  const referral = useReferral(fromMarket);
+  const fair = useSwapFair();
+
   const { isWrapSol, isUnwrapSol } = useIsWrapSol(fromMint, toMint);
-  const fromOpenOrders = fromMarket
-    ? openOrders.get(fromMarket?.address.toString())
-    : undefined;
-  const toOpenOrders = toMarket
-    ? openOrders.get(toMarket?.address.toString())
-    : undefined;
+
+  const fromOpenOrders = useMemo(() => {
+    return fromMarket
+      ? openOrders.get(fromMarket?.address.toString())
+      : undefined;
+  }, [fromMarket, openOrders]);
+
+  const toOpenOrders = useMemo(() => {
+    return toMarket ? openOrders.get(toMarket?.address.toString()) : undefined;
+  }, [toMarket, openOrders]);
+
   const disconnected = !swapClient.program.provider.wallet.publicKey;
+
   const needsCreateAccounts =
     !toWallet || !fromOpenOrders || (toMarket && !toOpenOrders);
 
   // Click handlers.
+
+  /**
+   * Find if OpenOrders or associated token accounts are required
+   * for the swap, then send a create transaction
+   */
   const sendCreateAccountsTransaction = async () => {
     if (!fromMintInfo || !toMintInfo) {
       throw new Error("Unable to calculate mint decimals");
@@ -380,95 +407,132 @@ export function SwapButton() {
     }
     const tx = new Transaction();
     const signers = [];
+
+    let toAssociatedPubkey!: PublicKey;
+    let quoteAssociatedPubkey!: PublicKey;
+
+    // Associated token account creation
     if (!toWallet) {
-      const associatedTokenPubkey = await Token.getAssociatedTokenAddress(
-        ASSOCIATED_TOKEN_PROGRAM_ID,
-        TOKEN_PROGRAM_ID,
-        toMint,
-        swapClient.program.provider.wallet.publicKey
-      );
-      tx.add(
-        Token.createAssociatedTokenAccountInstruction(
-          ASSOCIATED_TOKEN_PROGRAM_ID,
-          TOKEN_PROGRAM_ID,
+      const { tokenAddress, createTokenAddrIx } =
+        await getTokenAddrressAndCreateIx(
           toMint,
-          associatedTokenPubkey,
-          swapClient.program.provider.wallet.publicKey,
           swapClient.program.provider.wallet.publicKey
-        )
-      );
+        );
+      toAssociatedPubkey = tokenAddress;
+      tx.add(createTokenAddrIx);
     }
+
     if (!quoteWallet && !quoteMint.equals(toMint)) {
-      const quoteAssociatedPubkey = await Token.getAssociatedTokenAddress(
-        ASSOCIATED_TOKEN_PROGRAM_ID,
-        TOKEN_PROGRAM_ID,
-        quoteMint,
-        swapClient.program.provider.wallet.publicKey
-      );
-      tx.add(
-        Token.createAssociatedTokenAccountInstruction(
-          ASSOCIATED_TOKEN_PROGRAM_ID,
-          TOKEN_PROGRAM_ID,
+      const { tokenAddress, createTokenAddrIx } =
+        await getTokenAddrressAndCreateIx(
           quoteMint,
-          quoteAssociatedPubkey,
-          swapClient.program.provider.wallet.publicKey,
           swapClient.program.provider.wallet.publicKey
-        )
+        );
+      quoteAssociatedPubkey = tokenAddress;
+      tx.add(createTokenAddrIx);
+    }
+    // No point of initializing from wallet, as user won't have tokens there
+
+    // Helper functions for OpenOrders
+
+    /**
+     * Add instructions to init and create an OpenOrders account
+     * @param openOrdersKeypair
+     * @param market
+     * @param tx
+     */
+    async function getInitOpenOrdersIx(
+      openOrdersKeypair: Keypair,
+      market: PublicKey,
+      tx: Transaction
+    ) {
+      const createOoIx = await OpenOrders.makeCreateAccountTransaction(
+        swapClient.program.provider.connection,
+        market,
+        swapClient.program.provider.wallet.publicKey,
+        openOrdersKeypair.publicKey,
+        DEX_PID
       );
+      const initAcIx = swapClient.program.instruction.initAccount({
+        accounts: {
+          openOrders: openOrdersKeypair.publicKey,
+          authority: swapClient.program.provider.wallet.publicKey,
+          market: market,
+          dexProgram: DEX_PID,
+          rent: SYSVAR_RENT_PUBKEY,
+        },
+      });
+      tx.add(createOoIx);
+      tx.add(initAcIx);
+    }
+
+    /**
+     * Save data of newly created OpenOrders account in cache
+     * TODO: generate object client side to save a network call
+     * @param openOrdersAddress
+     */
+    async function saveOpenOrders(openOrdersAddress: PublicKey) {
+      const generatedOpenOrders = await OpenOrders.load(
+        swapClient.program.provider.connection,
+        openOrdersAddress,
+        DEX_PID
+      );
+      addOpenOrderAccount(generatedOpenOrders.market, generatedOpenOrders);
     }
+
+    // Open order accounts for to / from wallets. Generate if not already present
+    let ooFrom!: Keypair;
+    let ooTo!: Keypair;
     if (fromMarket && !fromOpenOrders) {
-      const ooFrom = Keypair.generate();
+      ooFrom = Keypair.generate();
+      await getInitOpenOrdersIx(ooFrom, fromMarket.address, tx);
       signers.push(ooFrom);
-      tx.add(
-        await OpenOrders.makeCreateAccountTransaction(
-          swapClient.program.provider.connection,
-          fromMarket.address,
-          swapClient.program.provider.wallet.publicKey,
-          ooFrom.publicKey,
-          DEX_PID
-        )
-      );
-      tx.add(
-        swapClient.program.instruction.initAccount({
-          accounts: {
-            openOrders: ooFrom.publicKey,
-            authority: swapClient.program.provider.wallet.publicKey,
-            market: fromMarket.address,
-            dexProgram: DEX_PID,
-            rent: SYSVAR_RENT_PUBKEY,
-          },
-        })
-      );
     }
     if (toMarket && !toOpenOrders) {
-      const ooTo = Keypair.generate();
+      ooTo = Keypair.generate();
+      await getInitOpenOrdersIx(ooTo, toMarket.address, tx);
       signers.push(ooTo);
-      tx.add(
-        await OpenOrders.makeCreateAccountTransaction(
-          swapClient.program.provider.connection,
-          toMarket.address,
-          swapClient.program.provider.wallet.publicKey,
-          ooTo.publicKey,
-          DEX_PID
-        )
-      );
-      tx.add(
-        swapClient.program.instruction.initAccount({
-          accounts: {
-            openOrders: ooTo.publicKey,
-            authority: swapClient.program.provider.wallet.publicKey,
-            market: toMarket.address,
-            dexProgram: DEX_PID,
-            rent: SYSVAR_RENT_PUBKEY,
-          },
-        })
-      );
     }
-    await swapClient.program.provider.send(tx, signers);
 
-    // TODO: update local data stores to add the newly created token
-    //       and open orders accounts.
+    try {
+      // Send transaction to create accounts
+      await swapClient.program.provider.send(tx, signers);
+
+      // Save OpenOrders to cache
+      if (ooFrom) {
+        await saveOpenOrders(ooFrom.publicKey);
+      }
+      if (ooTo) {
+        await saveOpenOrders(ooTo.publicKey);
+      }
+
+      // Save created associated token accounts to cache
+      const tokensToAdd: CachedToken[] = [];
+      if (toAssociatedPubkey) {
+        tokensToAdd.push(
+          getNewTokenAccountData(
+            toAssociatedPubkey,
+            toMint,
+            swapClient.program.provider.wallet.publicKey
+          )
+        );
+      }
+      if (quoteAssociatedPubkey && !quoteMint.equals(toMint)) {
+        tokensToAdd.push(
+          getNewTokenAccountData(
+            quoteAssociatedPubkey,
+            quoteMint,
+            swapClient.program.provider.wallet.publicKey
+          )
+        );
+      }
+      addTokensToCache(tokensToAdd);
+
+      // Refresh UI to display balance of the created token account
+      refreshTokenState();
+    } catch (error) {}
   };
+
   const sendWrapSolTransaction = async () => {
     if (!fromMintInfo || !toMintInfo) {
       throw new Error("Unable to calculate mint decimals");
@@ -528,6 +592,7 @@ export function SwapButton() {
     }
     await swapClient.program.provider.send(tx, signers);
   };
+
   const sendUnwrapSolTransaction = async () => {
     if (!fromMintInfo || !toMintInfo) {
       throw new Error("Unable to calculate mint decimals");
@@ -568,6 +633,7 @@ export function SwapButton() {
 
     await swapClient.program.provider.send(tx, signers);
   };
+
   const sendSwapTransaction = async () => {
     if (!fromMintInfo || !toMintInfo) {
       throw new Error("Unable to calculate mint decimals");
@@ -608,6 +674,12 @@ export function SwapButton() {
         ? toWallet.publicKey
         : undefined;
 
+      const fromOpenOrdersList = openOrders.get(fromMarket?.address.toString());
+      let fromOpenOrders: PublicKey | undefined = undefined;
+      if (fromOpenOrdersList) {
+        fromOpenOrders = fromOpenOrdersList[0].address;
+      }
+
       return await swapClient.swapTxs({
         fromMint,
         toMint,
@@ -618,7 +690,7 @@ export function SwapButton() {
         fromMarket,
         toMarket,
         // Automatically created if undefined.
-        fromOpenOrders: fromOpenOrders ? fromOpenOrders[0].address : undefined,
+        fromOpenOrders,
         toOpenOrders: toOpenOrders ? toOpenOrders[0].address : undefined,
         fromWallet: fromWalletAddr,
         toWallet: toWalletAddr,
@@ -675,11 +747,20 @@ export function SwapButton() {
         onClick={sendSwapTransaction}
         disabled={true}
       >
-        Swap
+        Loading
       </Button>
     );
   }
-  return needsCreateAccounts ? (
+  return !fromWallet ? (
+    <Button
+      variant="contained"
+      className={styles.swapButton}
+      onClick={sendSwapTransaction}
+      disabled={true}
+    >
+      No balance
+    </Button>
+  ) : needsCreateAccounts ? (
     <Button
       variant="contained"
       className={styles.swapButton}
@@ -799,3 +880,25 @@ function unwrapSol(
   );
   return { tx, signers: [] };
 }
+function getNewTokenAccountData(
+  toAssociatedPubkey: PublicKey,
+  mint: PublicKey,
+  owner: PublicKey
+): CachedToken {
+  return {
+    publicKey: toAssociatedPubkey,
+    account: {
+      address: toAssociatedPubkey,
+      mint,
+      owner,
+      amount: new u64(0),
+      delegate: null,
+      delegatedAmount: new u64(0),
+      isInitialized: true,
+      isFrozen: false,
+      isNative: false,
+      rentExemptReserve: null,
+      closeAuthority: null,
+    },
+  };
+}
diff --git a/src/context/Dex.tsx b/src/context/Dex.tsx
index c323cf4c..beae3fc5 100644
--- a/src/context/Dex.tsx
+++ b/src/context/Dex.tsx
@@ -35,10 +35,11 @@ type DexContext = {
   // Maps market address to open orders accounts.
   openOrders: Map<string, Array<OpenOrders>>;
   closeOpenOrders: (openOrder: OpenOrders) => void;
+  addOpenOrderAccount: (market: PublicKey, accountData: OpenOrders) => void;
   swapClient: SwapClient;
   isLoaded: boolean;
 };
-const _DexContext = React.createContext<DexContext | null>(null);
+export const _DexContext = React.createContext<DexContext | null>(null);
 
 export function DexContextProvider(props: any) {
   const [ooAccounts, setOoAccounts] = useState<Map<string, Array<OpenOrders>>>(
@@ -61,6 +62,16 @@ export function DexContextProvider(props: any) {
     setOoAccounts(newOoAccounts);
   };
 
+  const addOpenOrderAccount = async (
+    market: PublicKey,
+    accountData: OpenOrders
+  ) => {
+    const newOoAccounts = new Map(ooAccounts);
+    newOoAccounts.set(market.toString(), [accountData]);
+    setOoAccounts(newOoAccounts);
+    setIsLoaded(true);
+  };
+
   // Three operations:
   //
   // 1. Fetch all open orders accounts for the connected wallet.
@@ -169,6 +180,7 @@ export function DexContextProvider(props: any) {
       value={{
         openOrders: ooAccounts,
         closeOpenOrders,
+        addOpenOrderAccount,
         swapClient,
         isLoaded,
       }}
diff --git a/src/context/Swap.tsx b/src/context/Swap.tsx
index 88d2af95..56f248c7 100644
--- a/src/context/Swap.tsx
+++ b/src/context/Swap.tsx
@@ -217,6 +217,7 @@ export function useCanSwap(): boolean {
   const fromWallet = useOwnedTokenAccount(fromMint);
   const fair = useSwapFair();
   const route = useRouteVerbose(fromMint, toMint);
+
   if (route === null) {
     return false;
   }
diff --git a/src/context/Token.tsx b/src/context/Token.tsx
index 6826929e..cab8622a 100644
--- a/src/context/Token.tsx
+++ b/src/context/Token.tsx
@@ -1,46 +1,65 @@
 import React, { useContext, useState, useEffect } from "react";
-import * as assert from "assert";
 import { useAsync } from "react-async-hook";
 import { Provider, BN } from "@project-serum/anchor";
 import { PublicKey, Account } from "@solana/web3.js";
 import {
   MintInfo,
-  AccountInfo as TokenAccount,
+  AccountInfo as TokenAccountInfo,
   Token,
   TOKEN_PROGRAM_ID,
 } from "@solana/spl-token";
-import { getOwnedTokenAccounts, parseTokenAccountData } from "../utils/tokens";
+import assert from "assert";
+import {
+  getOwnedAssociatedTokenAccounts,
+  parseTokenAccountData,
+} from "../utils/tokens";
 import { SOL_MINT } from "../utils/pubkeys";
 
+export type CachedToken = {
+  publicKey: PublicKey; // Token account address
+  account: TokenAccountInfo;
+};
+
 export type TokenContext = {
   provider: Provider;
   isLoaded: boolean;
+  refreshTokenState(): void;
 };
 const _TokenContext = React.createContext<TokenContext | null>(null);
 
+const _OWNED_TOKEN_ACCOUNTS_CACHE: CachedToken[] = [];
+
+export function addTokensToCache(tokenList: CachedToken[]) {
+  _OWNED_TOKEN_ACCOUNTS_CACHE.push(...tokenList);
+}
+
 export function TokenContextProvider(props: any) {
   const provider = props.provider;
   const [, setRefresh] = useState(0);
   const [isLoaded, setIsLoaded] = useState(false);
 
+  function refreshTokenState() {
+    setRefresh((r) => r + 1);
+  }
+
   // Fetch all the owned token accounts for the wallet.
   useEffect(() => {
     if (!provider.wallet.publicKey) {
-      _OWNED_TOKEN_ACCOUNTS_CACHE.length = 0;
       setRefresh((r) => r + 1);
       return;
     }
-    // Fetch SPL tokens.
-    getOwnedTokenAccounts(provider.connection, provider.wallet.publicKey).then(
-      (accs) => {
-        if (accs) {
-          _OWNED_TOKEN_ACCOUNTS_CACHE.push(...accs);
-          setRefresh((r) => r + 1);
-        }
-        console.log("setting is loaded");
-        setIsLoaded(true);
+    // Fetch all SPL tokens belonging to the user
+    getOwnedAssociatedTokenAccounts(
+      provider.connection,
+      provider.wallet.publicKey
+    ).then((accs) => {
+      if (accs) {
+        // @ts-ignore
+        addTokensToCache(accs);
+        setRefresh((r) => r + 1);
       }
-    );
+      setIsLoaded(true);
+    });
     // Fetch SOL balance.
     provider.connection
       .getAccountInfo(provider.wallet.publicKey)
@@ -54,6 +73,7 @@ export function TokenContextProvider(props: any) {
               mint: SOL_MINT,
             },
           });
+
           setRefresh((r) => r + 1);
         }
       });
@@ -64,6 +84,7 @@ export function TokenContextProvider(props: any) {
       value={{
         provider,
         isLoaded,
+        refreshTokenState,
       }}
     >
       {props.children}
@@ -83,9 +104,11 @@ export function useTokenContext() {
 // Undefined => loading.
 export function useOwnedTokenAccount(
   mint?: PublicKey
-): { publicKey: PublicKey; account: TokenAccount } | null | undefined {
+): { publicKey: PublicKey; account: TokenAccountInfo } | null | undefined {
   const { provider } = useTokenContext();
+
   const [, setRefresh] = useState(0);
+
   const tokenAccounts = _OWNED_TOKEN_ACCOUNTS_CACHE.filter(
     (account) => mint && account.account.mint.equals(mint)
   );
@@ -100,6 +123,7 @@ export function useOwnedTokenAccount(
   );
 
   let tokenAccount = tokenAccounts[0];
+
   const isSol = mint?.equals(SOL_MINT);
 
   // Stream updates when the balance changes.
@@ -110,34 +134,39 @@ export function useOwnedTokenAccount(
       listener = provider.connection.onAccountChange(
         provider.wallet.publicKey,
         (info: { lamports: number }) => {
-          const token = {
+          const updatedTokenData = {
             amount: new BN(info.lamports),
             mint: SOL_MINT,
-          } as TokenAccount;
-          if (token.amount !== tokenAccount.account.amount) {
+          } as TokenAccountInfo;
+          if (updatedTokenData.amount !== tokenAccount.account.amount) {
             const index = _OWNED_TOKEN_ACCOUNTS_CACHE.indexOf(tokenAccount);
             assert.ok(index >= 0);
-            _OWNED_TOKEN_ACCOUNTS_CACHE[index].account = token;
+            _OWNED_TOKEN_ACCOUNTS_CACHE[index].account = updatedTokenData;
+
             setRefresh((r) => r + 1);
           }
         }
       );
     }
     // SPL tokens.
-    else if (tokenAccount) {
+    if (tokenAccount) {
       listener = provider.connection.onAccountChange(
         tokenAccount.publicKey,
         (info) => {
-          const token = parseTokenAccountData(info.data);
-          if (token.amount !== tokenAccount.account.amount) {
-            const index = _OWNED_TOKEN_ACCOUNTS_CACHE.indexOf(tokenAccount);
-            assert.ok(index >= 0);
-            _OWNED_TOKEN_ACCOUNTS_CACHE[index].account = token;
-            setRefresh((r) => r + 1);
+          if (info.data.length !== 0) {
+            const updatedTokenData = parseTokenAccountData(info.data);
+            if (updatedTokenData.amount !== tokenAccount.account.amount) {
+              const index = _OWNED_TOKEN_ACCOUNTS_CACHE.indexOf(tokenAccount);
+              assert.ok(index >= 0);
+              _OWNED_TOKEN_ACCOUNTS_CACHE[index].account = updatedTokenData;
+
+              setRefresh((r) => r + 1);
+            }
           }
         }
       );
     }
+    // Clean-up side effects. Called on re-rendering
     return () => {
       if (listener) {
         provider.connection.removeAccountChangeListener(listener);
@@ -145,10 +174,12 @@ export function useOwnedTokenAccount(
     };
   }, [provider.connection, tokenAccount]);
 
+  // Loading
   if (mint === undefined) {
     return undefined;
   }
 
+  // Account for given mint does not exist
   if (!isSol && tokenAccounts.length === 0) {
     return null;
   }
@@ -188,13 +219,9 @@ export function setMintCache(pk: PublicKey, account: MintInfo) {
   _MINT_CACHE.set(pk.toString(), new Promise((resolve) => resolve(account)));
 }
 
-// Cache storing all token accounts for the connected wallet provider.
-const _OWNED_TOKEN_ACCOUNTS_CACHE: Array<{
-  publicKey: PublicKey;
-  account: TokenAccount;
-}> = [];
-
 // Cache storing all previously fetched mint infos.
+// Initially SOL_MINT and mints of existing token accounts are stored
+// A mint is added for each new token opened
 // @ts-ignore
 const _MINT_CACHE = new Map<string, Promise<MintInfo>>([
   [SOL_MINT.toString(), { decimals: 9 }],
diff --git a/src/utils/tokens.ts b/src/utils/tokens.ts
index 2890c44c..fd2f3bae 100644
--- a/src/utils/tokens.ts
+++ b/src/utils/tokens.ts
@@ -5,62 +5,60 @@ import * as BufferLayout from "buffer-layout";
 import { BN } from "@project-serum/anchor";
 import {
   TOKEN_PROGRAM_ID,
+  ASSOCIATED_TOKEN_PROGRAM_ID,
+  Token,
   AccountInfo as TokenAccount,
 } from "@solana/spl-token";
 import { Connection, PublicKey } from "@solana/web3.js";
 import * as bs58 from "bs58";
+import { CachedToken } from "../context/Token";
 
-export async function getOwnedTokenAccounts(
+export async function getOwnedAssociatedTokenAccounts(
   connection: Connection,
   publicKey: PublicKey
 ) {
   let filters = getOwnedAccountsFilters(publicKey);
   // @ts-ignore
-  let resp = await connection._rpcRequest("getProgramAccounts", [
-    TOKEN_PROGRAM_ID.toBase58(),
-    {
-      commitment: connection.commitment,
-      filters,
-    },
-  ]);
-  if (resp.error) {
-    throw new Error(
-      "failed to get token accounts owned by " +
-        publicKey.toBase58() +
-        ": " +
-        resp.error.message
-    );
-  }
-  return resp.result
+  let resp = await connection.getProgramAccounts(TOKEN_PROGRAM_ID, {
+    commitment: connection.commitment,
+    filters,
+  });
+
+  const accs = resp
     .map(({ pubkey, account: { data, executable, owner, lamports } }: any) => ({
       publicKey: new PublicKey(pubkey),
       accountInfo: {
-        data: bs58.decode(data),
+        data,
         executable,
         owner: new PublicKey(owner),
         lamports,
       },
     }))
-    .filter(({ accountInfo }: any) => {
-      // TODO: remove this check once mainnet is updated
-      return filters.every((filter) => {
-        if (filter.dataSize) {
-          return accountInfo.data.length === filter.dataSize;
-        } else if (filter.memcmp) {
-          let filterBytes = bs58.decode(filter.memcmp.bytes);
-          return accountInfo.data
-            .slice(
-              filter.memcmp.offset,
-              filter.memcmp.offset + filterBytes.length
-            )
-            .equals(filterBytes);
-        }
-        return false;
-      });
-    })
     .map(({ publicKey, accountInfo }: any) => {
       return { publicKey, account: parseTokenAccountData(accountInfo.data) };
     });
+
+  return (
+    (
+      await Promise.all(
+        accs
+          // @ts-ignore
+          .map(async (ta) => {
+            const ata = await Token.getAssociatedTokenAddress(
+              ASSOCIATED_TOKEN_PROGRAM_ID,
+              TOKEN_PROGRAM_ID,
+              ta.account.mint,
+              publicKey
+            );
+            return [ta, ata];
+          })
+      )
+    )
+      // @ts-ignore
+      .filter(([ta, ata]) => ta.publicKey.equals(ata))
+      // @ts-ignore
+      .map(([ta]) => ta)
+  );
 }
 
 const ACCOUNT_LAYOUT = BufferLayout.struct([
@@ -95,3 +93,61 @@ function getOwnedAccountsFilters(publicKey: PublicKey) {
     },
   ];
 }
+
+/**
+ * Get associated token account for given mint, and instruction
+ * to generate this account
+ * @param mint
+ * @returns
+ */
+export async function getTokenAddrressAndCreateIx(
+  mint: PublicKey,
+  wallet: PublicKey
+) {
+  const tokenAddress = await Token.getAssociatedTokenAddress(
+    ASSOCIATED_TOKEN_PROGRAM_ID,
+    TOKEN_PROGRAM_ID,
+    mint,
+    wallet
+  );
+
+  const createTokenAddrIx = Token.createAssociatedTokenAccountInstruction(
+    ASSOCIATED_TOKEN_PROGRAM_ID,
+    TOKEN_PROGRAM_ID,
+    mint,
+    tokenAddress,
+    wallet,
+    wallet
+  );
+  return { tokenAddress, createTokenAddrIx };
+}
+
+/**
+ * Get account data for newly generated token account
+ * Object generated client side with balance 0 to save a network request
+ * @param tokenWallet
+ * @param mint
+ * @returns
+ */
+function getNewTokenAccountData(
+  tokenWallet: PublicKey,
+  mint: PublicKey,
+  owner: PublicKey
+): CachedToken {
+  return {
+    publicKey: tokenWallet,
+    account: {
+      address: tokenWallet,
+      mint,
+      owner,
+      amount: new BN(0),
+      delegate: null,
+      delegatedAmount: new BN(0),
+      isFrozen: false,
+      isInitialized: true,
+      isNative: false,
+      closeAuthority: null,
+      rentExemptReserve: null,
+    },
+  };
+}