feat: add convenience methods to EpochSchedule (#17810)

* first try, failing test

* fix implementation and tests

* lint:fix

* move method tests to seperate test

* lint fix

* apply starry's comments and grab the bonus points

* minor fixes after starry's second review

Co-authored-by: Arrowana <8245419+Arrowana@users.noreply.github.com>

web3.js/src/connection.ts

@@ -27,6 +27,7 @@ import RpcClient from 'jayson/lib/client/browser';
 import {IWSRequestParams} from 'rpc-websockets/dist/lib/client';
 
 import {AgentManager} from './agent-manager';
+import {EpochSchedule} from './epoch-schedule';
 import {NonceAccount} from './nonce-account';
 import {PublicKey} from './publickey';
 import {Signer} from './keypair';
@@ -373,23 +374,6 @@ const GetEpochInfoResult = pick({
   transactionCount: optional(number()),
 });
 
-/**
- * Epoch schedule
- * (see https://docs.solana.com/terminology#epoch)
- */
-export type EpochSchedule = {
-  /** The maximum number of slots in each epoch */
-  slotsPerEpoch: number;
-  /** The number of slots before beginning of an epoch to calculate a leader schedule for that epoch */
-  leaderScheduleSlotOffset: number;
-  /** Indicates whether epochs start short and grow */
-  warmup: boolean;
-  /** The first epoch with `slotsPerEpoch` slots */
-  firstNormalEpoch: number;
-  /** The first slot of `firstNormalEpoch` */
-  firstNormalSlot: number;
-};
-
 const GetEpochScheduleResult = pick({
   slotsPerEpoch: number(),
   leaderScheduleSlotOffset: number(),
@@ -2788,7 +2772,14 @@ export class Connection {
     if ('error' in res) {
       throw new Error('failed to get epoch schedule: ' + res.error.message);
     }
-    return res.result;
+    const epochSchedule = res.result;
+    return new EpochSchedule(
+      epochSchedule.slotsPerEpoch,
+      epochSchedule.leaderScheduleSlotOffset,
+      epochSchedule.warmup,
+      epochSchedule.firstNormalEpoch,
+      epochSchedule.firstNormalSlot,
+    );
   }
 
   /**

web3.js/src/epoch-schedule.ts

@@ -0,0 +1,102 @@
+const MINIMUM_SLOT_PER_EPOCH = 32;
+
+// Returns the number of trailing zeros in the binary representation of self.
+function trailingZeros(n: number) {
+  let trailingZeros = 0;
+  while (n > 1) {
+    n /= 2;
+    trailingZeros++;
+  }
+  return trailingZeros;
+}
+
+// Returns the smallest power of two greater than or equal to n
+function nextPowerOfTwo(n: number) {
+  if (n === 0) return 1;
+  n--;
+  n |= n >> 1;
+  n |= n >> 2;
+  n |= n >> 4;
+  n |= n >> 8;
+  n |= n >> 16;
+  n |= n >> 32;
+  return n + 1;
+}
+
+/**
+ * Epoch schedule
+ * (see https://docs.solana.com/terminology#epoch)
+ * Can be retrieved with the {@link connection.getEpochSchedule} method
+ */
+export class EpochSchedule {
+  /** The maximum number of slots in each epoch */
+  public slotsPerEpoch: number;
+  /** The number of slots before beginning of an epoch to calculate a leader schedule for that epoch */
+  public leaderScheduleSlotOffset: number;
+  /** Indicates whether epochs start short and grow */
+  public warmup: boolean;
+  /** The first epoch with `slotsPerEpoch` slots */
+  public firstNormalEpoch: number;
+  /** The first slot of `firstNormalEpoch` */
+  public firstNormalSlot: number;
+
+  constructor(
+    slotsPerEpoch: number,
+    leaderScheduleSlotOffset: number,
+    warmup: boolean,
+    firstNormalEpoch: number,
+    firstNormalSlot: number,
+  ) {
+    this.slotsPerEpoch = slotsPerEpoch;
+    this.leaderScheduleSlotOffset = leaderScheduleSlotOffset;
+    this.warmup = warmup;
+    this.firstNormalEpoch = firstNormalEpoch;
+    this.firstNormalSlot = firstNormalSlot;
+  }
+
+  getEpoch(slot: number): number {
+    return this.getEpochAndSlotIndex(slot)[0];
+  }
+
+  getEpochAndSlotIndex(slot: number): [number, number] {
+    if (slot < this.firstNormalSlot) {
+      const epoch =
+        trailingZeros(nextPowerOfTwo(slot + MINIMUM_SLOT_PER_EPOCH + 1)) -
+        trailingZeros(MINIMUM_SLOT_PER_EPOCH) -
+        1;
+
+      const epochLen = this.getSlotsInEpoch(epoch);
+      const slotIndex = slot - (epochLen - MINIMUM_SLOT_PER_EPOCH);
+      return [epoch, slotIndex];
+    } else {
+      const normalSlotIndex = slot - this.firstNormalSlot;
+      const normalEpochIndex = Math.floor(normalSlotIndex / this.slotsPerEpoch);
+      const epoch = this.firstNormalEpoch + normalEpochIndex;
+      const slotIndex = normalSlotIndex % this.slotsPerEpoch;
+      return [epoch, slotIndex];
+    }
+  }
+
+  getFirstSlotInEpoch(epoch: number): number {
+    if (epoch <= this.firstNormalEpoch) {
+      return (Math.pow(2, epoch) - 1) * MINIMUM_SLOT_PER_EPOCH;
+    } else {
+      return (
+        (epoch - this.firstNormalEpoch) * this.slotsPerEpoch +
+        this.firstNormalSlot
+      );
+    }
+  }
+
+  getLastSlotInEpoch(epoch: number): number {
+    return this.getFirstSlotInEpoch(epoch) + this.getSlotsInEpoch(epoch) - 1;
+  }
+
+  getSlotsInEpoch(epoch: number) {
+    if (epoch < this.firstNormalEpoch) {
+      return Math.pow(2, epoch + trailingZeros(MINIMUM_SLOT_PER_EPOCH));
+    } else {
+      return this.slotsPerEpoch;
+    }
+  }
+}

web3.js/src/index.ts

@@ -3,6 +3,7 @@ export * from './blockhash';
 export * from './bpf-loader-deprecated';
 export * from './bpf-loader';
 export * from './connection';
+export * from './epoch-schedule';
 export * from './fee-calculator';
 export * from './keypair';
 export * from './loader';

web3.js/test/connection.test.ts

@@ -9,6 +9,7 @@ import {
   Account,
   Authorized,
   Connection,
+  EpochSchedule,
   SystemProgram,
   Transaction,
   LAMPORTS_PER_SOL,
@@ -24,7 +25,6 @@ import {
   BLOCKHASH_CACHE_TIMEOUT_MS,
   Commitment,
   EpochInfo,
-  EpochSchedule,
   InflationGovernor,
   SlotInfo,
 } from '../src/connection';

web3.js/test/epoch-schedule.test.ts

@@ -0,0 +1,46 @@
+import {expect} from 'chai';
+
+import {EpochSchedule} from '../src';
+
+describe('EpochSchedule', () => {
+  it('slot methods work', () => {
+    const firstNormalEpoch = 14;
+    const firstNormalSlot = 524_256;
+    const leaderScheduleSlotOffset = 432_000;
+    const slotsPerEpoch = 432_000;
+    const warmup = true;
+
+    const epochSchedule = new EpochSchedule(
+      slotsPerEpoch,
+      leaderScheduleSlotOffset,
+      warmup,
+      firstNormalEpoch,
+      firstNormalSlot,
+    );
+
+    expect(epochSchedule.getEpoch(35)).to.be.equal(1);
+    expect(epochSchedule.getEpochAndSlotIndex(35)).to.be.eql([1, 3]);
+
+    expect(
+      epochSchedule.getEpoch(firstNormalSlot + 3 * slotsPerEpoch + 12345),
+    ).to.be.equal(17);
+    expect(
+      epochSchedule.getEpochAndSlotIndex(
+        firstNormalSlot + 3 * slotsPerEpoch + 12345,
+      ),
+    ).to.be.eql([17, 12345]);
+
+    expect(epochSchedule.getSlotsInEpoch(4)).to.be.equal(512);
+    expect(epochSchedule.getSlotsInEpoch(100)).to.be.equal(slotsPerEpoch);
+
+    expect(epochSchedule.getFirstSlotInEpoch(2)).to.be.equal(96);
+    expect(epochSchedule.getLastSlotInEpoch(2)).to.be.equal(223);
+
+    expect(epochSchedule.getFirstSlotInEpoch(16)).to.be.equal(
+      firstNormalSlot + 2 * slotsPerEpoch,
+    );
+    expect(epochSchedule.getLastSlotInEpoch(16)).to.be.equal(
+      firstNormalSlot + 3 * slotsPerEpoch - 1,
+    );
+  });
+});