feat: support additional compute budget ixs (#25104)

solana-labs · May 20, 2022 · a7e6f4d · a7e6f4d
1 parent f6da78a
commit a7e6f4d
Show file tree

Hide file tree

Showing 2 changed files with 212 additions and 89 deletions.
diff --git a/web3.js/src/compute-budget.ts b/web3.js/src/compute-budget.ts
@@ -1,4 +1,5 @@
 import * as BufferLayout from '@solana/buffer-layout';
+import {u64} from '@solana/buffer-layout-utils';
 
 import {
   encodeData,
@@ -76,6 +77,34 @@ export class ComputeBudgetInstruction {
     return {bytes};
   }
 
+  /**
+   * Decode set compute unit limit compute budget instruction and retrieve the instruction params.
+   */
+  static decodeSetComputeUnitLimit(
+    instruction: TransactionInstruction,
+  ): SetComputeUnitLimitParams {
+    this.checkProgramId(instruction.programId);
+    const {units} = decodeData(
+      COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitLimit,
+      instruction.data,
+    );
+    return {units};
+  }
+
+  /**
+   * Decode set compute unit price compute budget instruction and retrieve the instruction params.
+   */
+  static decodeSetComputeUnitPrice(
+    instruction: TransactionInstruction,
+  ): SetComputeUnitPriceParams {
+    this.checkProgramId(instruction.programId);
+    const {microLamports} = decodeData(
+      COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitPrice,
+      instruction.data,
+    );
+    return {microLamports};
+  }
+
   /**
    * @internal
    */
@@ -96,11 +125,18 @@ export type ComputeBudgetInstructionType =
   // It would be preferable for this type to be `keyof ComputeBudgetInstructionInputData`
   // but Typedoc does not transpile `keyof` expressions.
   // See https://github.com/TypeStrong/typedoc/issues/1894
-  'RequestUnits' | 'RequestHeapFrame';
+  | 'RequestUnits'
+  | 'RequestHeapFrame'
+  | 'SetComputeUnitLimit'
+  | 'SetComputeUnitPrice';
 
 type ComputeBudgetInstructionInputData = {
   RequestUnits: IInstructionInputData & Readonly<RequestUnitsParams>;
   RequestHeapFrame: IInstructionInputData & Readonly<RequestHeapFrameParams>;
+  SetComputeUnitLimit: IInstructionInputData &
+    Readonly<SetComputeUnitLimitParams>;
+  SetComputeUnitPrice: IInstructionInputData &
+    Readonly<SetComputeUnitPriceParams>;
 };
 
 /**
@@ -109,8 +145,7 @@ type ComputeBudgetInstructionInputData = {
 export interface RequestUnitsParams {
   /** Units to request for transaction-wide compute */
   units: number;
-
-  /** Additional fee to pay */
+  /** Prioritization fee lamports */
   additionalFee: number;
 }
 
@@ -122,6 +157,22 @@ export type RequestHeapFrameParams = {
   bytes: number;
 };
 
+/**
+ * Set compute unit limit instruction params
+ */
+export interface SetComputeUnitLimitParams {
+  /** Transaction-wide compute unit limit */
+  units: number;
+}
+
+/**
+ * Set compute unit price instruction params
+ */
+export interface SetComputeUnitPriceParams {
+  /** Transaction compute unit price used for prioritization fees */
+  microLamports: number | bigint;
+}
+
 /**
  * An enumeration of valid ComputeBudget InstructionType's
  * @internal
@@ -147,6 +198,18 @@ export const COMPUTE_BUDGET_INSTRUCTION_LAYOUTS = Object.freeze<{
       ComputeBudgetInstructionInputData['RequestHeapFrame']
     >([BufferLayout.u8('instruction'), BufferLayout.u32('bytes')]),
   },
+  SetComputeUnitLimit: {
+    index: 2,
+    layout: BufferLayout.struct<
+      ComputeBudgetInstructionInputData['SetComputeUnitLimit']
+    >([BufferLayout.u8('instruction'), BufferLayout.u32('units')]),
+  },
+  SetComputeUnitPrice: {
+    index: 3,
+    layout: BufferLayout.struct<
+      ComputeBudgetInstructionInputData['SetComputeUnitPrice']
+    >([BufferLayout.u8('instruction'), u64('microLamports')]),
+  },
 });
 
 /**
@@ -186,4 +249,30 @@ export class ComputeBudgetProgram {
       data,
     });
   }
+
+  static setComputeUnitLimit(
+    params: SetComputeUnitLimitParams,
+  ): TransactionInstruction {
+    const type = COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitLimit;
+    const data = encodeData(type, params);
+    return new TransactionInstruction({
+      keys: [],
+      programId: this.programId,
+      data,
+    });
+  }
+
+  static setComputeUnitPrice(
+    params: SetComputeUnitPriceParams,
+  ): TransactionInstruction {
+    const type = COMPUTE_BUDGET_INSTRUCTION_LAYOUTS.SetComputeUnitPrice;
+    const data = encodeData(type, {
+      microLamports: BigInt(params.microLamports),
+    });
+    return new TransactionInstruction({
+      keys: [],
+      programId: this.programId,
+      data,
+    });
+  }
 }