/
token-balances.ts
169 lines (154 loc) · 4.83 KB
/
token-balances.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import {
Connection,
ParsedMessageAccount,
PublicKey,
TransactionSignature,
} from '@solana/web3.js'
import BN from 'bn.js'
import table from 'text-table'
import { AddressLabels } from './address-labels'
import { resolveTokenRegistry } from './token-registry'
/**
* Interface to query token balances of a particular transaction.
*
* @category diagnostics
*/
export class TokenBalances {
private constructor(
private readonly connection: Connection,
private readonly signature: TransactionSignature,
private readonly addressLabels?: AddressLabels
) {}
/**
* Provides an interfact to query token balances for the transaction with the
* provided {@link signature}.
*
* If {@link addressLabels} are provided then they are used to resolve
* account and mint addresses to more meaningful labels.
*/
static forTransaction(
connection: Connection,
signature: TransactionSignature,
addressLabels?: AddressLabels
) {
return new TokenBalances(connection, signature, addressLabels)
}
/**
* Gets token balance for the provided account and mint.
*/
async balance(
account: PublicKey,
mint: PublicKey
): Promise<
{ amountPre: BN | number; amountPost: BN | number } | null | undefined
> {
const tokenBalances = await this.byAccountMap(true)
const forAccount = tokenBalances.get(account.toBase58())
if (forAccount == null) return null
return forAccount[mint.toBase58()]
}
/**
* Gets all token balances for the transaction mapped by account and then grouped
* by mint.
*/
async byAccountMap(
rawAddresses = false
): Promise<
Map<
string,
Record<
string,
{ amountPre: BN | number; amountPost: BN | number; rawMint: string }
>
>
> {
const parsed = await this.connection.getParsedTransaction(this.signature)
const accounts = parsed?.transaction.message.accountKeys
const preTokenBalances = parsed?.meta?.preTokenBalances
const postTokenBalances = parsed?.meta?.postTokenBalances
if (
(preTokenBalances == null && postTokenBalances == null) ||
accounts == null
) {
return new Map()
}
const byAccount = new Map()
for (let {
mint: rawMint,
uiTokenAmount,
accountIndex,
} of preTokenBalances ?? []) {
const account = this.resolveAccount(accounts, accountIndex, rawAddresses)
if (account == null) continue
const mint = rawAddresses
? rawMint
: this.addressLabels?.resolve(rawMint) ?? rawMint
byAccount.set(account, {
[mint]: { amountPre: new BN(uiTokenAmount.amount), rawMint },
})
}
for (let {
mint: rawMint,
uiTokenAmount,
accountIndex,
} of postTokenBalances ?? []) {
const account = this.resolveAccount(accounts, accountIndex, rawAddresses)
const mint = rawAddresses
? rawMint
: this.addressLabels?.resolve(rawMint) ?? rawMint
if (account == null) continue
if (!byAccount.has(account)) {
// The account has never been minted to before at all, thus has no pre balances
byAccount.set(account, {})
}
const current = byAccount.get(account)!
let currentMint = current[mint]
if (currentMint == null) {
// The account has not been minted the mint to before and thus it has no
// pre balance. We denote this as `0` pre balance which is what the
// solana explorer does as well.
currentMint = current[mint] = { amountPre: new BN(0), rawMint }
}
currentMint.amountPost = new BN(uiTokenAmount.amount)
}
return byAccount
}
/**
* Dumps all token balances to the console.
*/
async dump(
log: Console['log'] & { enabled?: boolean } = console.log
): Promise<TokenBalances> {
if (typeof log?.enabled !== 'undefined' && !log?.enabled) return this
const tokenRegistry = await resolveTokenRegistry()
const balances = await this.byAccountMap()
const rows: any[] = [
['Address', 'Token', 'Change', 'Post Balance'],
['-------', '-----', '------', '------------'],
]
for (const [account, mints] of balances) {
for (const [
mintAddress,
{ amountPre, amountPost, rawMint },
] of Object.entries(mints)) {
const delta = new BN(amountPost).sub(new BN(amountPre))
const unit = tokenRegistry.get(rawMint)?.name ?? 'tokens'
const row = [account, mintAddress, delta, `${amountPost} ${unit}`]
rows.push(row)
}
}
log(table(rows))
return this
}
private resolveAccount(
accounts: ParsedMessageAccount[],
accountIndex: number,
rawAddresses: boolean
) {
const parsedAccount = accounts[accountIndex]
return rawAddresses
? parsedAccount.pubkey.toBase58()
: this.addressLabels?.resolve(parsedAccount.pubkey) ??
parsedAccount.pubkey.toBase58()
}
}