@wireio/stake 0.1.0 → 0.1.2

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

@@ -1,481 +1,354 @@
-import { AnchorProvider, Program, BN } from '@coral-xyz/anchor';
+import { AnchorProvider, Program } from '@coral-xyz/anchor';
 import {
+    ParsedAccountData,
     PublicKey,
-    SystemProgram,
     SYSVAR_INSTRUCTIONS_PUBKEY,
     Transaction,
-    TransactionInstruction,
+    SystemProgram,
 } from '@solana/web3.js';
 import {
-    getAccount,
-    getAssociatedTokenAddressSync,
+    TOKEN_2022_PROGRAM_ID,
+    getAssociatedTokenAddress,
 } from '@solana/spl-token';
 
 import { SolanaProgramService } from '../program';
+import type { LiqsolCore } from '../../../assets/solana/types/liqsol_core';
 import {
-    TOKEN_2022_PROGRAM_ID,
-    ASSOCIATED_TOKEN_PROGRAM_ID,
-    STAKE_CONTROLLER_PROGRAM_ID,
-    YIELD_ORACLE_PROGRAM_ID,
-    deriveDistributionStatePDA,
-    deriveUserRecordPDA,
-    deriveStakeControllerStatePDA,
-    deriveBucketAuthorityPDA,
-    derivePayRateHistoryPDA,
-    deriveStakeControllerAuthorityPDA,
+    deriveBucketAuthorityPda,
+    deriveDistributionStatePda,
+    deriveLiqsolMintPda,
+    derivePayRateHistoryPda,
+    deriveUserRecordPda,
 } from '../constants';
-import { Distribution } from '../../../assets/solana/types/distribution';
-import { MismatchCandidate, CorrectAndRegisterBuild, CorrectionPlan } from '../types';
-
+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 {
-    private programs: SolanaProgramService;
+    private program: Program<LiqsolCore>;
 
     constructor(private provider: AnchorProvider) {
-        this.programs = new SolanaProgramService(provider);
+        const svc = new SolanaProgramService(provider);
+        this.program = svc.getProgram('liqsolCore');
     }
 
-    /** Anchor Program<Distribution> (address comes from SolanaProgramService) */
-    private get program(): Program<Distribution> {
-        return this.programs.getProgram('distribution') as Program<Distribution>;
+    get connection() {
+        return this.provider.connection;
     }
 
-    // ───────────────────────────────────────────────────────────────────────────────
-    // BASIC READS
-    // ───────────────────────────────────────────────────────────────────────────────
-
     /**
-     * Fetch Distribution global state.
+     * Fetch the global distribution state account.
      */
-    async getDistributionState(): Promise<any> {
-        const [pda] = deriveDistributionStatePDA();
-        return this.program.account.distributionState.fetch(pda);
+    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 UserRecord or null if it doesn’t exist yet.
+     * Fetch a user's distribution userRecord (or null if missing).
      */
-    async getUserRecord(user: PublicKey): Promise<any | null> {
-        const [pda] = deriveUserRecordPDA(user);
-        try { return await this.program.account.userRecord.fetch(pda); }
-        catch { return null; }
+    async getUserRecord(user: PublicKey): Promise<UserRecord | null> {
+        const pda = deriveUserRecordPda(user);
+        try {
+            return await this.program.account.userRecord.fetchNullable(pda);
+        } catch {
+            return null;
+        }
     }
 
     /**
-     * Read protocol-tracked balance for a user (from userRecord).
-     * Returns (amount, decimals).
-     * @default decimals=9
+     * Helper: get actual liqSOL balances for all token holders.
+     *
+     * Returns a map: owner (base58) -> actual balance (BigInt)
      */
-    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 };
-    }
+    private async getActualBalancesByOwner(): Promise<Map<string, bigint>> {
+        const liqsolMint = deriveLiqsolMintPda();
+        const mintStr = liqsolMint.toBase58();
 
-    /**
-     * 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,
+        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',
+            },
         );
-        try {
-            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 { amount: ZERO, decimals: 9, ata };
+
+        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;
+
+            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);
         }
-    }
 
-    // ───────────────────────────────────────────────────────────────────────────────
-    // STATE-CHANGING ACTIONS (single-user)
-    // ───────────────────────────────────────────────────────────────────────────────
+        return byOwner;
+    }
 
     /**
-     * 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()`.
+     * 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
      */
-    async updateUser(targetUser: PublicKey): Promise<string> {
-        const { builder } = await this.prepareUpdateUser(targetUser);
-        return builder.rpc();
-    }
+    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;
 
-    /** 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();
+            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 map;
     }
 
     /**
-     * Claim rewards for a user.
-     * @returns tx signature
+     * Discover all mismatch candidates where tracked > actual.
+     *
+     * - actual balances are derived from token accounts for the liqSOL mint
+     * - tracked balances come from Distribution.userRecord
      */
-    async claimRewards(user: PublicKey): Promise<string> {
-        const state = await this.getDistributionState();
-        const liqsolMint: PublicKey = state.liqsolMint;
-
-        const userAta = getAssociatedTokenAddressSync(
-            liqsolMint,
-            user,
-            false,
-            TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID
-        );
-
-        const [userRecordPDA] = deriveUserRecordPDA(user);
-        const [distributionStatePDA] = deriveDistributionStatePDA();
-        const [bucketAuthority] = deriveBucketAuthorityPDA();
-        const [payRateHistory] = derivePayRateHistoryPDA();
+    private async deriveMismatchCandidates(): Promise<MismatchCandidate[]> {
+        const [actualByOwner, trackedByOwner] = await Promise.all([
+            this.getActualBalancesByOwner(),
+            this.getTrackedBalances(),
+        ]);
 
-        const bucketTokenAccount = getAssociatedTokenAddressSync(
-            liqsolMint,
-            bucketAuthority,
-            true,
-            TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID
-        );
+        const out: MismatchCandidate[] = [];
 
-        const sig = await this.program.methods
-            .claimRewards()
-            .accounts({
-                user,
-                // @ts-ignore
-                userAta,
-                userRecord: userRecordPDA,
-                distributionState: distributionStatePDA,
-                liqsolMint,
-                stakeControllerProgram: STAKE_CONTROLLER_PROGRAM_ID,
-                bucketAuthority,
-                bucketTokenAccount,
-                payRateHistory,
-                tokenProgram: TOKEN_2022_PROGRAM_ID,
-                associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
-                systemProgram: SystemProgram.programId,
-            })
-            .rpc();
+        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 });
+            }
+        }
 
-        return sig;
+        // Largest discrepancy first
+        out.sort((a, b) => (b.delta > a.delta ? 1 : b.delta < a.delta ? -1 : 0));
+        return out;
     }
 
     /**
-     * Withdraw liqSOL (amount in base units; e.g. 1e9 = 1 liqSOL).
-     * @returns tx signature
+     * Canonical helper to build an `updateUser` instruction for a given user,
+     * matching the standalone update-user.ts script.
      */
-    async withdraw(user: PublicKey, amount: bigint | number): Promise<string> {
-        const state = await this.getDistributionState();
-        const liqsolMint: PublicKey = state.liqsolMint;
+    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 = getAssociatedTokenAddressSync(
-            liqsolMint,
-            user,
+        const userAta = await getAssociatedTokenAddress(
+            liqsolMintPDA,
+            targetUser,
             false,
             TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID
         );
 
-        const [userRecordPDA] = deriveUserRecordPDA(user);
-        const [distributionStatePDA] = deriveDistributionStatePDA();
-        const [stakeControllerState] = deriveStakeControllerStatePDA();
-        const [controllerAuthority] = deriveStakeControllerAuthorityPDA();
+        const bucketTokenAccount = await getAssociatedTokenAddress(
+            liqsolMintPDA,
+            bucketAuthorityPDA,
+            true, // allowOwnerOffCurve
+            TOKEN_2022_PROGRAM_ID,
+        );
 
-        const sig = await this.program.methods
-            .withdraw(new BN(amount.toString()))
+        return (this.program.methods as any)
+            .updateUser()
             .accounts({
-                user,
-                // @ts-ignore
+                user: targetUser,
                 userAta,
                 userRecord: userRecordPDA,
+                authority: walletPk,
+                payer: walletPk,
                 distributionState: distributionStatePDA,
-                liqsolMint,
+                liqsolMint: liqsolMintPDA,
+                instructionsSysvar: SYSVAR_INSTRUCTIONS_PUBKEY,
                 tokenProgram: TOKEN_2022_PROGRAM_ID,
-                stakeControllerProgram: STAKE_CONTROLLER_PROGRAM_ID,
-                stakeControllerState,
-                controllerAuthority,
-                yieldOracleProgram: YIELD_ORACLE_PROGRAM_ID,
+                bucketAuthority: bucketAuthorityPDA,
+                bucketTokenAccount,
+                payRateHistory: payRateHistoryPDA,
                 systemProgram: SystemProgram.programId,
             })
-            .rpc();
-
-        return sig;
-    }
-
-    // ───────────────────────────────────────────────────────────────────────────────
-    // 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');
-
-        const state = await this.getDistributionState();
-        const liqsolMint: PublicKey = state.liqsolMint;
-
-        const targetUserAta = getAssociatedTokenAddressSync(
-            liqsolMint,
-            targetUser,
-            false,
-            TOKEN_2022_PROGRAM_ID,
-            ASSOCIATED_TOKEN_PROGRAM_ID
-        );
-
-        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 };
+            .instruction();
     }
 
-    // ───────────────────────────────────────────────────────────────────────────────
-    // 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).
+     * 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).
      *
-     * It refuses to build if it can’t free enough to cover the full positive untracked amount
-     * (since updateUser has no partial-amount arg).
+     * 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?: {
-        /** optional override of computed mismatch; may be positive (register) or negative (self-correct) */
+    async buildCorrectRegisterTx(opts: {
         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;
+        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: [] },
+            };
         }
 
-        // CASE A: Negative mismatch → just correct self
-        if (desired < ZERO) {
-            console.log('Building Correct transaction for self with amount:', desired);
+        const availableBalance = BigInt(distState.availableBalance.toString());
 
-            const tx = new Transaction().add(await this.buildUpdateUserIx(self));
-            return {
-                canSucceed: true,
-                transaction: tx,
-                liqsolMint,
-                needToRegister: desired, // negative indicates self-correction
-                availableBefore,
-                candidates: [],
-                plan: { selected: [], willFree: ZERO, deficit: ZERO },
-                accounts: { selfUserRecordPda, selfUserAta },
+        const trackedEntry =
+            trackedByOwner.get(callerStr) ?? {
+                owner: walletPk,
+                tracked: BigInt(0),
             };
-        }
 
-        // Reuse the existing “positive mismatch” path
-        const needToRegister = desired; // >= 0
-        if (needToRegister === ZERO) {
+        const actual = actualByOwner.get(callerStr) ?? BigInt(0);
+        const tracked = trackedEntry.tracked;
+        const untracked = actual - tracked;
+
+        // Nothing to register
+        if (untracked <= BigInt(0)) {
             return {
-                canSucceed: false,
-                reason: 'No mismatch to resolve.',
-                liqsolMint,
-                needToRegister,
-                availableBefore,
-                candidates: [],
-                plan: { selected: [], willFree: ZERO, deficit: ZERO },
-                accounts: { selfUserRecordPda, selfUserAta },
+                needToRegister: false,
+                canSucceed: true,
+                reason: 'No untracked liqSOL to register',
+                transaction: undefined,
+                plan: { deficit: BigInt(0), willFree: BigInt(0), selected: [] },
             };
         }
 
-        // 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);
+        // 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,
-                liqsolMint,
-                needToRegister,
-                availableBefore,
-                candidates: [],
-                plan: { selected: [], willFree: ZERO, deficit: ZERO },
-                accounts: { selfUserRecordPda, selfUserAta },
+                plan: { deficit: BigInt(0), willFree: BigInt(0), selected: [] },
             };
         }
 
-        // Need to free more
-        const missing = needToRegister - availableBefore;
-        const candidates = opts?.preloadCandidates ?? (await this.fetchMismatchCandidates());
+        // Need to free up more availableBalance by correcting others.
+        const deficit = targetRegister - availableBalance;
 
-        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 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;
         }
 
-        const plan = this.chooseCandidatesFor(missing, candidates);
-        if (plan.deficit > ZERO) {
+        if (willFree < deficit) {
             return {
+                needToRegister: true,
                 canSucceed: false,
-                reason: `Insufficient freeable balance (need ${missing}, can free ${plan.willFree}).`,
-                liqsolMint,
-                needToRegister,
-                availableBefore,
-                candidates,
-                plan,
-                accounts: { selfUserRecordPda, selfUserAta },
+                reason: 'Not enough mismatched candidates to cover deficit',
+                transaction: undefined,
+                plan: { deficit, willFree, selected },
             };
         }
 
-        // Build: correct N others, then register self
         const tx = new Transaction();
-        for (const c of plan.selected) {
-            tx.add(await this.buildUpdateUserIx(c.owner));
+
+        // 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);
         }
-        tx.add(await this.buildUpdateUserIx(self));
 
         return {
+            needToRegister: true,
             canSucceed: true,
             transaction: tx,
-            liqsolMint,
-            needToRegister,
-            availableBefore,
-            candidates,
-            plan,
-            accounts: { selfUserRecordPda, selfUserAta },
+            plan: { deficit, willFree, selected },
         };
     }
-
-
-    // ───────────────────────────────────────────────────────────────────────────────
-    // 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;
-    }
-
-    /**
-     * 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);
+}