@wireio/stake 0.0.5 → 0.1.0

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

@@ -1,93 +1,134 @@
-// src/solana/clients/DistributionClient.ts
-import { AnchorProvider, Program, BN, web3 } from '@coral-xyz/anchor';
+import { AnchorProvider, Program, BN } from '@coral-xyz/anchor';
 import {
     PublicKey,
+    SystemProgram,
+    SYSVAR_INSTRUCTIONS_PUBKEY,
     Transaction,
     TransactionInstruction,
-    SystemProgram,
-    VersionedTransaction,
 } from '@solana/web3.js';
 import {
-    TOKEN_2022_PROGRAM_ID,
-    ASSOCIATED_TOKEN_PROGRAM_ID,
+    getAccount,
     getAssociatedTokenAddressSync,
 } from '@solana/spl-token';
+
+import { SolanaProgramService } from '../program';
 import {
-    DISTRIBUTION_PROGRAM_ID,
+    TOKEN_2022_PROGRAM_ID,
+    ASSOCIATED_TOKEN_PROGRAM_ID,
     STAKE_CONTROLLER_PROGRAM_ID,
-    DistributionIDL,
-    Distribution,
     YIELD_ORACLE_PROGRAM_ID,
-    LIQSOL_TOKEN_PROGRAM_ID,
-} from '../constants';
-import {
     deriveDistributionStatePDA,
     deriveUserRecordPDA,
-    deriveLiqsolMintAuthorityPDA,
-} from '../utils';
+    deriveStakeControllerStatePDA,
+    deriveBucketAuthorityPDA,
+    derivePayRateHistoryPDA,
+    deriveStakeControllerAuthorityPDA,
+} from '../constants';
+import { Distribution } from '../../../assets/solana/types/distribution';
+import { MismatchCandidate, CorrectAndRegisterBuild, CorrectionPlan } from '../types';
 
 export class DistributionClient {
-    constructor(private provider: AnchorProvider) { }
+    private programs: SolanaProgramService;
 
-    /** Wrapped Anchor Program for Distribution */
+    constructor(private provider: AnchorProvider) {
+        this.programs = new SolanaProgramService(provider);
+    }
+
+    /** Anchor Program<Distribution> (address comes from SolanaProgramService) */
     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>;
+        return this.programs.getProgram('distribution') as Program<Distribution>;
     }
 
-    /** Derive the PDA for global distribution state */
-    deriveStatePDA(): PublicKey {
-        return deriveDistributionStatePDA()[0];
+    // ───────────────────────────────────────────────────────────────────────────────
+    // BASIC READS
+    // ───────────────────────────────────────────────────────────────────────────────
+
+    /**
+     * Fetch Distribution global state.
+     */
+    async getDistributionState(): Promise<any> {
+        const [pda] = deriveDistributionStatePDA();
+        return this.program.account.distributionState.fetch(pda);
     }
 
-    /** Derive the PDA for a user’s record */
-    deriveUserRecordPDA(user: PublicKey): PublicKey {
-        return deriveUserRecordPDA(user)[0];
+    /**
+     * Fetch a user’s UserRecord or null if it doesn’t exist yet.
+     */
+    async getUserRecord(user: PublicKey): Promise<any | null> {
+        const [pda] = deriveUserRecordPDA(user);
+        try { return await this.program.account.userRecord.fetch(pda); }
+        catch { return null; }
     }
 
-    /** Fetch on-chain distribution state */
-    async getState(): Promise<any> {
-        return this.program.account.distributionState.fetch(this.deriveStatePDA());
+    /**
+     * Read protocol-tracked balance for a user (from userRecord).
+     * Returns (amount, decimals).
+     * @default decimals=9
+     */
+    async getTrackedBalance(user: PublicKey): Promise<{ amount: bigint; decimals: number }> {
+        const rec = await this.getUserRecord(user);
+        if (!rec) return { amount: ZERO, decimals: 9 };
+        const amountStr: string = rec.trackedBalance?.toString?.() ?? '0';
+        return { amount: BigInt(amountStr), decimals: 9 };
     }
 
-    /** Fetch a user’s record or return null if it doesn’t exist */
-    async getUserRecord(user: PublicKey): Promise<any | null> {
+    /**
+     * Read *actual* liqSOL token balance (ATA) for a user.
+     * If user has no ATA yet, returns 0.
+     * @returns { amount, decimals, ata }
+     */
+    async getActualBalance(user: PublicKey): Promise<{ amount: bigint; decimals: number; ata: PublicKey }> {
+        const state = await this.getDistributionState();
+        const liqsolMint: PublicKey = state.liqsolMint;
+        const ata = getAssociatedTokenAddressSync(
+            liqsolMint,
+            user,
+            false,
+            TOKEN_2022_PROGRAM_ID,
+            ASSOCIATED_TOKEN_PROGRAM_ID
+        );
         try {
-            return await this.program.account.userRecord.fetch(
-                this.deriveUserRecordPDA(user)
-            );
+            const acc = await getAccount(this.provider.connection, ata, 'confirmed', TOKEN_2022_PROGRAM_ID);
+            // spl-token returns bigint amount
+            return { amount: acc.amount as unknown as bigint, decimals: 9, ata };
         } catch {
-            return null;
+            return { amount: ZERO, decimals: 9, ata };
         }
     }
 
-    /** 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();
+    // ───────────────────────────────────────────────────────────────────────────────
+    // STATE-CHANGING ACTIONS (single-user)
+    // ───────────────────────────────────────────────────────────────────────────────
 
-        return new Transaction().add(ix);
+    /**
+     * Register / Update a user’s record.
+     * - If updating someone else, pass their pubkey as `targetUser` and let the connected wallet be `authority`.
+     * - Optionally pass an explicit `authorityUser` if you need to override (must be the connected wallet).
+     * Returns the tx signature.
+     */
+    /**
+     * Register / Update a user’s record with the connected wallet as authority.
+     * Keeps a single source of truth via `prepareUpdateUser()`.
+     */
+    async updateUser(targetUser: PublicKey): Promise<string> {
+        const { builder } = await this.prepareUpdateUser(targetUser);
+        return builder.rpc();
     }
 
-    /** 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);
+    /** Build-only variant (returns the IX) so you can batch. */
+    async buildUpdateUserIx(targetUser: PublicKey): Promise<TransactionInstruction> {
+        const { builder } = await this.prepareUpdateUser(targetUser);
+        return builder.instruction();
+    }
 
-        // Fetch state to get the liqSOL mint address
-        const state: any = await this.getState();
+    /**
+     * Claim rewards for a user.
+     * @returns tx signature
+     */
+    async claimRewards(user: PublicKey): Promise<string> {
+        const state = await this.getDistributionState();
         const liqsolMint: PublicKey = state.liqsolMint;
+
         const userAta = getAssociatedTokenAddressSync(
             liqsolMint,
             user,
@@ -96,38 +137,49 @@ export class DistributionClient {
             ASSOCIATED_TOKEN_PROGRAM_ID
         );
 
-        const [stakeControllerStatePda] = PublicKey.findProgramAddressSync(
-            [Buffer.from('stake_controller')],
-            STAKE_CONTROLLER_PROGRAM_ID
+        const [userRecordPDA] = deriveUserRecordPDA(user);
+        const [distributionStatePDA] = deriveDistributionStatePDA();
+        const [bucketAuthority] = deriveBucketAuthorityPDA();
+        const [payRateHistory] = derivePayRateHistoryPDA();
+
+        const bucketTokenAccount = getAssociatedTokenAddressSync(
+            liqsolMint,
+            bucketAuthority,
+            true,
+            TOKEN_2022_PROGRAM_ID,
+            ASSOCIATED_TOKEN_PROGRAM_ID
         );
 
-        const ix = await this.program.methods
-            .updateUser()
+        const sig = await this.program.methods
+            .claimRewards()
             .accounts({
                 user,
+                // @ts-ignore
                 userAta,
-                userRecord: userPda,
-                authority: user,
-                payer: user,
-                distributionState: statePda,
-                stakeControllerState: stakeControllerStatePda,
+                userRecord: userRecordPDA,
+                distributionState: distributionStatePDA,
+                liqsolMint,
+                stakeControllerProgram: STAKE_CONTROLLER_PROGRAM_ID,
+                bucketAuthority,
+                bucketTokenAccount,
+                payRateHistory,
                 tokenProgram: TOKEN_2022_PROGRAM_ID,
-                yieldOracleProgram: YIELD_ORACLE_PROGRAM_ID,
+                associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
                 systemProgram: SystemProgram.programId,
-            } as any)
-            .instruction();
+            })
+            .rpc();
 
-        return new Transaction().add(ix);
+        return sig;
     }
 
-    /** 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();
+    /**
+     * Withdraw liqSOL (amount in base units; e.g. 1e9 = 1 liqSOL).
+     * @returns tx signature
+     */
+    async withdraw(user: PublicKey, amount: bigint | number): Promise<string> {
+        const state = await this.getDistributionState();
         const liqsolMint: PublicKey = state.liqsolMint;
+
         const userAta = getAssociatedTokenAddressSync(
             liqsolMint,
             user,
@@ -136,95 +188,294 @@ export class DistributionClient {
             ASSOCIATED_TOKEN_PROGRAM_ID
         );
 
-        const [stakeControllerStatePda] = PublicKey.findProgramAddressSync(
-            [Buffer.from('stake_controller')],
-            STAKE_CONTROLLER_PROGRAM_ID
-        );
-        const [controllerAuthPda] = PublicKey.findProgramAddressSync(
-            [Buffer.from('stake_authority')],
-            STAKE_CONTROLLER_PROGRAM_ID
-        );
+        const [userRecordPDA] = deriveUserRecordPDA(user);
+        const [distributionStatePDA] = deriveDistributionStatePDA();
+        const [stakeControllerState] = deriveStakeControllerStatePDA();
+        const [controllerAuthority] = deriveStakeControllerAuthorityPDA();
 
-        const ix = await this.program.methods
-            .withdraw(new BN(amount))
+        const sig = await this.program.methods
+            .withdraw(new BN(amount.toString()))
             .accounts({
                 user,
+                // @ts-ignore
                 userAta,
-                userRecord: userPda,
-                distributionState: statePda,
+                userRecord: userRecordPDA,
+                distributionState: distributionStatePDA,
                 liqsolMint,
                 tokenProgram: TOKEN_2022_PROGRAM_ID,
                 stakeControllerProgram: STAKE_CONTROLLER_PROGRAM_ID,
-                stakeControllerState: stakeControllerStatePda,
-                controllerAuthority: controllerAuthPda,
+                stakeControllerState,
+                controllerAuthority,
                 yieldOracleProgram: YIELD_ORACLE_PROGRAM_ID,
                 systemProgram: SystemProgram.programId,
-            } as any)
-            .instruction();
+            })
+            .rpc();
 
-        return new Transaction().add(ix);
+        return sig;
     }
 
-    /** Build a `claimRewards` transaction */
-    async buildClaimRewardsTransaction(user: PublicKey): Promise<Transaction> {
-        const statePda = this.deriveStatePDA();
-        const userPda = this.deriveUserRecordPDA(user);
+    // ───────────────────────────────────────────────────────────────────────────────
+    // INSTRUCTION BUILDERS (for bundling into 1 tx)
+    // ───────────────────────────────────────────────────────────────────────────────
+
+    /** Single prep path used by both send & build. */
+    private async prepareUpdateUser(targetUser: PublicKey) {
+        const authority = this.provider.wallet?.publicKey;
+        if (!authority) throw new Error('Wallet not connected');
 
-        // Fetch state for the mint
-        const state: any = await this.getState();
+        const state = await this.getDistributionState();
         const liqsolMint: PublicKey = state.liqsolMint;
-        const userAta = getAssociatedTokenAddressSync(
+
+        const targetUserAta = getAssociatedTokenAddressSync(
             liqsolMint,
-            user,
+            targetUser,
             false,
             TOKEN_2022_PROGRAM_ID,
             ASSOCIATED_TOKEN_PROGRAM_ID
         );
-        const [mintAuthPda] = deriveLiqsolMintAuthorityPDA();
 
-        const ix = await this.program.methods
-            .claimRewards()
-            .accounts({
-                user,
-                userAta,
-                userRecord: userPda,
-                distributionState: statePda,
+        const [targetUserRecordPDA] = deriveUserRecordPDA(targetUser);
+        const [distributionStatePDA] = deriveDistributionStatePDA();
+        const [payRateHistory] = derivePayRateHistoryPDA();
+
+        const accounts = {
+            user: targetUser,
+            userAta: targetUserAta,
+            // @ts-ignore (account name casing differences)
+            userRecord: targetUserRecordPDA,
+            authority,
+            payer: authority,
+            distributionState: distributionStatePDA,
+            liqsolMint,
+            instructionsSysvar: SYSVAR_INSTRUCTIONS_PUBKEY,
+            tokenProgram: TOKEN_2022_PROGRAM_ID,
+            stakeControllerProgram: STAKE_CONTROLLER_PROGRAM_ID,
+            payRateHistory,
+            systemProgram: SystemProgram.programId,
+        };
+
+        const builder = this.program.methods.updateUser().accounts(accounts as any);
+        return { builder, accounts };
+    }
+
+    // ───────────────────────────────────────────────────────────────────────────────
+    // ONE-SHOT BUILDER: correct 0..N others, then register self
+    // ───────────────────────────────────────────────────────────────────────────────
+
+    /**
+     * Build a single transaction that:
+     *  1) If self mismatch < 0 → just correct self (single updateUser).
+     *  2) Else (mismatch > 0) → correct top candidates (if needed) to free “available” balance
+     *     and then register self (updateUser).
+     *
+     * It refuses to build if it can’t free enough to cover the full positive untracked amount
+     * (since updateUser has no partial-amount arg).
+     */
+    async buildCorrectRegisterTx(opts?: {
+        /** optional override of computed mismatch; may be positive (register) or negative (self-correct) */
+        amount?: bigint;
+        /** optionally pass pre-fetched candidates to avoid re-reading on UI flows */
+        preloadCandidates?: MismatchCandidate[];
+    }): Promise<CorrectAndRegisterBuild> {
+        const self = this.provider.wallet?.publicKey;
+        if (!self) throw new Error('Wallet not connected');
+
+        const state = await this.getDistributionState();
+        const liqsolMint: PublicKey = state.liqsolMint;
+        const availableBefore = BigInt(state.availableBalance.toString());
+
+        // self balances
+        const [selfUserRecordPda] = deriveUserRecordPDA(self);
+        const { amount: selfTracked } = await this.getTrackedBalance(self);
+        const { amount: selfActual, ata: selfUserAta } = await this.getActualBalance(self);
+
+        // True on-chain mismatch (can be negative or positive)
+        const computedMismatch = selfActual - selfTracked; // + = need to register, - = need to correct self
+
+        // Allow caller to suggest an amount (mainly for planning); clamp to real mismatch
+        let desired = opts?.amount ?? computedMismatch;
+        if (computedMismatch >= ZERO) {
+            // clamp down to computedMismatch
+            if (desired < ZERO) desired = ZERO;
+            if (desired > computedMismatch) desired = computedMismatch;
+        } else {
+            // negative side: clamp up to computedMismatch (remember computedMismatch is negative)
+            if (desired > ZERO) desired = ZERO;
+            if (desired < computedMismatch) desired = computedMismatch;
+        }
+
+        // CASE A: Negative mismatch → just correct self
+        if (desired < ZERO) {
+            console.log('Building Correct transaction for self with amount:', desired);
+
+            const tx = new Transaction().add(await this.buildUpdateUserIx(self));
+            return {
+                canSucceed: true,
+                transaction: tx,
                 liqsolMint,
-                liqsolProgram: LIQSOL_TOKEN_PROGRAM_ID,
-                liqsolMintAuthority: mintAuthPda,
-                instructionsSysvar: web3.SYSVAR_INSTRUCTIONS_PUBKEY,
-                yieldOracleProgram: YIELD_ORACLE_PROGRAM_ID,
-                tokenProgram: TOKEN_2022_PROGRAM_ID,
-                associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
-                systemProgram: SystemProgram.programId,
-            } as any)
-            .instruction();
+                needToRegister: desired, // negative indicates self-correction
+                availableBefore,
+                candidates: [],
+                plan: { selected: [], willFree: ZERO, deficit: ZERO },
+                accounts: { selfUserRecordPda, selfUserAta },
+            };
+        }
+
+        // Reuse the existing “positive mismatch” path
+        const needToRegister = desired; // >= 0
+        if (needToRegister === ZERO) {
+            return {
+                canSucceed: false,
+                reason: 'No mismatch to resolve.',
+                liqsolMint,
+                needToRegister,
+                availableBefore,
+                candidates: [],
+                plan: { selected: [], willFree: ZERO, deficit: ZERO },
+                accounts: { selfUserRecordPda, selfUserAta },
+            };
+        }
 
-        return new Transaction().add(ix);
+        // Already enough available → just register self
+        if (availableBefore >= needToRegister) {
+            console.log('Available balance', availableBefore, 'is enough to register self:', needToRegister);
+            const registerIx = await this.buildUpdateUserIx(self);
+            const tx = new Transaction().add(registerIx);
+            return {
+                canSucceed: true,
+                transaction: tx,
+                liqsolMint,
+                needToRegister,
+                availableBefore,
+                candidates: [],
+                plan: { selected: [], willFree: ZERO, deficit: ZERO },
+                accounts: { selfUserRecordPda, selfUserAta },
+            };
+        }
+
+        // Need to free more
+        const missing = needToRegister - availableBefore;
+        const candidates = opts?.preloadCandidates ?? (await this.fetchMismatchCandidates());
+
+        if (!candidates.length) {
+            return {
+                canSucceed: false,
+                reason: 'No candidates to free available balance.',
+                liqsolMint,
+                needToRegister,
+                availableBefore,
+                candidates: [],
+                plan: { selected: [], willFree: ZERO, deficit: missing },
+                accounts: { selfUserRecordPda, selfUserAta },
+            };
+        }
+
+        const plan = this.chooseCandidatesFor(missing, candidates);
+        if (plan.deficit > ZERO) {
+            return {
+                canSucceed: false,
+                reason: `Insufficient freeable balance (need ${missing}, can free ${plan.willFree}).`,
+                liqsolMint,
+                needToRegister,
+                availableBefore,
+                candidates,
+                plan,
+                accounts: { selfUserRecordPda, selfUserAta },
+            };
+        }
+
+        // Build: correct N others, then register self
+        const tx = new Transaction();
+        for (const c of plan.selected) {
+            tx.add(await this.buildUpdateUserIx(c.owner));
+        }
+        tx.add(await this.buildUpdateUserIx(self));
+
+        return {
+            canSucceed: true,
+            transaction: tx,
+            liqsolMint,
+            needToRegister,
+            availableBefore,
+            candidates,
+            plan,
+            accounts: { selfUserRecordPda, selfUserAta },
+        };
     }
 
-    /** 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);
+
+    // ───────────────────────────────────────────────────────────────────────────────
+    // MISMATCH DISCOVERY
+    // ───────────────────────────────────────────────────────────────────────────────
+
+    /**
+     * Fetch all distribution user records and turn them into mismatch candidates:
+     * rows where `tracked > actual` (delta > 0). Sorted by largest delta first.
+     *
+     * NOTE: This reads each user’s token account (`getAccount`) to recover
+     * the owner and actual balance — simpler and reliable for Token-2022 ATAs.
+     */
+    async fetchMismatchCandidates(): Promise<MismatchCandidate[]> {
+        const state = await this.getDistributionState();
+        const liqsolMint: PublicKey = state.liqsolMint;
+
+        const records = await this.program.account.userRecord.all();
+        if (!records.length) return [];
+
+        const out: MismatchCandidate[] = [];
+
+        // For each userRecord, read its token account to recover owner + actual
+        // (small N is fine; for very large N you can batch or shard).
+        for (const rec of records) {
+            const userRecordPda = rec.publicKey;
+            const userAta: PublicKey = rec.account.userAta;
+            const tracked = BigInt(rec.account.trackedBalance.toString());
+
+            let actual = ZERO;
+            let owner: PublicKey | null = null;
+
+            try {
+                const acc = await getAccount(this.provider.connection, userAta, 'confirmed', TOKEN_2022_PROGRAM_ID);
+                actual = acc.amount as unknown as bigint;
+                owner = acc.owner as unknown as PublicKey;
+            } catch {
+                // missing/closed ATA → actual = 0
+                owner = null;
+            }
+
+            if (!owner) continue;
+
+            const delta = tracked - actual;
+            if (delta > ZERO) {
+                // we only care about freeable deltas
+                out.push({ owner, userRecordPda, userAta, tracked, actual, delta });
+            }
+        }
+
+        // Largest first
+        out.sort((a, b) => (a.delta === b.delta ? 0 : a.delta > b.delta ? -1 : 1));
+        return out;
     }
 
-    /** 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! };
+    /**
+     * Given a required amount to free (in base units), choose the
+     * smallest prefix of candidates (already sorted desc) that can cover it.
+     */
+    chooseCandidatesFor(required: bigint, candidates: MismatchCandidate[]): CorrectionPlan {
+        if (required <= ZERO) return { selected: [], willFree: ZERO, deficit: ZERO };
+
+        let willFree = ZERO;
+        const selected: MismatchCandidate[] = [];
+
+        for (const c of candidates) {
+            if (willFree >= required) break;
+            selected.push(c);
+            willFree += c.delta;
+        }
+
+        const deficit = willFree >= required ? ZERO : required - willFree;
+        return { selected, willFree, deficit };
     }
-}
+}
+
+const ZERO = BigInt(0);