stake-pool-js: Add deserializer for FutureEpoch (#6745)

* stake-pool-js: Add deserializer for `FutureEpoch`

* Fixup for new types

Author: joncinque

pnpm-lock.yaml

@@ -49,6 +49,7 @@
     "@solana/web3.js": "^1.91.8",
     "bn.js": "^5.2.0",
     "buffer": "^6.0.3",
+    "buffer-layout": "^1.2.2",
     "superstruct": "^1.0.4"
   },
   "devDependencies": {

stake-pool/js/package.json

@@ -139,43 +139,46 @@ export async function getStakeAccount(
 export async function getStakePoolAccounts(
   connection: Connection,
   stakePoolProgramAddress: PublicKey,
-): Promise<(StakePoolAccount | ValidatorListAccount)[] | undefined> {
+): Promise<(StakePoolAccount | ValidatorListAccount | undefined)[] | undefined> {
   const response = await connection.getProgramAccounts(stakePoolProgramAddress);
 
-  return response.map((a) => {
-    let decodedData;
-
-    if (a.account.data.readUInt8() === 1) {
+  return response
+    .map((a) => {
       try {
-        decodedData = StakePoolLayout.decode(a.account.data);
+        if (a.account.data.readUInt8() === 1) {
+          const data = StakePoolLayout.decode(a.account.data);
+          return {
+            pubkey: a.pubkey,
+            account: {
+              data,
+              executable: a.account.executable,
+              lamports: a.account.lamports,
+              owner: a.account.owner,
+            },
+          };
+        } else if (a.account.data.readUInt8() === 2) {
+          const data = ValidatorListLayout.decode(a.account.data);
+          return {
+            pubkey: a.pubkey,
+            account: {
+              data,
+              executable: a.account.executable,
+              lamports: a.account.lamports,
+              owner: a.account.owner,
+            },
+          };
+        } else {
+          console.error(
+            `Could not decode. StakePoolAccount Enum is ${a.account.data.readUInt8()}, expected 1 or 2!`,
+          );
+          return undefined;
+        }
       } catch (error) {
-        console.log('Could not decode StakeAccount. Error:', error);
-        decodedData = undefined;
+        console.error('Could not decode account. Error:', error);
+        return undefined;
       }
-    } else if (a.account.data.readUInt8() === 2) {
-      try {
-        decodedData = ValidatorListLayout.decode(a.account.data);
-      } catch (error) {
-        console.log('Could not decode ValidatorList. Error:', error);
-        decodedData = undefined;
-      }
-    } else {
-      console.error(
-        `Could not decode. StakePoolAccount Enum is ${a.account.data.readUInt8()}, expected 1 or 2!`,
-      );
-      decodedData = undefined;
-    }
-
-    return {
-      pubkey: a.pubkey,
-      account: {
-        data: decodedData,
-        executable: a.account.executable,
-        lamports: a.account.lamports,
-        owner: a.account.owner,
-      },
-    };
-  });
+    })
+    .filter((a) => a !== undefined);
 }
 
 /**

stake-pool/js/src/index.ts

@@ -1,4 +1,5 @@
-import { publicKey, struct, u32, u64, u8, option, vec } from '@coral-xyz/borsh';
+import { Layout, publicKey, struct, u32, u64, u8, option, vec } from '@coral-xyz/borsh';
+import { Layout as LayoutCls, u8 as u8Cls } from 'buffer-layout';
 import { PublicKey } from '@solana/web3.js';
 import BN from 'bn.js';
 import {
@@ -37,6 +38,50 @@ export const PublicKeyFromString = coerce(
   (value) => new PublicKey(value),
 );
 
+export class FutureEpochLayout<T> extends LayoutCls<T | null> {
+  layout: Layout<T>;
+  discriminator: Layout<number>;
+
+  constructor(layout: Layout<T>, property?: string) {
+    super(-1, property);
+    this.layout = layout;
+    this.discriminator = u8Cls();
+  }
+
+  encode(src: T | null, b: Buffer, offset = 0): number {
+    if (src === null || src === undefined) {
+      return this.discriminator.encode(0, b, offset);
+    }
+    // This isn't right, but we don't typically encode outside of tests
+    this.discriminator.encode(2, b, offset);
+    return this.layout.encode(src, b, offset + 1) + 1;
+  }
+
+  decode(b: Buffer, offset = 0): T | null {
+    const discriminator = this.discriminator.decode(b, offset);
+    if (discriminator === 0) {
+      return null;
+    } else if (discriminator === 1 || discriminator === 2) {
+      return this.layout.decode(b, offset + 1);
+    }
+    throw new Error('Invalid future epoch ' + this.property);
+  }
+
+  getSpan(b: Buffer, offset = 0): number {
+    const discriminator = this.discriminator.decode(b, offset);
+    if (discriminator === 0) {
+      return 1;
+    } else if (discriminator === 1 || discriminator === 2) {
+      return this.layout.getSpan(b, offset + 1) + 1;
+    }
+    throw new Error('Invalid future epoch ' + this.property);
+  }
+}
+
+export function futureEpoch<T>(layout: Layout<T>, property?: string): Layout<T | null> {
+  return new FutureEpochLayout<T>(layout, property);
+}
+
 export type StakeAccountType = Infer<typeof StakeAccountType>;
 export const StakeAccountType = enums(['uninitialized', 'initialized', 'delegated', 'rewardsPool']);
 
@@ -131,19 +176,19 @@ export const StakePoolLayout = struct<StakePool>([
   u64('lastUpdateEpoch'),
   struct([u64('unixTimestamp'), u64('epoch'), publicKey('custodian')], 'lockup'),
   struct(feeFields, 'epochFee'),
-  option(struct(feeFields), 'nextEpochFee'),
+  futureEpoch(struct(feeFields), 'nextEpochFee'),
   option(publicKey(), 'preferredDepositValidatorVoteAddress'),
   option(publicKey(), 'preferredWithdrawValidatorVoteAddress'),
   struct(feeFields, 'stakeDepositFee'),
   struct(feeFields, 'stakeWithdrawalFee'),
-  option(struct(feeFields), 'nextStakeWithdrawalFee'),
+  futureEpoch(struct(feeFields), 'nextStakeWithdrawalFee'),
   u8('stakeReferralFee'),
   option(publicKey(), 'solDepositAuthority'),
   struct(feeFields, 'solDepositFee'),
   u8('solReferralFee'),
   option(publicKey(), 'solWithdrawAuthority'),
   struct(feeFields, 'solWithdrawalFee'),
-  option(struct(feeFields), 'nextSolWithdrawalFee'),
+  futureEpoch(struct(feeFields), 'nextSolWithdrawalFee'),
   u64('lastEpochPoolTokenSupply'),
   u64('lastEpochTotalLamports'),
 ]);

stake-pool/js/src/layouts.ts

@@ -1,12 +1,22 @@
 declare module 'buffer-layout' {
-  export class Layout {}
-  export class UInt {}
-  /* eslint-disable  @typescript-eslint/no-unused-vars */
-  export function struct<T>(fields: any, property?: string, decodePrefixes?: boolean): any;
-  export function s32(property?: string): UInt;
-  export function u32(property?: string): UInt;
-  export function s16(property?: string): UInt;
-  export function u16(property?: string): UInt;
-  export function s8(property?: string): UInt;
-  export function u8(property?: string): UInt;
+  export class Layout<T = any> {
+    span: number;
+    property?: string;
+    constructor(span: number, property?: string);
+    decode(b: Buffer | undefined, offset?: number): T;
+    encode(src: T, b: Buffer, offset?: number): number;
+    getSpan(b: Buffer, offset?: number): number;
+    replicate(name: string): this;
+  }
+  export function struct<T>(
+    fields: Layout<any>[],
+    property?: string,
+    decodePrefixes?: boolean,
+  ): Layout<T>;
+  export function s32(property?: string): Layout<number>;
+  export function u32(property?: string): Layout<number>;
+  export function s16(property?: string): Layout<number>;
+  export function u16(property?: string): Layout<number>;
+  export function s8(property?: string): Layout<number>;
+  export function u8(property?: string): Layout<number>;
 }