@wireio/stake 0.1.2 → 0.1.69

package/src/networks/solana/clients/distribution.client.ts

@@ -1,354 +1,358 @@
-import { AnchorProvider, Program } from '@coral-xyz/anchor';
-import {
-    ParsedAccountData,
-    PublicKey,
-    SYSVAR_INSTRUCTIONS_PUBKEY,
-    Transaction,
-    SystemProgram,
-} from '@solana/web3.js';
-import {
-    TOKEN_2022_PROGRAM_ID,
-    getAssociatedTokenAddress,
-} from '@solana/spl-token';
-
-import { SolanaProgramService } from '../program';
-import type { LiqsolCore } from '../../../assets/solana/types/liqsol_core';
-import {
-    deriveBucketAuthorityPda,
-    deriveDistributionStatePda,
-    deriveLiqsolMintPda,
-    derivePayRateHistoryPda,
-    deriveUserRecordPda,
-} from '../constants';
-import type { CorrectRegisterBuildResult, DistributionState, MismatchCandidate, ParsedAccountInfo, UserRecord } from '../types';
+import { AnchorProvider, Program, BN } from '@coral-xyz/anchor';
+import { PublicKey, SystemProgram, TransactionInstruction } from '@solana/web3.js';
+
+import { LiqsolCoreClientIdl, SolanaProgramService } from '../program';
+import type { DistributionState, DistributionUserRecord, GlobalConfig, PayRateEntry, PayRateHistory } from '../types';
+import { getAssociatedTokenAddressSync, TOKEN_2022_PROGRAM_ID } from '@solana/spl-token';
+import { ceilDiv } from '../utils';
+
+const INDEX_SCALE_BN = new BN('1000000000000');
 
 /**
  * Distribution client – wraps the distribution portion of the liqsol_core
- * program. Responsible for:
- *  - Reading DistributionState + UserRecord
- *  - Computing mismatch candidates (tracked > actual)
- *  - Building a "correct then register" transaction for the caller
+ * program in the *new* shares-only model.
+ *
+ * Responsibilities:
+ *  - Read DistributionState (global index/share pool)
+ *  - Read per-user UserRecord (shares)
+ *  - Provide simple helpers for querying user share positions
  *
- * Aligned with the on-chain `update_user` script:
- *   - Single `updateUser()` entrypoint that:
- *       * Can create userRecord if missing
- *       * Reconciles tracked vs actual using user, userAta, bucket, pay-rate
+ * There is no longer any notion of:
+ *  - tracked vs actual balances
+ *  - mismatch candidates
+ *  - updateUser / correct-&-register flow
+ *
+ * On-chain truth:
+ *  - User balances live in their liqSOL ATA.
+ *  - DistributionState + UserRecord track *shares* for yield.
  */
 export class DistributionClient {
-    private program: Program<LiqsolCore>;
+    private program: Program<LiqsolCoreClientIdl>;
 
-    constructor(private provider: AnchorProvider) {
-        const svc = new SolanaProgramService(provider);
-        this.program = svc.getProgram('liqsolCore');
+    constructor(
+        private readonly provider: AnchorProvider,
+        private readonly pgs: SolanaProgramService
+    ) {
+        this.program = pgs.getProgram('liqsolCore');
     }
 
     get connection() {
         return this.provider.connection;
     }
 
+    private async getTokenBalance(ata: PublicKey): Promise<BN> {
+        try {
+            const bal = await this.connection.getTokenAccountBalance(ata);
+            return new BN(bal.value.amount);
+        } catch {
+            return new BN(0);
+        }
+    }
+
     /**
      * Fetch the global distribution state account.
+     *
+     * IDL account name: "distributionState"
      */
     async getDistributionState(): Promise<DistributionState | null> {
-        const pda = deriveDistributionStatePda();
+        const pda = this.pgs.deriveDistributionStatePda();
         try {
-            // IDL account name: "distributionState"
             return await this.program.account.distributionState.fetch(pda);
         } catch {
             return null;
         }
     }
 
+
     /**
-     * Fetch a user's distribution userRecord (or null if missing).
+     * Fetch the global payRateHistory account (circular buffer of scaled pay rates).
+     *
+     * IDL account name: "payRateHistory"
      */
-    async getUserRecord(user: PublicKey): Promise<UserRecord | null> {
-        const pda = deriveUserRecordPda(user);
+    async getPayRateHistory(): Promise<PayRateHistory | null> {
+        const pda = this.pgs.derivePayRateHistoryPda();
         try {
-            return await this.program.account.userRecord.fetchNullable(pda);
+            // Anchor types map directly onto our PayRateHistory TS type
+            return (await this.program.account.payRateHistory.fetch(
+                pda,
+            )) as PayRateHistory;
         } catch {
             return null;
         }
     }
 
     /**
-     * Helper: get actual liqSOL balances for all token holders.
+     * Fetch the globalConfig account (contains depositFeeMultiplier, etc).
      *
-     * Returns a map: owner (base58) -> actual balance (BigInt)
+     * IDL account name: "globalConfig"
      */
-    private async getActualBalancesByOwner(): Promise<Map<string, bigint>> {
-        const liqsolMint = deriveLiqsolMintPda();
-        const mintStr = liqsolMint.toBase58();
-
-        const accounts = await this.connection.getParsedProgramAccounts(
-            TOKEN_2022_PROGRAM_ID,
-            {
-                filters: [
-                    // SPL token layout: mint at offset 0
-                    { memcmp: { offset: 0, bytes: mintStr } },
-                ],
-                commitment: 'confirmed',
-            },
-        );
-
-        const byOwner = new Map<string, bigint>();
-
-        for (const acct of accounts) {
-            const data = acct.account.data as ParsedAccountData;
-            const parsed = data.parsed;
-            if (!parsed || parsed.type !== 'account') continue;
+    async getGlobalConfig(): Promise<GlobalConfig | null> {
+        const pda = this.pgs.deriveGlobalConfigPda();
+        try {
+            return (await this.program.account.globalConfig.fetch(
+                pda,
+            )) as GlobalConfig;
+        } catch {
+            return null;
+        }
+    }
+    /**
+     * Fetch a user's distribution userRecord (or null if missing).
+     *
+     * IDL account name: "userRecord"
+     */
+    async getUserRecord(
+        ownerOrAta: PublicKey,
+    ): Promise<DistributionUserRecord | null> {
+        const liqsolMint = this.pgs.deriveLiqsolMintPda();
 
-            const info: ParsedAccountInfo = parsed.info;
-            const ownerStr = info.owner;
-            const amountStr = info.tokenAmount.amount;
-            const amount = BigInt(amountStr);
+        // 1) Try treating ownerOrAta as a WALLET: derive its liqSOL ATA, then userRecord.
+        try {
+            const ata = getAssociatedTokenAddressSync(
+                liqsolMint,
+                ownerOrAta,
+                true,
+                TOKEN_2022_PROGRAM_ID,
+            );
+            const pdaFromWallet = this.pgs.deriveUserRecordPda(ata);
+            const recFromWallet =
+                await this.program.account.userRecord.fetchNullable(pdaFromWallet);
+            if (recFromWallet) {
+                return recFromWallet;
+            }
+        } catch {
+            // fall through
+        }
 
-            const prev = byOwner.get(ownerStr) ?? BigInt(0);
-            byOwner.set(ownerStr, prev + amount);
+        // 2) Fallback: treat ownerOrAta as already being the TOKEN ACCOUNT.
+        try {
+            const pdaFromAta = this.pgs.deriveUserRecordPda(ownerOrAta);
+            return await this.program.account.userRecord.fetchNullable(pdaFromAta);
+        } catch {
+            return null;
         }
+    }
 
-        return byOwner;
+    /**
+     * Convenience: get the user's current distribution shares.
+     *
+     * Returns BN(0) if the user has no userRecord yet.
+     */
+    async getUserShares(user: PublicKey): Promise<BN> {
+        const record = await this.getUserRecord(user);
+        return record ? record.shares : new BN(0);
     }
 
     /**
-     * Helper: get tracked balances from all userRecord accounts,
-     * keyed by *actual wallet owner*, not the userRecord PDA.
+     * Convenience: get both the user's shares and totalShares, plus a
+     * floating-point ratio (userShares / totalShares).
      *
-     * userRecord struct:
-     *  - userAta: pubkey
-     *  - trackedBalance: u64
-     *  - claimBalance: u64
-     *  - lastClaimTimestamp: i64
-     *  - bump: u8
+     * NOTE: ratio is a JS number and may lose precision for very large values,
+     * but is fine for display / percentage UIs.
      */
-    private async getTrackedBalances(): Promise<
-        Map<string, { owner: PublicKey; tracked: bigint }>
-    > {
-        const records = await this.program.account.userRecord.all();
-        const map = new Map<string, { owner: PublicKey; tracked: bigint }>();
-
-        for (const r of records) {
-            const ur = r.account as UserRecord;
-            const userAta = ur.userAta as PublicKey;
-
-            // Resolve the *wallet* that owns this ATA
-            const ataInfo = await this.connection.getParsedAccountInfo(userAta);
-            const parsedData = ataInfo.value?.data as ParsedAccountData | undefined;
-            if (!parsedData) continue;
-
-            const parsed = parsedData.parsed;
-            if (!parsed || (parsed as any).type !== 'account') continue;
-
-            const info: ParsedAccountInfo = parsed.info;
-            const ownerStr = info.owner;
-            const owner = new PublicKey(ownerStr);
-
-            const tracked = BigInt(ur.trackedBalance.toString());
-            map.set(ownerStr, { owner, tracked });
+    async getUserSharePosition(user: PublicKey): Promise<{
+        shares: BN;
+        totalShares: BN;
+        ratio: number;
+    }> {
+        const [dist, userShares] = await Promise.all([
+            this.getDistributionState(),
+            this.getUserShares(user),
+        ]);
+
+        const totalShares = dist ? dist.totalShares : new BN(0);
+
+        let ratio = 0;
+        if (!totalShares.isZero()) {
+            // use toNumber() for UI-friendly ratio; if you need exact math,
+            // do it with BN or bigint at the call site instead.
+            ratio = userShares.toNumber() / totalShares.toNumber();
         }
 
-        return map;
+        return { shares: userShares, totalShares, ratio };
     }
 
     /**
-     * Discover all mismatch candidates where tracked > actual.
+     * Compute claimable liqSOL for a wallet:
+     * claimable = max(0, entitled - actual)
+     * entitled = floor(shares * syncedIndex / 1e12)
      *
-     * - actual balances are derived from token accounts for the liqSOL mint
-     * - tracked balances come from Distribution.userRecord
+     * `syncedIndex` mirrors on-chain `sync_index` behavior by incorporating any
+     * unsynced bucket delta (bucket_balance - last_bucket_balance).
      */
-    private async deriveMismatchCandidates(): Promise<MismatchCandidate[]> {
-        const [actualByOwner, trackedByOwner] = await Promise.all([
-            this.getActualBalancesByOwner(),
-            this.getTrackedBalances(),
+    async getClaimableLiqsol(user: PublicKey): Promise<BN> {
+        const liqsolMint = this.pgs.deriveLiqsolMintPda();
+        const bucketAuthority = this.pgs.deriveBucketAuthorityPda();
+
+        const userAta = getAssociatedTokenAddressSync(
+            liqsolMint,
+            user,
+            true,
+            TOKEN_2022_PROGRAM_ID,
+        );
+        const bucketTokenAccount = getAssociatedTokenAddressSync(
+            liqsolMint,
+            bucketAuthority,
+            true,
+            TOKEN_2022_PROGRAM_ID,
+        );
+
+        const [distributionState, userRecord, actualBalance, bucketBalance] = await Promise.all([
+            this.getDistributionState(),
+            this.getUserRecord(user),
+            this.getTokenBalance(userAta),
+            this.getTokenBalance(bucketTokenAccount),
         ]);
 
-        const out: MismatchCandidate[] = [];
+        if (!distributionState || !userRecord) {
+            return new BN(0);
+        }
+
+        let syncedIndex = new BN(distributionState.currentIndex.toString());
+        const totalShares = new BN(distributionState.totalShares.toString());
+        const lastBucketBalance = new BN(distributionState.lastBucketBalance.toString());
 
-        for (const [ownerStr, { owner, tracked }] of trackedByOwner.entries()) {
-            const actual = actualByOwner.get(ownerStr) ?? BigInt(0);
-            const delta = tracked - actual;
-            if (delta > BigInt(0)) {
-                out.push({ owner, actual, tracked, delta });
+        if (totalShares.gt(new BN(0)) && bucketBalance.gt(lastBucketBalance)) {
+            const delta = bucketBalance.sub(lastBucketBalance);
+            const indexDelta = delta.mul(INDEX_SCALE_BN).div(totalShares);
+            if (indexDelta.gt(new BN(0))) {
+                syncedIndex = syncedIndex.add(indexDelta);
             }
         }
 
-        // Largest discrepancy first
-        out.sort((a, b) => (b.delta > a.delta ? 1 : b.delta < a.delta ? -1 : 0));
-        return out;
+        const shares = new BN(userRecord.shares.toString());
+        const entitled = shares.mul(syncedIndex).div(INDEX_SCALE_BN);
+
+        if (entitled.lte(actualBalance)) {
+            return new BN(0);
+        }
+        return entitled.sub(actualBalance);
     }
 
     /**
-     * Canonical helper to build an `updateUser` instruction for a given user,
-     * matching the standalone update-user.ts script.
+     * Build claim_rewards instruction for a wallet.
      */
-    private async buildUpdateUserIx(targetUser: PublicKey) {
-        const walletPk = this.provider.wallet.publicKey;
-
-        const distributionStatePDA = deriveDistributionStatePda();
-        const userRecordPDA = deriveUserRecordPda(targetUser);
-        const liqsolMintPDA = deriveLiqsolMintPda();
-        const payRateHistoryPDA = derivePayRateHistoryPda();
-        const bucketAuthorityPDA = deriveBucketAuthorityPda();
-
-        const userAta = await getAssociatedTokenAddress(
-            liqsolMintPDA,
-            targetUser,
-            false,
+    async buildClaimRewardsIx(user: PublicKey): Promise<TransactionInstruction> {
+        const liqsolMint = this.pgs.deriveLiqsolMintPda();
+        const distributionState = this.pgs.deriveDistributionStatePda();
+        const bucketAuthority = this.pgs.deriveBucketAuthorityPda();
+
+        const userAta = getAssociatedTokenAddressSync(
+            liqsolMint,
+            user,
+            true,
             TOKEN_2022_PROGRAM_ID,
         );
-
-        const bucketTokenAccount = await getAssociatedTokenAddress(
-            liqsolMintPDA,
-            bucketAuthorityPDA,
-            true, // allowOwnerOffCurve
+        const bucketTokenAccount = getAssociatedTokenAddressSync(
+            liqsolMint,
+            bucketAuthority,
+            true,
             TOKEN_2022_PROGRAM_ID,
         );
 
-        return (this.program.methods as any)
-            .updateUser()
+        const userRecord = this.pgs.deriveUserRecordPda(userAta);
+        const bucketUserRecord = this.pgs.deriveUserRecordPda(bucketTokenAccount);
+        const extraAccountMetaList = this.pgs.deriveExtraAccountMetaListPda(liqsolMint);
+
+        return this.program.methods
+            .claimRewards()
             .accounts({
-                user: targetUser,
+                user,
                 userAta,
-                userRecord: userRecordPDA,
-                authority: walletPk,
-                payer: walletPk,
-                distributionState: distributionStatePDA,
-                liqsolMint: liqsolMintPDA,
-                instructionsSysvar: SYSVAR_INSTRUCTIONS_PUBKEY,
-                tokenProgram: TOKEN_2022_PROGRAM_ID,
-                bucketAuthority: bucketAuthorityPDA,
+                userRecord,
+                bucketUserRecord,
+                distributionState,
+                extraAccountMetaList,
+                liqsolCoreProgram: this.pgs.PROGRAM_IDS.LIQSOL_CORE,
+                transferHookProgram: this.pgs.PROGRAM_IDS.TRANSFER_HOOK,
+                liqsolMint,
+                bucketAuthority,
                 bucketTokenAccount,
-                payRateHistory: payRateHistoryPDA,
-                systemProgram: SystemProgram.programId,
+                tokenProgram: TOKEN_2022_PROGRAM_ID,
+                systemProgram: SystemProgram.programId
             })
             .instruction();
     }
 
     /**
-     * Build the "correct & register" transaction.
+     * Compute an average scaled pay rate over the most recent `windowSize`
+     * valid entries in the pay-rate history circular buffer.
      *
-     * - Fetches DistributionState + all userRecords + token holders
-     * - Computes caller's untracked amount (actual - tracked)
-     * - If DistributionState.availableBalance already covers that, we only
-     *   send updateUser(caller).
-     * - Otherwise we select top mismatch candidates until their freed deltas
-     *   cover the shortfall, then build updateUser(target) for each,
-     *   followed by updateUser(caller).
-     *
-     * NOTE:
-     *  - This no longer uses a separate updateSpecificUser; the single
-     *    updateUser entrypoint accepts any `user` as long as authority/payer
-     *    are valid, per your script.
+     * Returns a BN scaled by 1e12 (same as on-chain).
      */
-    async buildCorrectRegisterTx(opts: {
-        amount?: bigint;
-        maxCandidates?: number;
-    } = {}): Promise<CorrectRegisterBuildResult> {
-        const walletPk = this.provider.wallet.publicKey;
-        const callerStr = walletPk.toBase58();
-
-        const [distState, actualByOwner, trackedByOwner] = await Promise.all([
-            this.getDistributionState(),
-            this.getActualBalancesByOwner(),
-            this.getTrackedBalances(),
-        ]);
-
-        if (!distState) {
-            return {
-                needToRegister: false,
-                canSucceed: false,
-                reason: 'DistributionState not initialized',
-                transaction: undefined,
-                plan: { deficit: BigInt(0), willFree: BigInt(0), selected: [] },
-            };
-        }
-
-        const availableBalance = BigInt(distState.availableBalance.toString());
-
-        const trackedEntry =
-            trackedByOwner.get(callerStr) ?? {
-                owner: walletPk,
-                tracked: BigInt(0),
-            };
-
-        const actual = actualByOwner.get(callerStr) ?? BigInt(0);
-        const tracked = trackedEntry.tracked;
-        const untracked = actual - tracked;
-
-        // Nothing to register
-        if (untracked <= BigInt(0)) {
-            return {
-                needToRegister: false,
-                canSucceed: true,
-                reason: 'No untracked liqSOL to register',
-                transaction: undefined,
-                plan: { deficit: BigInt(0), willFree: BigInt(0), selected: [] },
-            };
+    async getAverageScaledPayRate(windowSize = 5): Promise<BN> {
+        const history = await this.getPayRateHistory();
+        if (!history) {
+            return new BN(0);
         }
 
-        // Optional user-specified cap on how much to register.
-        const targetRegister =
-            opts.amount && opts.amount > BigInt(0) ? opts.amount : untracked;
-
-        // Simple case: availableBalance already covers what we want to register
-        if (availableBalance >= targetRegister) {
-            const tx = new Transaction();
-            const ix = await this.buildUpdateUserIx(walletPk); // caller only
-            tx.add(ix);
-
-            return {
-                needToRegister: true,
-                canSucceed: true,
-                transaction: tx,
-                plan: { deficit: BigInt(0), willFree: BigInt(0), selected: [] },
-            };
+        const entries: PayRateEntry[] = history.entries ?? [];
+        if (!entries.length) {
+            return new BN(0);
         }
 
-        // Need to free up more availableBalance by correcting others.
-        const deficit = targetRegister - availableBalance;
+        const maxEntries: number =
+            typeof history.maxEntries === 'number'
+                ? history.maxEntries
+                : entries.length;
 
-        const allCandidates = await this.deriveMismatchCandidates();
-        const maxCandidates = opts.maxCandidates ?? 10;
+        const rawTotalAdded: any = history.totalEntriesAdded ?? 0;
+        const totalAddedBn = new BN(rawTotalAdded.toString());
 
-        const selected: MismatchCandidate[] = [];
-        let willFree = BigInt(0);
-
-        for (const c of allCandidates) {
-            if (c.owner.equals(walletPk)) continue; // don't self-correct here
-            selected.push(c);
-            willFree += c.delta;
-            if (willFree >= deficit || selected.length >= maxCandidates) break;
+        // No valid entries written yet
+        if (totalAddedBn.isZero()) {
+            return new BN(0);
         }
 
-        if (willFree < deficit) {
-            return {
-                needToRegister: true,
-                canSucceed: false,
-                reason: 'Not enough mismatched candidates to cover deficit',
-                transaction: undefined,
-                plan: { deficit, willFree, selected },
-            };
+        const totalAdded = Math.min(
+            totalAddedBn.toNumber(),
+            maxEntries,
+            entries.length,
+        );
+
+        if (!totalAdded) {
+            return new BN(0);
         }
 
-        const tx = new Transaction();
+        const COUNT = Math.max(1, Math.min(windowSize, totalAdded));
+
+        const currentIndexNum = Number(history.currentIndex ?? 0);
+
+        // Most recently written entry is at currentIndex - 1 (mod maxEntries)
+        let idx =
+            currentIndexNum === 0
+                ? maxEntries - 1
+                : currentIndexNum - 1;
+
+        let sum = new BN(0);
+        let valid = 0;
+        const zero = new BN(0);
+
+        for (let i = 0; i < COUNT; i++) {
+            const entry: any = entries[idx];
+            if (entry) {
+                // Support both camelCase and snake_case (for safety)
+                const rawScaled =
+                    entry.scaledRate ??
+                    entry.scaled_rate ??
+                    0;
+
+                const rate = new BN(rawScaled.toString());
+                if (rate.gt(zero)) {
+                    sum = sum.add(rate);
+                    valid += 1;
+                }
+            }
 
-        // 1) Correct selected mismatched users.
-        for (const c of selected) {
-            const ix = await this.buildUpdateUserIx(c.owner);
-            tx.add(ix);
+            // Walk backwards through the circular buffer
+            idx = idx === 0 ? maxEntries - 1 : idx - 1;
         }
 
-        // 2) Register caller (updateUser(caller)).
-        {
-            const ix = await this.buildUpdateUserIx(walletPk);
-            tx.add(ix);
+        if (!valid) {
+            return new BN(0);
         }
 
-        return {
-            needToRegister: true,
-            canSucceed: true,
-            transaction: tx,
-            plan: { deficit, willFree, selected },
-        };
+        // Same behavior as the dashboard: use a ceiling-like average
+        return ceilDiv(sum, new BN(valid));
     }
-}
+}