refactor(sdk-coin-sol): change marinade and jito flags to enum

Ticket: SC-2620

Author: haritkapadia

examples/ts/sol/stake-jito.ts

@@ -59,8 +59,8 @@ async function main() {
     .sender(account.publicKey.toBase58())
     .stakingAddress(JITO_STAKE_POOL_ADDRESS)
     .validator(JITO_STAKE_POOL_ADDRESS)
-    .isJito(true)
-    .jitoParams({
+    .stakingTypeParams({
+      type: 'JITO',
       stakePoolData: {
         managerFeeAccount: stakePoolAccount.account.data.managerFeeAccount.toString(),
         poolMint: stakePoolAccount.account.data.poolMint.toString(),

modules/sdk-coin-sol/src/lib/iface.ts

@@ -120,18 +120,32 @@ export interface Approve {
   };
 }
 
+export interface JitoStakingActivateParams {
+  type: 'JITO';
+  stakePoolData: DepositSolStakePoolData;
+}
+
+export interface MarinadeStakingActivateParams {
+  type: 'MARINADE';
+}
+
+export interface NativeStakingActivateParams {
+  type: 'NATIVE';
+}
+
+export type StakingActivateStakingTypeParams =
+  | JitoStakingActivateParams
+  | MarinadeStakingActivateParams
+  | NativeStakingActivateParams;
+
 export interface StakingActivate {
   type: InstructionBuilderTypes.StakingActivate;
   params: {
     fromAddress: string;
     stakingAddress: string;
     amount: string;
     validator: string;
-    isMarinade?: boolean;
-    isJito?: boolean;
-    jitoParams?: {
-      stakePoolData: DepositSolStakePoolData;
-    };
+    stakingTypeParams: StakingActivateStakingTypeParams;
   };
 }
 
@@ -140,20 +154,34 @@ export interface StakingDelegate {
   params: { stakingAddress: string; fromAddress: string; validator: string };
 }
 
+export interface JitoStakingDeactivateParams {
+  type: 'JITO';
+  stakePoolData: WithdrawStakeStakePoolData;
+  validatorAddress: string;
+  transferAuthorityAddress: string;
+}
+
+export interface MarinadeStakingDeactivateParams {
+  type: 'MARINADE';
+}
+
+export interface NativeStakingDeactivateParams {
+  type: 'NATIVE';
+}
+
+export type StakingDeactivateStakingTypeParams =
+  | JitoStakingDeactivateParams
+  | MarinadeStakingDeactivateParams
+  | NativeStakingDeactivateParams;
+
 export interface StakingDeactivate {
   type: InstructionBuilderTypes.StakingDeactivate;
   params: {
     fromAddress: string;
     stakingAddress: string;
     amount?: string;
     unstakingAddress?: string;
-    isMarinade?: boolean;
-    isJito?: boolean;
-    jitoParams?: {
-      stakePoolData: WithdrawStakeStakePoolData;
-      validatorAddress: string;
-      transferAuthorityAddress: string;
-    };
+    stakingTypeParams: StakingDeactivateStakingTypeParams;
     recipients?: Recipient[];
   };
 }

modules/sdk-coin-sol/src/lib/instructionParamsFactory.ts

@@ -421,19 +421,22 @@ function parseStakingActivateInstructions(
         stakingInstructions.initialize?.authorized.staker.toString() ||
         stakingInstructions.depositSol?.stakePool.toString() ||
         '',
-      isMarinade: stakingInstructionsIsMarinade(stakingInstructions),
-      isJito: stakingInstructionsIsJito(stakingInstructions),
-      ...(stakingInstructions.depositSol && stakingInstructionsIsJito(stakingInstructions)
+      stakingTypeParams: stakingInstructionsIsMarinade(stakingInstructions)
         ? {
-            jitoParams: {
-              stakePoolData: {
-                managerFeeAccount: stakingInstructions.depositSol.managerFeeAccount.toString(),
-                poolMint: stakingInstructions.depositSol.poolMint.toString(),
-                reserveStake: stakingInstructions.depositSol.reserveStake.toString(),
-              },
+            type: 'MARINADE',
+          }
+        : stakingInstructions.depositSol && stakingInstructionsIsJito(stakingInstructions)
+        ? {
+            type: 'JITO',
+            stakePoolData: {
+              managerFeeAccount: stakingInstructions.depositSol.managerFeeAccount.toString(),
+              poolMint: stakingInstructions.depositSol.poolMint.toString(),
+              reserveStake: stakingInstructions.depositSol.reserveStake.toString(),
             },
           }
-        : {}),
+        : {
+            type: 'NATIVE',
+          },
     },
   };
   instructionData.push(stakingActivate);
@@ -653,7 +656,24 @@ function parseStakingDeactivateInstructions(
         unstakingAddress:
           unstakingInstruction.split?.splitStakePubkey.toString() ||
           unstakingInstruction.withdrawStake?.destinationStake.toString(),
-        isMarinade: isMarinade,
+        stakingTypeParams: isMarinade
+          ? {
+              type: 'MARINADE',
+            }
+          : unstakingInstruction.withdrawStake !== undefined
+          ? {
+              type: 'JITO',
+              stakePoolData: {
+                managerFeeAccount: unstakingInstruction.withdrawStake.managerFeeAccount.toString(),
+                poolMint: unstakingInstruction.withdrawStake.poolMint.toString(),
+                validatorList: unstakingInstruction.withdrawStake.validatorList.toString(),
+              },
+              validatorAddress: unstakingInstruction.withdrawStake.validatorStake.toString(),
+              transferAuthorityAddress: unstakingInstruction.withdrawStake.sourceTransferAuthority.toString(),
+            }
+          : {
+              type: 'NATIVE',
+            },
         recipients: isMarinade
           ? [
               {
@@ -662,20 +682,6 @@ function parseStakingDeactivateInstructions(
               },
             ]
           : undefined,
-        ...(unstakingInstruction.withdrawStake !== undefined
-          ? {
-              isJito: unstakingInstruction.withdrawStake !== undefined,
-              jitoParams: unstakingInstruction.withdrawStake && {
-                stakePoolData: {
-                  managerFeeAccount: unstakingInstruction.withdrawStake.managerFeeAccount.toString(),
-                  poolMint: unstakingInstruction.withdrawStake.poolMint.toString(),
-                  validatorList: unstakingInstruction.withdrawStake.validatorList.toString(),
-                },
-                validatorAddress: unstakingInstruction.withdrawStake.validatorStake.toString(),
-                transferAuthorityAddress: unstakingInstruction.withdrawStake.sourceTransferAuthority.toString(),
-              },
-            }
-          : {}),
       },
     };
     instructionData.push(stakingDeactivate);

modules/sdk-coin-sol/src/lib/solInstructionFactory.ts

@@ -274,34 +274,30 @@ function createNonceAccountInstruction(data: WalletInit): TransactionInstruction
  */
 function stakingInitializeInstruction(data: StakingActivate): TransactionInstruction[] {
   const {
-    params: { fromAddress, stakingAddress, amount, validator, isMarinade, isJito, jitoParams },
+    params: { fromAddress, stakingAddress, amount, validator, stakingTypeParams },
   } = data;
   assert(fromAddress, 'Missing fromAddress param');
   assert(stakingAddress, 'Missing stakingAddress param');
   assert(amount, 'Missing amount param');
   assert(validator, 'Missing validator param');
-  assert(isMarinade !== undefined, 'Missing isMarinade param');
-  assert(isJito !== undefined, 'Missing isJito param');
-  assert([isMarinade, isJito].filter((x) => x).length <= 1, 'At most one of isMarinade and isJito can be true');
+  assert(stakingTypeParams !== undefined, 'Missing stakingTypeParams param');
 
   const fromPubkey = new PublicKey(fromAddress);
   const stakePubkey = new PublicKey(stakingAddress);
   const validatorPubkey = new PublicKey(validator);
   const tx = new Transaction();
 
-  if (isJito) {
-    assert(jitoParams, 'missing jitoParams param');
-
+  if (stakingTypeParams.type === 'JITO') {
     const instructions = depositSolInstructions(
       {
         stakePoolAddress: stakePubkey,
         from: fromPubkey,
         lamports: BigInt(amount),
       },
-      jitoParams.stakePoolData
+      stakingTypeParams.stakePoolData
     );
     tx.add(...instructions);
-  } else if (isMarinade) {
+  } else if (stakingTypeParams.type === 'MARINADE') {
     const walletInitStaking = StakeProgram.createAccount({
       fromPubkey,
       stakePubkey,
@@ -339,18 +335,16 @@ function stakingInitializeInstruction(data: StakingActivate): TransactionInstruc
  */
 function stakingDeactivateInstruction(data: StakingDeactivate): TransactionInstruction[] {
   const {
-    params: { fromAddress, stakingAddress, amount, unstakingAddress, isMarinade, isJito, recipients, jitoParams },
+    params: { fromAddress, stakingAddress, amount, unstakingAddress, recipients, stakingTypeParams },
   } = data;
   assert(fromAddress, 'Missing fromAddress param');
-  if (!isMarinade) {
+  if (stakingTypeParams.type !== 'MARINADE') {
     assert(stakingAddress, 'Missing stakingAddress param');
   }
-  assert([isMarinade, isJito].filter((x) => x).length <= 1, 'At most one of isMarinade and isJito can be true');
 
-  if (isJito) {
+  if (stakingTypeParams.type === 'JITO') {
     assert(unstakingAddress, 'Missing unstakingAddress param');
     assert(amount, 'Missing amount param');
-    assert(jitoParams, 'Missing jitoParams param');
 
     const tx = new Transaction();
     tx.add(
@@ -359,15 +353,15 @@ function stakingDeactivateInstruction(data: StakingDeactivate): TransactionInstr
           stakePoolAddress: new PublicKey(stakingAddress),
           tokenOwner: new PublicKey(fromAddress),
           destinationStakeAccount: new PublicKey(unstakingAddress),
-          validatorAddress: new PublicKey(jitoParams.validatorAddress),
-          transferAuthority: new PublicKey(jitoParams.transferAuthorityAddress),
+          validatorAddress: new PublicKey(stakingTypeParams.validatorAddress),
+          transferAuthority: new PublicKey(stakingTypeParams.transferAuthorityAddress),
           poolAmount: amount,
         },
-        jitoParams.stakePoolData
+        stakingTypeParams.stakePoolData
       )
     );
     return tx.instructions;
-  } else if (isMarinade) {
+  } else if (stakingTypeParams.type === 'MARINADE') {
     assert(recipients, 'Missing recipients param');
 
     const tx = new Transaction();

modules/sdk-coin-sol/src/lib/stakingActivateBuilder.ts

@@ -5,16 +5,14 @@ import { TransactionBuilder } from './transactionBuilder';
 import { InstructionBuilderTypes } from './constants';
 
 import assert from 'assert';
-import { StakingActivate } from './iface';
+import { StakingActivate, StakingActivateStakingTypeParams } from './iface';
 import { isValidStakingAmount, validateAddress } from './utils';
 
 export class StakingActivateBuilder extends TransactionBuilder {
   protected _amount: string;
   protected _stakingAddress: string;
   protected _validator: string;
-  protected _isMarinade = false;
-  protected _isJito = false;
-  protected _jitoParams?: StakingActivate['params']['jitoParams'];
+  protected _stakingTypeParams: StakingActivateStakingTypeParams = { type: 'NATIVE' };
 
   constructor(_coinConfig: Readonly<CoinConfig>) {
     super(_coinConfig);
@@ -34,9 +32,7 @@ export class StakingActivateBuilder extends TransactionBuilder {
         this.stakingAddress(activateInstruction.params.stakingAddress);
         this.amount(activateInstruction.params.amount);
         this.validator(activateInstruction.params.validator);
-        this.isMarinade(activateInstruction.params.isMarinade ?? false);
-        this.isJito(activateInstruction.params.isJito ?? false);
-        this.jitoParams(activateInstruction.params.jitoParams);
+        this.stakingTypeParams(activateInstruction.params.stakingTypeParams);
       }
     }
   }
@@ -84,34 +80,13 @@ export class StakingActivateBuilder extends TransactionBuilder {
   }
 
   /**
-   * Set isMarinade flag
-   * @param {boolean} flag - true if the transaction is for Marinade, false by default if not set
-   * @returns {StakingActivateBuilder} This staking builder
-   */
-  isMarinade(flag: boolean): this {
-    this._isMarinade = flag;
-    return this;
-  }
-
-  /**
-   * Set isJito flag
-   * @param {boolean} flag - true if the transaction is for Jito, false by default if not set
-   * @returns {StakingActivateBuilder} This staking builder
-   */
-  isJito(flag: boolean): this {
-    this._isJito = flag;
-    // this._coinConfig.network.type === NetworkType.TESTNET
-    return this;
-  }
-
-  /**
-   * Set parameters specific to Jito staking.
+   * Set parameters specific to a staking type.
    *
-   * @param {string} jitoParams parameters specific to Jito staking.
+   * @param {StakingActivateStakingTypeParams} stakingTypeParams parameters specific to a staking type.
    * @returns {StakingActivateBuilder} This staking builder.
    */
-  jitoParams(jitoParams: StakingActivate['params']['jitoParams']): this {
-    this._jitoParams = jitoParams;
+  stakingTypeParams(stakingTypeParams: StakingActivateStakingTypeParams): this {
+    this._stakingTypeParams = stakingTypeParams;
     return this;
   }
 
@@ -121,12 +96,7 @@ export class StakingActivateBuilder extends TransactionBuilder {
     assert(this._stakingAddress, 'Staking Address must be set before building the transaction');
     assert(this._validator, 'Validator must be set before building the transaction');
     assert(this._amount, 'Amount must be set before building the transaction');
-    assert(this._isMarinade !== undefined, 'isMarinade must be set before building the transaction');
-    assert(this._isJito !== undefined, 'isJito must be set before building the transaction');
-    assert(
-      [this._isMarinade, this._isJito].filter((x) => x).length <= 1,
-      'At most one of isMarinade and isJito can be true'
-    );
+    assert(this._stakingTypeParams, 'Staking type params must be set before building the transaction');
 
     if (this._sender === this._stakingAddress) {
       throw new BuildTransactionError('Sender address cannot be the same as the Staking address');
@@ -139,9 +109,7 @@ export class StakingActivateBuilder extends TransactionBuilder {
         stakingAddress: this._stakingAddress,
         amount: this._amount,
         validator: this._validator,
-        isMarinade: this._isMarinade,
-        isJito: this._isJito,
-        jitoParams: this._jitoParams,
+        stakingTypeParams: this._stakingTypeParams,
       },
     };
     this._instructionsData = [stakingAccountData];