spb/os-vault Public

1public extension SolanaAPIClient {2    func getAccountBalances(3        for address: String,4        withToken2022: Bool,5        tokensRepository: TokenRepository,6        commitment: String = "confirmed"7    ) async throws -> (8        resolved: [AccountBalance],9        unresolved: [UnknownAccountBalance]10    ) {11        // old token standard12        async let oldTokenAccountsResult = getAccountBalances(13            for: address,14            tokensRepository: tokensRepository,15            commitment: commitment,16            programId: TokenProgram.id.base58EncodedString,17            accountStateType: TokenAccountState.self,18            mintType: TokenMintState.self19        )2021        // token 202222        let token2022Accounts: (resolved: [AccountBalance], unresolved: [UnknownAccountBalance])23        if withToken2022 {24            async let tokens2022AccountsResult = getAccountBalances(25                for: address,26                tokensRepository: tokensRepository,27                commitment: commitment,28                programId: Token2022Program.id.base58EncodedString,29                accountStateType: Token2022AccountState.self,30                mintType: Token2022MintState.self31            )32            token2022Accounts = try await tokens2022AccountsResult33        } else {34            token2022Accounts = (resolved: [], unresolved: [])35        }3637        // get result3839        return try await(40            oldTokenAccountsResult.resolved + token2022Accounts.resolved,41            oldTokenAccountsResult.unresolved + token2022Accounts.unresolved42        )43    }4445    // MARK: - Helpers4647    private func getAccountBalances<48        T: TokenAccountLayoutState,49        M: MintLayoutState50    >(51        for address: String,52        tokensRepository: TokenRepository,53        commitment: String = "confirmed",54        programId: String,55        accountStateType _: T.Type,56        mintType _: M.Type57    ) async throws -> (58        resolved: [AccountBalance],59        unresolved: [UnknownAccountBalance]60    ) {61        let tokenAccounts = try await getTokenAccountsByOwner(62            pubkey: address,63            params: .init(64                mint: nil,65                programId: programId66            ),67            configs: .init(68                commitment: commitment,69                encoding: "base64"70            ),71            decodingTo: T.self72        )7374        return try await convertResult(75            tokenAccounts: tokenAccounts,76            tokensRepository: tokensRepository,77            tokenProgramId: programId,78            mintType: M.self79        )80    }8182    private func convertResult<83        T: TokenAccountLayoutState,84        M: MintLayoutState85    >(86        tokenAccounts: [TokenAccount<T>],87        tokensRepository: TokenRepository,88        commitment: String = "confirmed",89        tokenProgramId: String,90        mintType _: M.Type91    ) async throws -> (92        resolved: [AccountBalance],93        unresolved: [UnknownAccountBalance]94    ) {95        var resolvedAccountBalances: [AccountBalance] = []96        var unresolvedAccountBalances: [UnknownAccountBalance] = []9798        var unknownTokenAccountBalances: [TokenAccount<T>] = []99100        let tokenMetadatas = try await tokensRepository101            .get(addresses: tokenAccounts.map(\.account.data.mint.base58EncodedString))102103        for tokenAccount in tokenAccounts {104            let token: TokenMetadata105106            let tokenMintAddress: String = tokenAccount.account.data.mint.base58EncodedString107            if let resolvedToken = tokenMetadatas[tokenMintAddress] {108                token = resolvedToken109            } else {110                unknownTokenAccountBalances.append(tokenAccount)111                continue112            }113114            let accountBalance = AccountBalance(115                pubkey: tokenAccount.pubkey,116                lamports: tokenAccount.account.data.lamports,117                token: token,118                tokenProgramId: tokenProgramId,119                minimumBalanceForRentExemption: tokenAccount.account.lamports120            )121122            resolvedAccountBalances.append(accountBalance)123        }124125        // For case when token metadata in not in repository. We try to load on-chain126        let onChainTokenMetadataList = try await getMultipleMintDatas(127            mintAddresses: unknownTokenAccountBalances.map(128                \.account.data.mint.base58EncodedString129            ),130            commitment: commitment,131            mintType: M.self132        )133134        for tokenAccount in unknownTokenAccountBalances {135            let tokenMetadata = onChainTokenMetadataList[tokenAccount.account.data.mint.base58EncodedString]136137            if let tokenMetadata {138                // We have onchain token metadata139                let accountBalance = AccountBalance(140                    pubkey: tokenAccount.pubkey,141                    lamports: tokenAccount.account.data.lamports,142                    token: .unsupported(143                        tags: nil,144                        mint: tokenAccount.account.data.mint.base58EncodedString,145                        decimals: tokenMetadata.decimals,146                        symbol: "",147                        supply: tokenMetadata.supply148                    ),149                    tokenProgramId: tokenProgramId,150                    minimumBalanceForRentExemption: tokenAccount.account.lamports151                )152153                resolvedAccountBalances.append(accountBalance)154            } else {155                // We don't have onchain token metadata156                unresolvedAccountBalances.append(157                    .init(158                        pubkey: tokenAccount.pubkey,159                        lamports: tokenAccount.account.data.lamports,160                        mintAddress: tokenAccount.account.data.mint.base58EncodedString,161                        decimals: 0,162                        supply: nil,163                        tokenProgramId: tokenProgramId,164                        minimumBalanceForRentExemption: tokenAccount.account.lamports165                    )166                )167            }168        }169170        return (resolvedAccountBalances, unresolvedAccountBalances)171    }172}173