[Pool] Support Pool Deposit / Withdraw (#35)

* pool deposit/withdraw interface + deposit impl

* pool withdraw impl + minor fixes

* move OrcaPoolToken to types.ts

* fix: wrong pool token mint

* deposit/withdraw interface comments

* edit comments and params for deposit/withdraw interface

* Rebase: support wallet signing for deposit/withdraw

* minor name changes

* remove .DS_Store

Co-authored-by: scuba <scuba@example.com>

src/model/orca/pool/orca-pool.ts

@@ -14,14 +14,17 @@ import {
   TransactionBuilder,
   OrcaPool,
   OrcaToken,
+  OrcaPoolToken,
   Quote,
   TransactionPayload,
   Percentage,
   resolveOrCreateAssociatedTokenAddress,
 } from "../../../public";
 import {
   createApprovalInstruction,
+  createDepositInstruction,
   createSwapInstruction,
+  createWithdrawInstruction,
 } from "../../../public/utils/web3/instructions/pool-instructions";
 import { Owner } from "../../../public/utils/web3/key-utils";
 import { QuotePoolParams, QuoteBuilderFactory } from "../../quote/quote-builder";
@@ -36,12 +39,12 @@ export class OrcaPoolImpl implements OrcaPool {
     this.poolParams = config;
   }
 
-  public getTokenA(): OrcaToken {
+  public getTokenA(): OrcaPoolToken {
     const tokenId = this.poolParams.tokenIds[0];
     return this.poolParams.tokens[tokenId];
   }
 
-  public getTokenB(): OrcaToken {
+  public getTokenB(): OrcaPoolToken {
     const tokenId = this.poolParams.tokenIds[1];
     return this.poolParams.tokens[tokenId];
   }
@@ -185,4 +188,167 @@ export class OrcaPoolImpl implements OrcaPool {
       .addInstruction(swapInstruction)
       .build();
   }
+
+  public async deposit(
+    owner: Keypair | PublicKey,
+    maxTokenAIn: Decimal | OrcaU64,
+    maxTokenBIn: Decimal | OrcaU64,
+    minPoolTokenAmountOut: Decimal | OrcaU64
+  ): Promise<TransactionPayload> {
+    const _owner = new Owner(owner);
+    const ownerAddress = _owner.publicKey;
+
+    const tokenA = this.getTokenA();
+    const tokenB = this.getTokenB();
+
+    const maxTokenAIn_U64 = U64Utils.toTokenU64(maxTokenAIn, tokenA, "maxTokenAIn");
+    const maxTokenBIn_U64 = U64Utils.toTokenU64(maxTokenBIn, tokenB, "maxTokenBIn");
+    const minPoolTokenAmountOut_U64 = U64Utils.toPoolU64(
+      minPoolTokenAmountOut,
+      this.poolParams,
+      "poolTokenAmount"
+    );
+
+    // If tokenA is SOL, this will create a new wSOL account
+    // Otherwise, get tokenA's associated token account
+    const { address: userTokenAPublicKey, ...resolveTokenAInstrucitons } =
+      await resolveOrCreateAssociatedTokenAddress(
+        this.connection,
+        _owner,
+        tokenA.mint,
+        maxTokenAIn_U64
+      );
+
+    // If tokenB is SOL, this will create a new wSOL account
+    // Otherwise, get tokenB's associated token account
+    const { address: userTokenBPublicKey, ...resolveTokenBInstrucitons } =
+      await resolveOrCreateAssociatedTokenAddress(
+        this.connection,
+        _owner,
+        tokenB.mint,
+        maxTokenBIn_U64
+      );
+
+    // If the user lacks the pool token account, create it
+    const { address: userPoolTokenPublicKey, ...resolvePoolTokenInstructions } =
+      await resolveOrCreateAssociatedTokenAddress(
+        this.connection,
+        _owner,
+        this.poolParams.poolTokenMint
+      );
+
+    // Approve transfer of the tokens being deposited
+    const { userTransferAuthority, ...transferTokenAInstruction } = createApprovalInstruction(
+      ownerAddress,
+      maxTokenAIn_U64,
+      userTokenAPublicKey
+    );
+    const { ...transferTokenBInstruction } = createApprovalInstruction(
+      ownerAddress,
+      maxTokenBIn_U64,
+      userTokenBPublicKey,
+      userTransferAuthority
+    );
+
+    // Create the deposit instruction
+    const depositInstruction = await createDepositInstruction(
+      this.poolParams,
+      userTransferAuthority.publicKey,
+      userTokenAPublicKey,
+      userTokenBPublicKey,
+      userPoolTokenPublicKey,
+      minPoolTokenAmountOut_U64,
+      maxTokenAIn_U64,
+      maxTokenBIn_U64,
+      tokenA.addr,
+      tokenB.addr,
+      _owner
+    );
+
+    return await new TransactionBuilder(this.connection, ownerAddress, _owner)
+      .addInstruction(resolveTokenAInstrucitons)
+      .addInstruction(resolveTokenBInstrucitons)
+      .addInstruction(resolvePoolTokenInstructions)
+      .addInstruction(transferTokenAInstruction)
+      .addInstruction(transferTokenBInstruction)
+      .addInstruction(depositInstruction)
+      .build();
+  }
+
+  public async withdraw(
+    owner: Keypair | PublicKey,
+    poolTokenAmountIn: Decimal | OrcaU64,
+    minTokenAOut: Decimal | OrcaU64,
+    minTokenBOut: Decimal | OrcaU64
+  ): Promise<TransactionPayload> {
+    const _owner = new Owner(owner);
+    const ownerAddress = _owner.publicKey;
+
+    const tokenA = this.getTokenA();
+    const tokenB = this.getTokenB();
+
+    const minTokenAOut_U64 = U64Utils.toTokenU64(minTokenAOut, tokenA, "minTokenAOut");
+    const minTokenBOut_U64 = U64Utils.toTokenU64(minTokenBOut, tokenB, "minTokenBOut");
+    const poolTokenAmountIn_U64 = U64Utils.toPoolU64(
+      poolTokenAmountIn,
+      this.poolParams,
+      "poolTokenAmount"
+    );
+
+    // Create a token account for tokenA, if necessary
+    const { address: userTokenAPublicKey, ...resolveTokenAInstrucitons } =
+      await resolveOrCreateAssociatedTokenAddress(
+        this.connection,
+        _owner,
+        tokenA.mint,
+        minTokenAOut_U64
+      );
+
+    // Create a token account for tokenB, if necessary
+    const { address: userTokenBPublicKey, ...resolveTokenBInstrucitons } =
+      await resolveOrCreateAssociatedTokenAddress(
+        this.connection,
+        _owner,
+        tokenB.mint,
+        minTokenBOut_U64
+      );
+
+    // Get user's poolToken token account
+    const { address: userPoolTokenPublicKey, ...resolvePoolTokenInstructions } =
+      await resolveOrCreateAssociatedTokenAddress(
+        this.connection,
+        _owner,
+        this.poolParams.poolTokenMint
+      );
+
+    // Approve transfer of pool token
+    const { userTransferAuthority, ...transferPoolTokenInstruction } = createApprovalInstruction(
+      ownerAddress,
+      poolTokenAmountIn_U64,
+      userPoolTokenPublicKey
+    );
+
+    // Create the withdraw instruction
+    const withdrawInstruction = await createWithdrawInstruction(
+      this.poolParams,
+      userTransferAuthority.publicKey,
+      userTokenAPublicKey,
+      userTokenBPublicKey,
+      userPoolTokenPublicKey,
+      poolTokenAmountIn_U64,
+      minTokenAOut_U64,
+      minTokenBOut_U64,
+      tokenA.addr,
+      tokenB.addr,
+      _owner
+    );
+
+    return await new TransactionBuilder(this.connection, ownerAddress, _owner)
+      .addInstruction(resolveTokenAInstrucitons)
+      .addInstruction(resolveTokenBInstrucitons)
+      .addInstruction(resolvePoolTokenInstructions)
+      .addInstruction(transferPoolTokenInstruction)
+      .addInstruction(withdrawInstruction)
+      .build();
+  }
 }

src/model/orca/pool/pool-types.ts

@@ -1,5 +1,5 @@
 import { PublicKey } from "@solana/web3.js";
-import { OrcaToken } from "../../..";
+import { OrcaPoolToken } from "../../..";
 import { Percentage } from "../../../public";
 
 /**
@@ -30,14 +30,6 @@ export type OrcaPoolParams = {
   amp?: number;
 };
 
-/**
- * An Orca Token within an OrcaPool
- * @param addr The public key for this token for this Orca Pool
- */
-export type OrcaPoolToken = OrcaToken & {
-  addr: PublicKey;
-};
-
 export enum CurveType {
   ConstantProduct,
   ConstantPrice,

src/public/pools/types.ts

@@ -1,4 +1,4 @@
-import { Keypair, PublicKey, TransactionSignature } from "@solana/web3.js";
+import { Keypair, PublicKey } from "@solana/web3.js";
 import Decimal from "decimal.js";
 import { OrcaU64 } from "..";
 import { TransactionPayload } from "../utils";
@@ -11,13 +11,13 @@ export type OrcaPool = {
    * Query the token of tokenA in this pool.
    * @returns Returns the token id of tokenA in this pool
    */
-  getTokenA: () => OrcaToken;
+  getTokenA: () => OrcaPoolToken;
 
   /**
    * Query the token of tokenB in this pool.
    * @returns Returns the token id of tokenB in this pool
    */
-  getTokenB: () => OrcaToken;
+  getTokenB: () => OrcaPoolToken;
 
   /**
    * Query the balance for an user address
@@ -68,6 +68,48 @@ export type OrcaPool = {
     amountIn: Decimal | OrcaU64,
     minimumAmountOut: Decimal | OrcaU64
   ) => Promise<TransactionPayload>;
+
+  /**
+   * Perform a deposit: send tokenA and tokenB, and receive a poolToken in return.
+   * Fee for the transaction will be paid by the owner's wallet.
+   *
+   * NOTE:
+   * 1. Associated Token Address initialization instructions will be appended if the ATA of the specified token does not exist in the user's wallet
+   * 2. OrcaU64 must have the same scale as the corresponding token scale value
+   *
+   * @param owner The keypair for the user's wallet or just the user's public key
+   * @param maxTokenAIn The maximum amount of tokenA to send
+   * @param maxTokenBIn The maximum amount of tokenB to send
+   * @param minPoolTokenAmountOut The amount of poolToken to receive
+   * @return The transaction signature of the deposit instruction
+   */
+  deposit: (
+    owner: Keypair | PublicKey,
+    maxTokenAIn: Decimal | OrcaU64,
+    maxTokenBIn: Decimal | OrcaU64,
+    minPoolTokenAmountOut: Decimal | OrcaU64
+  ) => Promise<TransactionPayload>;
+
+  /**
+   * Perform a withdraw: send poolToken, and receive tokenA and tokenB in return.
+   * Fee for the transaction will be paid by the owner's wallet.
+   *
+   * NOTE:
+   * 1. Associated Token Address initialization instructions will be appended if the ATA of the specified token does not exist in the user's wallet
+   * 2. OrcaU64 must have the same scale as the corresponding token scale value
+   *
+   * @param owner The keypair for the user's wallet or just the user's public key
+   * @param poolTokenAmountIn The amount of poolToken to send
+   * @param minTokenAOut The minimum amount of tokenA to receive
+   * @param minTokenBOut The minimum amount of tokenB to receive
+   * @return The transaction signature of the withdraw instruction
+   */
+  withdraw: (
+    owner: Keypair | PublicKey,
+    poolTokenAmountIn: Decimal | OrcaU64,
+    minTokenAOut: Decimal | OrcaU64,
+    minTokenBOut: Decimal | OrcaU64
+  ) => Promise<TransactionPayload>;
 };
 
 /**
@@ -84,6 +126,14 @@ export type OrcaToken = {
   scale: number;
 };
 
+/**
+ * An Orca Token within an OrcaPool
+ * @param addr The public key for this token for this Orca Pool
+ */
+export type OrcaPoolToken = OrcaToken & {
+  addr: PublicKey;
+};
+
 export type Quote = {
   /**
    * Returns the rate of exchange given the trade amount. Fees are included.

src/public/utils/numbers/u64-utils.ts

@@ -1,6 +1,7 @@
 import { u64 } from "@solana/spl-token";
 import Decimal from "decimal.js";
 import { OrcaToken, OrcaU64 } from "../..";
+import { OrcaPoolParams } from "../../../model/orca/pool/pool-types";
 import { DecimalUtil } from "./decimal-utils";
 
 export const ZERO = new u64(0);
@@ -20,6 +21,19 @@ export class U64Utils {
     return DecimalUtil.toU64(input, token.scale);
   }
 
+  public static toPoolU64(input: Decimal | OrcaU64, pool: OrcaPoolParams, varName: string) {
+    if (input instanceof OrcaU64) {
+      if (input.scale !== pool.poolTokenDecimals) {
+        throw new Error(
+          `${varName}'s scale of ${input.scale} does not match pool's decimal of ${pool.poolTokenDecimals}`
+        );
+      }
+      return input.toU64();
+    }
+
+    return DecimalUtil.toU64(input, pool.poolTokenDecimals);
+  }
+
   // Note: divisor input variable modified in place
   // https://github.com/solana-labs/solana-program-library/blob/master/libraries/math/src/checked_ceil_div.rs#L5-L22
   public static ceilingDivision(dividend: u64, divisor: u64): [u64, u64] {

src/public/utils/web3/get-token-count.ts

@@ -1,6 +1,7 @@
 import { u64 } from "@solana/spl-token";
 import { Connection } from "@solana/web3.js";
-import { OrcaPoolParams, OrcaPoolToken } from "../../../model/orca/pool/pool-types";
+import { OrcaPoolToken } from "../..";
+import { OrcaPoolParams } from "../../../model/orca/pool/pool-types";
 import { deserializeAccount } from "./deserialize-account";
 
 export type PoolTokenCount = {