@wireio/stake 0.0.6 → 0.1.1

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

@@ -1,230 +1,354 @@
-// src/solana/clients/DistributionClient.ts
-import { AnchorProvider, Program, BN, web3 } from '@coral-xyz/anchor';
+import { AnchorProvider, Program } from '@coral-xyz/anchor';
 import {
+    ParsedAccountData,
     PublicKey,
+    SYSVAR_INSTRUCTIONS_PUBKEY,
     Transaction,
-    TransactionInstruction,
     SystemProgram,
-    VersionedTransaction,
 } from '@solana/web3.js';
 import {
     TOKEN_2022_PROGRAM_ID,
-    ASSOCIATED_TOKEN_PROGRAM_ID,
-    getAssociatedTokenAddressSync,
+    getAssociatedTokenAddress,
 } from '@solana/spl-token';
+
+import { SolanaProgramService } from '../program';
+import type { LiqsolCore } from '../../../assets/solana/types/liqsol_core';
 import {
-    DISTRIBUTION_PROGRAM_ID,
-    STAKE_CONTROLLER_PROGRAM_ID,
-    DistributionIDL,
-    Distribution,
-    YIELD_ORACLE_PROGRAM_ID,
-    LIQSOL_TOKEN_PROGRAM_ID,
+    deriveBucketAuthorityPda,
+    deriveDistributionStatePda,
+    deriveLiqsolMintPda,
+    derivePayRateHistoryPda,
+    deriveUserRecordPda,
 } from '../constants';
-import {
-    deriveDistributionStatePDA,
-    deriveUserRecordPDA,
-    deriveLiqsolMintAuthorityPDA,
-} from '../utils';
+import type { CorrectRegisterBuildResult, DistributionState, MismatchCandidate, ParsedAccountInfo, UserRecord } from '../types';
 
+/**
+ * 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
+ *
+ * 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
+ */
 export class DistributionClient {
-    constructor(private provider: AnchorProvider) { }
-
-    /** Wrapped Anchor Program for Distribution */
-    private get program(): Program<Distribution> {
-        const idlWithAddress = {
-            ...JSON.parse(JSON.stringify(DistributionIDL)),
-            address: DISTRIBUTION_PROGRAM_ID.toString(),
-        };
-        return new Program(idlWithAddress as any, this.provider) as Program<Distribution>;
-    }
+    private program: Program<LiqsolCore>;
 
-    /** Derive the PDA for global distribution state */
-    deriveStatePDA(): PublicKey {
-        return deriveDistributionStatePDA()[0];
+    constructor(private provider: AnchorProvider) {
+        const svc = new SolanaProgramService(provider);
+        this.program = svc.getProgram('liqsolCore');
     }
 
-    /** Derive the PDA for a user’s record */
-    deriveUserRecordPDA(user: PublicKey): PublicKey {
-        return deriveUserRecordPDA(user)[0];
+    get connection() {
+        return this.provider.connection;
     }
 
-    /** Fetch on-chain distribution state */
-    async getState(): Promise<any> {
-        return this.program.account.distributionState.fetch(this.deriveStatePDA());
+    /**
+     * Fetch the global distribution state account.
+     */
+    async getDistributionState(): Promise<DistributionState | null> {
+        const pda = deriveDistributionStatePda();
+        try {
+            // IDL account name: "distributionState"
+            return await this.program.account.distributionState.fetch(pda);
+        } catch {
+            return null;
+        }
     }
 
-    /** Fetch a user’s record or return null if it doesn’t exist */
-    async getUserRecord(user: PublicKey): Promise<any | null> {
+    /**
+     * Fetch a user's distribution userRecord (or null if missing).
+     */
+    async getUserRecord(user: PublicKey): Promise<UserRecord | null> {
+        const pda = deriveUserRecordPda(user);
         try {
-            return await this.program.account.userRecord.fetch(
-                this.deriveUserRecordPDA(user)
-            );
+            return await this.program.account.userRecord.fetchNullable(pda);
         } catch {
             return null;
         }
     }
 
-    /** Build an `initialize` transaction for the distribution program */
-    async buildInitializeTransaction(user: PublicKey): Promise<Transaction> {
-        const [statePda] = deriveDistributionStatePDA();
-        const ix = await this.program.methods
-            .initialize()
-            .accounts({
-                authority: user,
-                distributionState: statePda,
-                systemProgram: SystemProgram.programId,
-                rent: web3.SYSVAR_RENT_PUBKEY,
-            } as any)
-            .instruction();
-
-        return new Transaction().add(ix);
-    }
+    /**
+     * Helper: get actual liqSOL balances for all token holders.
+     *
+     * Returns a map: owner (base58) -> actual balance (BigInt)
+     */
+    private async getActualBalancesByOwner(): Promise<Map<string, bigint>> {
+        const liqsolMint = deriveLiqsolMintPda();
+        const mintStr = liqsolMint.toBase58();
 
-    /** Build an `updateUser` transaction (register or refresh a user record) */
-    async buildUpdateUserTransaction(user: PublicKey): Promise<Transaction> {
-        const statePda = this.deriveStatePDA();
-        const userPda = this.deriveUserRecordPDA(user);
-
-        // Fetch state to get the liqSOL mint address
-        const state: any = await this.getState();
-        const liqsolMint: PublicKey = state.liqsolMint;
-        const userAta = getAssociatedTokenAddressSync(
-            liqsolMint,
-            user,
-            false,
+        const accounts = await this.connection.getParsedProgramAccounts(
             TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID
+            {
+                filters: [
+                    // SPL token layout: mint at offset 0
+                    { memcmp: { offset: 0, bytes: mintStr } },
+                ],
+                commitment: 'confirmed',
+            },
         );
 
-        const [stakeControllerStatePda] = PublicKey.findProgramAddressSync(
-            [Buffer.from('stake_controller')],
-            STAKE_CONTROLLER_PROGRAM_ID
-        );
+        const byOwner = new Map<string, bigint>();
 
-        const ix = await this.program.methods
-            .updateUser()
-            .accounts({
-                user,
-                userAta,
-                userRecord: userPda,
-                authority: user,
-                payer: user,
-                distributionState: statePda,
-                stakeControllerState: stakeControllerStatePda,
-                tokenProgram: TOKEN_2022_PROGRAM_ID,
-                yieldOracleProgram: YIELD_ORACLE_PROGRAM_ID,
-                systemProgram: SystemProgram.programId,
-            } as any)
-            .instruction();
+        for (const acct of accounts) {
+            const data = acct.account.data as ParsedAccountData;
+            const parsed = data.parsed;
+            if (!parsed || parsed.type !== 'account') continue;
+
+            const info: ParsedAccountInfo = parsed.info;
+            const ownerStr = info.owner;
+            const amountStr = info.tokenAmount.amount;
+            const amount = BigInt(amountStr);
+
+            const prev = byOwner.get(ownerStr) ?? BigInt(0);
+            byOwner.set(ownerStr, prev + amount);
+        }
 
-        return new Transaction().add(ix);
+        return byOwner;
     }
 
-    /** Build a `withdraw` transaction */
-    async buildWithdrawTransaction(user: PublicKey, amount: number): Promise<Transaction> {
-        const statePda = this.deriveStatePDA();
-        const userPda = this.deriveUserRecordPDA(user);
-
-        // Fetch state for the mint
-        const state: any = await this.getState();
-        const liqsolMint: PublicKey = state.liqsolMint;
-        const userAta = getAssociatedTokenAddressSync(
-            liqsolMint,
-            user,
-            false,
-            TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID
-        );
+    /**
+     * Helper: get tracked balances from all userRecord accounts,
+     * keyed by *actual wallet owner*, not the userRecord PDA.
+     *
+     * userRecord struct:
+     *  - userAta: pubkey
+     *  - trackedBalance: u64
+     *  - claimBalance: u64
+     *  - lastClaimTimestamp: i64
+     *  - bump: u8
+     */
+    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 }>();
 
-        const [stakeControllerStatePda] = PublicKey.findProgramAddressSync(
-            [Buffer.from('stake_controller')],
-            STAKE_CONTROLLER_PROGRAM_ID
-        );
-        const [controllerAuthPda] = PublicKey.findProgramAddressSync(
-            [Buffer.from('stake_authority')],
-            STAKE_CONTROLLER_PROGRAM_ID
-        );
+        for (const r of records) {
+            const ur = r.account as UserRecord;
+            const userAta = ur.userAta as PublicKey;
 
-        const ix = await this.program.methods
-            .withdraw(new BN(amount))
-            .accounts({
-                user,
-                userAta,
-                userRecord: userPda,
-                distributionState: statePda,
-                liqsolMint,
-                tokenProgram: TOKEN_2022_PROGRAM_ID,
-                stakeControllerProgram: STAKE_CONTROLLER_PROGRAM_ID,
-                stakeControllerState: stakeControllerStatePda,
-                controllerAuthority: controllerAuthPda,
-                yieldOracleProgram: YIELD_ORACLE_PROGRAM_ID,
-                systemProgram: SystemProgram.programId,
-            } as any)
-            .instruction();
+            // 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 });
+        }
 
-        return new Transaction().add(ix);
+        return map;
     }
 
-    /** Build a `claimRewards` transaction */
-    async buildClaimRewardsTransaction(user: PublicKey): Promise<Transaction> {
-        const statePda = this.deriveStatePDA();
-        const userPda = this.deriveUserRecordPDA(user);
-
-        // Fetch state for the mint
-        const state: any = await this.getState();
-        const liqsolMint: PublicKey = state.liqsolMint;
-        const userAta = getAssociatedTokenAddressSync(
-            liqsolMint,
-            user,
+    /**
+     * Discover all mismatch candidates where tracked > actual.
+     *
+     * - actual balances are derived from token accounts for the liqSOL mint
+     * - tracked balances come from Distribution.userRecord
+     */
+    private async deriveMismatchCandidates(): Promise<MismatchCandidate[]> {
+        const [actualByOwner, trackedByOwner] = await Promise.all([
+            this.getActualBalancesByOwner(),
+            this.getTrackedBalances(),
+        ]);
+
+        const out: MismatchCandidate[] = [];
+
+        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 });
+            }
+        }
+
+        // Largest discrepancy first
+        out.sort((a, b) => (b.delta > a.delta ? 1 : b.delta < a.delta ? -1 : 0));
+        return out;
+    }
+
+    /**
+     * Canonical helper to build an `updateUser` instruction for a given user,
+     * matching the standalone update-user.ts script.
+     */
+    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,
             TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID
         );
-        const [mintAuthPda] = deriveLiqsolMintAuthorityPDA();
 
-        const ix = await this.program.methods
-            .claimRewards()
+        const bucketTokenAccount = await getAssociatedTokenAddress(
+            liqsolMintPDA,
+            bucketAuthorityPDA,
+            true, // allowOwnerOffCurve
+            TOKEN_2022_PROGRAM_ID,
+        );
+
+        return (this.program.methods as any)
+            .updateUser()
             .accounts({
-                user,
+                user: targetUser,
                 userAta,
-                userRecord: userPda,
-                distributionState: statePda,
-                liqsolMint,
-                liqsolProgram: LIQSOL_TOKEN_PROGRAM_ID,
-                liqsolMintAuthority: mintAuthPda,
-                instructionsSysvar: web3.SYSVAR_INSTRUCTIONS_PUBKEY,
-                yieldOracleProgram: YIELD_ORACLE_PROGRAM_ID,
+                userRecord: userRecordPDA,
+                authority: walletPk,
+                payer: walletPk,
+                distributionState: distributionStatePDA,
+                liqsolMint: liqsolMintPDA,
+                instructionsSysvar: SYSVAR_INSTRUCTIONS_PUBKEY,
                 tokenProgram: TOKEN_2022_PROGRAM_ID,
-                associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
+                bucketAuthority: bucketAuthorityPDA,
+                bucketTokenAccount,
+                payRateHistory: payRateHistoryPDA,
                 systemProgram: SystemProgram.programId,
-            } as any)
+            })
             .instruction();
-
-        return new Transaction().add(ix);
     }
 
-    /** Send & confirm a single-IX transaction; returns the tx signature */
-    async sendTransaction(
-        tx: Transaction,
-        signers: import('@solana/web3.js').Signer[] = []
-    ): Promise<string> {
-        tx.feePayer = this.provider.wallet.publicKey;
-        const { blockhash } = await this.provider.connection.getLatestBlockhash();
-        tx.recentBlockhash = blockhash;
-        return this.provider.sendAndConfirm(tx, signers);
-    }
+    /**
+     * Build the "correct & register" transaction.
+     *
+     * - 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.
+     */
+    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;
 
-    /** Simulate a single-IX transaction; returns error (if any) + compute units used */
-    async simulateTransaction(
-        tx: Transaction
-    ): Promise<{ err: any; unitsConsumed: number }> {
-        tx.feePayer = this.provider.wallet.publicKey;
-        const { blockhash } = await this.provider.connection.getLatestBlockhash();
-        tx.recentBlockhash = blockhash;
-        const versioned = new VersionedTransaction(tx.compileMessage());
-        const sim = await this.provider.connection.simulateTransaction(versioned, {
-            sigVerify: false,
-        });
-        return { err: sim.value.err, unitsConsumed: sim.value.unitsConsumed! };
+        // 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: [] },
+            };
+        }
+
+        // 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: [] },
+            };
+        }
+
+        // Need to free up more availableBalance by correcting others.
+        const deficit = targetRegister - availableBalance;
+
+        const allCandidates = await this.deriveMismatchCandidates();
+        const maxCandidates = opts.maxCandidates ?? 10;
+
+        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;
+        }
+
+        if (willFree < deficit) {
+            return {
+                needToRegister: true,
+                canSucceed: false,
+                reason: 'Not enough mismatched candidates to cover deficit',
+                transaction: undefined,
+                plan: { deficit, willFree, selected },
+            };
+        }
+
+        const tx = new Transaction();
+
+        // 1) Correct selected mismatched users.
+        for (const c of selected) {
+            const ix = await this.buildUpdateUserIx(c.owner);
+            tx.add(ix);
+        }
+
+        // 2) Register caller (updateUser(caller)).
+        {
+            const ix = await this.buildUpdateUserIx(walletPk);
+            tx.add(ix);
+        }
+
+        return {
+            needToRegister: true,
+            canSucceed: true,
+            transaction: tx,
+            plan: { deficit, willFree, selected },
+        };
     }
 }