1use crate::core::events::*;
2use crate::core::invoke_context::{InvokeContext, InvokeLookup};
3use solana_sdk::pubkey::Pubkey;
4use std::collections::HashMap;
5use yellowstone_grpc_proto::prelude::{Transaction, TransactionStatusMeta};
6
7#[inline]
10pub(crate) fn pumpswap_is_pump_pool_from_fee_instruction(data: &[u8]) -> Option<bool> {
11 const GET_FEES: [u8; 8] = [231, 37, 126, 85, 207, 91, 63, 52];
12 const GET_FEES_WITH_QUOTE: [u8; 8] = [154, 237, 138, 92, 162, 2, 162, 187];
13 let discriminator: [u8; 8] = data.get(..8)?.try_into().ok()?;
14 if discriminator != GET_FEES && discriminator != GET_FEES_WITH_QUOTE {
15 return None;
16 }
17 match data.get(8) {
18 Some(0) => Some(false),
19 Some(1) => Some(true),
20 _ => None,
21 }
22}
23
24#[inline]
25fn set_pumpswap_is_pump_pool_from_fees_ix<L: InvokeLookup + ?Sized>(
26 meta: &TransactionStatusMeta,
27 transaction: &Option<Transaction>,
28 program_invokes: &L,
29 pool: Pubkey,
30 user: Pubkey,
31 buy: bool,
32 is_pump_pool: &mut bool,
33) {
34 let Some(swaps) = program_invokes.get_invokes(&crate::grpc::program_ids::PUMPSWAP_PROGRAM) else {
37 return;
38 };
39 let Some(swap) = crate::core::account_dispatcher::find_pumpswap_trade_invoke(
40 swaps, meta, transaction, pool, user, buy,
41 ) else {
42 return;
43 };
44 let Some(group) = meta.inner_instructions.iter().find(|g| g.index == swap.0 as u32) else {
45 return;
46 };
47 let (start, depth) = if swap.1 < 0 {
48 (0, 1)
49 } else {
50 let Some(depth) = group.instructions.get(swap.1 as usize).and_then(|ix| ix.stack_height) else {
51 return;
52 };
53 (swap.1 as usize + 1, depth)
54 };
55 let Some(fees) = program_invokes.get_invokes(&crate::grpc::program_ids::PUMPSWAP_FEES_PROGRAM) else {
56 return;
57 };
58 let Some(child_depth) = depth.checked_add(1) else {
59 return;
60 };
61 let mut flag = None;
62 for (inner_index, ix) in group.instructions.iter().enumerate().skip(start) {
63 let Some(height) = ix.stack_height else {
64 return;
66 };
67 if height <= depth {
68 break;
69 }
70 if height == child_depth {
71 if let Some(value) = pumpswap_is_pump_pool_from_fee_instruction(&ix.data) {
72 if !fees.contains(&(swap.0, inner_index as i32)) {
73 continue;
74 }
75 if flag.is_some_and(|previous| previous != value) {
76 return;
77 }
78 flag = Some(value);
79 }
80 }
81 }
82 if let Some(flag) = flag {
83 *is_pump_pool = flag;
84 }
85}
86
87#[inline]
88fn fill_data_with_lookup<L: InvokeLookup + ?Sized>(
89 event: &mut DexEvent,
90 meta: &TransactionStatusMeta,
91 transaction: &Option<Transaction>,
92 program_invokes: &L,
93) {
94 match event {
95 DexEvent::PumpSwapBuy(ref mut e) => {
96 set_pumpswap_is_pump_pool_from_fees_ix(
97 meta,
98 transaction,
99 program_invokes,
100 e.pool,
101 e.user,
102 true,
103 &mut e.is_pump_pool,
104 );
105 }
106 DexEvent::PumpSwapSell(ref mut e) => {
107 set_pumpswap_is_pump_pool_from_fees_ix(
108 meta,
109 transaction,
110 program_invokes,
111 e.pool,
112 e.user,
113 false,
114 &mut e.is_pump_pool,
115 );
116 }
117 _ => {}
118 }
119}
120
121pub fn fill_data(
122 event: &mut DexEvent,
123 meta: &TransactionStatusMeta,
124 transaction: &Option<Transaction>,
125 program_invokes: &HashMap<Pubkey, Vec<(i32, i32)>>,
126) {
127 fill_data_with_lookup(event, meta, transaction, program_invokes);
128}
129
130#[inline]
131pub(crate) fn fill_data_with_invoke_context(
132 event: &mut DexEvent,
133 meta: &TransactionStatusMeta,
134 transaction: &Option<Transaction>,
135 program_invokes: &InvokeContext,
136) {
137 fill_data_with_lookup(event, meta, transaction, program_invokes);
138}
139
140#[inline]
146pub fn fill_token_balances(
147 event: &mut DexEvent,
148 meta: &TransactionStatusMeta,
149 transaction: &Option<Transaction>,
150) {
151 let trade = match event {
152 DexEvent::PumpFunTrade(e)
153 | DexEvent::PumpFunBuy(e)
154 | DexEvent::PumpFunSell(e)
155 | DexEvent::PumpFunBuyExactSolIn(e) => e,
156 _ => return,
157 };
158
159 if let Some(user_index) = account_index(transaction, meta, &trade.user) {
160 trade.sol_balance = meta.post_balances.get(user_index).copied();
161 }
162
163 if let Some(balance) = post_token_balance(transaction, meta, &trade.associated_user) {
164 trade.token_balance = balance;
165 }
166}
167
168#[inline]
169fn account_index(
170 transaction: &Option<Transaction>,
171 meta: &TransactionStatusMeta,
172 account: &Pubkey,
173) -> Option<usize> {
174 if *account == Pubkey::default() {
175 return None;
176 }
177
178 let message = transaction.as_ref()?.message.as_ref()?;
179 message
180 .account_keys
181 .iter()
182 .chain(meta.loaded_writable_addresses.iter())
183 .chain(meta.loaded_readonly_addresses.iter())
184 .position(|key| key.as_slice() == account.as_ref())
185}
186
187#[inline]
188fn post_token_balance(
189 transaction: &Option<Transaction>,
190 meta: &TransactionStatusMeta,
191 account: &Pubkey,
192) -> Option<Option<u64>> {
193 if *account == Pubkey::default() {
194 return None;
195 }
196
197 let message = transaction.as_ref()?.message.as_ref()?;
198 let static_len = message.account_keys.len();
199 let writable_len = meta.loaded_writable_addresses.len();
200
201 let matches_account = |balance: &yellowstone_grpc_proto::prelude::TokenBalance| {
202 let index = balance.account_index as usize;
203 let key = if index < static_len {
204 message.account_keys.get(index)
205 } else if index < static_len + writable_len {
206 meta.loaded_writable_addresses.get(index - static_len)
207 } else {
208 meta.loaded_readonly_addresses.get(index - static_len - writable_len)
209 };
210
211 key.is_some_and(|key| key.as_slice() == account.as_ref())
212 };
213
214 if let Some(balance) = meta.post_token_balances.iter().find(|balance| matches_account(balance))
215 {
216 return Some(crate::grpc::transaction_meta::try_token_balance_raw_amount(balance));
217 }
218
219 meta.pre_token_balances.iter().any(matches_account).then_some(Some(0))
220}
221
222pub fn get_instruction_data<'a>(
223 meta: &'a TransactionStatusMeta,
224 transaction: &'a Option<Transaction>,
225 index: &(i32, i32), ) -> Option<&'a [u8]> {
227 let data = if index.1 >= 0 {
228 meta.inner_instructions
229 .iter()
230 .find(|i| i.index == index.0 as u32)?
231 .instructions
232 .get(index.1 as usize)?
233 .data
234 .as_slice()
235 } else {
236 transaction.as_ref()?.message.as_ref()?.instructions.get(index.0 as usize)?.data.as_slice()
237 };
238 Some(data)
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244 use yellowstone_grpc_proto::prelude::{Message, TokenBalance, UiTokenAmount};
245
246 fn token_balance(account_index: u32, amount: u64) -> TokenBalance {
247 TokenBalance {
248 account_index,
249 ui_token_amount: Some(UiTokenAmount {
250 amount: amount.to_string(),
251 ..Default::default()
252 }),
253 ..Default::default()
254 }
255 }
256
257 fn pumpfun_event(user: Pubkey, associated_user: Pubkey) -> DexEvent {
258 DexEvent::PumpFunBuy(PumpFunTradeEvent {
259 user,
260 associated_user,
261 is_buy: true,
262 ..Default::default()
263 })
264 }
265
266 fn transaction(accounts: &[Pubkey]) -> Option<Transaction> {
267 Some(Transaction {
268 message: Some(Message {
269 account_keys: accounts.iter().map(|key| key.to_bytes().to_vec()).collect(),
270 ..Default::default()
271 }),
272 ..Default::default()
273 })
274 }
275
276 #[test]
277 fn fills_pumpfun_user_post_transaction_balances() {
278 let user = Pubkey::new_unique();
279 let associated_user = Pubkey::new_unique();
280 let mut event = pumpfun_event(user, associated_user);
281 let meta = TransactionStatusMeta {
282 pre_balances: vec![50_000, 1_000],
283 post_balances: vec![40_000, 1_000],
284 pre_token_balances: vec![token_balance(1, 10)],
285 post_token_balances: vec![token_balance(1, 35)],
286 ..Default::default()
287 };
288
289 fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
290
291 let DexEvent::PumpFunBuy(trade) = event else { panic!("expected PumpFun buy") };
292 assert_eq!(trade.sol_balance, Some(40_000));
293 assert_eq!(trade.token_balance, Some(35));
294 }
295
296 #[test]
297 fn missing_token_balance_side_means_zero_for_created_or_closed_account() {
298 let user = Pubkey::new_unique();
299 let associated_user = Pubkey::new_unique();
300 let tx = transaction(&[user, associated_user]);
301
302 let mut created = pumpfun_event(user, associated_user);
303 let created_meta = TransactionStatusMeta {
304 pre_balances: vec![50, 0],
305 post_balances: vec![40, 1],
306 post_token_balances: vec![token_balance(1, 25)],
307 ..Default::default()
308 };
309 fill_token_balances(&mut created, &created_meta, &tx);
310 let DexEvent::PumpFunBuy(created) = created else { unreachable!() };
311 assert_eq!(created.token_balance, Some(25));
312
313 let mut closed = pumpfun_event(user, associated_user);
314 let closed_meta = TransactionStatusMeta {
315 pre_balances: vec![40, 1],
316 post_balances: vec![45, 0],
317 pre_token_balances: vec![token_balance(1, 25)],
318 ..Default::default()
319 };
320 fill_token_balances(&mut closed, &closed_meta, &tx);
321 let DexEvent::PumpFunBuy(closed) = closed else { unreachable!() };
322 assert_eq!(closed.token_balance, Some(0));
323 }
324
325 #[test]
326 fn ignores_non_pumpfun_events() {
327 let mut event = DexEvent::Error("unchanged".to_string());
328
329 fill_token_balances(&mut event, &TransactionStatusMeta::default(), &None);
330
331 assert!(matches!(event, DexEvent::Error(ref message) if message == "unchanged"));
332 }
333
334 #[test]
335 fn resolves_user_token_account_from_loaded_writable_addresses() {
336 let user = Pubkey::new_unique();
337 let associated_user = Pubkey::new_unique();
338 let mut event = pumpfun_event(user, associated_user);
339 let tx = transaction(&[user]);
340 let meta = TransactionStatusMeta {
341 pre_balances: vec![50, 1],
342 post_balances: vec![40, 1],
343 loaded_writable_addresses: vec![associated_user.to_bytes().to_vec()],
344 pre_token_balances: vec![token_balance(1, 10)],
345 post_token_balances: vec![token_balance(1, 35)],
346 ..Default::default()
347 };
348
349 fill_token_balances(&mut event, &meta, &tx);
350
351 let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
352 assert_eq!(trade.token_balance, Some(35));
353 }
354
355 #[test]
356 fn ignores_out_of_range_token_balance_account_index() {
357 let user = Pubkey::new_unique();
358 let associated_user = Pubkey::new_unique();
359 let mut event = pumpfun_event(user, associated_user);
360 let meta = TransactionStatusMeta {
361 pre_balances: vec![50, 1],
362 post_balances: vec![40, 1],
363 pre_token_balances: vec![token_balance(99, 10)],
364 post_token_balances: vec![token_balance(99, 35)],
365 ..Default::default()
366 };
367
368 fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
369
370 let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
371 assert_eq!(trade.token_balance, None);
372 }
373
374 #[test]
375 fn malformed_post_token_amount_is_unknown_not_zero() {
376 let user = Pubkey::new_unique();
377 let associated_user = Pubkey::new_unique();
378 let mut event = pumpfun_event(user, associated_user);
379 let meta = TransactionStatusMeta {
380 pre_balances: vec![50, 1],
381 post_balances: vec![40, 1],
382 post_token_balances: vec![TokenBalance {
383 account_index: 1,
384 ui_token_amount: Some(UiTokenAmount {
385 amount: "invalid".to_string(),
386 ..Default::default()
387 }),
388 ..Default::default()
389 }],
390 ..Default::default()
391 };
392
393 fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
394
395 let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
396 assert_eq!(trade.token_balance, None);
397 assert_eq!(trade.sol_balance, Some(40));
398 }
399}
400
401#[cfg(test)]
402mod fee_instruction_layout_tests {
403 use super::pumpswap_is_pump_pool_from_fee_instruction;
404
405 #[test]
406 fn pump_pool_flag_is_independent_of_market_cap_and_other_fee_instructions() {
407 for disc in [[231, 37, 126, 85, 207, 91, 63, 52], [154, 237, 138, 92, 162, 2, 162, 187]] {
408 for (flag, market_cap) in [(0u8, 1u128), (1, 0)] {
409 let mut data = disc.to_vec();
410 data.push(flag);
411 data.extend_from_slice(&market_cap.to_le_bytes());
412 assert_eq!(pumpswap_is_pump_pool_from_fee_instruction(&data), Some(flag != 0));
413 }
414 assert_eq!(pumpswap_is_pump_pool_from_fee_instruction(&disc), None);
415 }
416 let mut unknown = vec![0; 9];
417 unknown[8] = 1;
418 assert_eq!(pumpswap_is_pump_pool_from_fee_instruction(&unknown), None);
419 }
420}
421
422#[cfg(test)]
423mod review_fee_scope_regressions {
424 use super::*;
425 use yellowstone_grpc_proto::prelude::{CompiledInstruction, InnerInstruction, InnerInstructions, Message};
426
427 #[test]
428 fn fees_only_enrich_their_own_swap_and_unknown_parentage_preserves_logs() {
429 let pool0 = Pubkey::new_unique();
430 let pool1 = Pubkey::new_unique();
431 let user = Pubkey::new_unique();
432 let swap = |pool_index| CompiledInstruction {
433 accounts: std::iter::once(pool_index).chain(std::iter::repeat(2).take(22)).collect(),
434 data: crate::instr::pump_amm::discriminators::BUY.to_vec(),
435 ..Default::default()
436 };
437 let fee = |flag| {
438 let mut data = vec![231, 37, 126, 85, 207, 91, 63, 52];
439 data.push(flag);
440 data
441 };
442 let tx = Some(Transaction {
443 message: Some(Message {
444 account_keys: vec![pool0.to_bytes().to_vec(), pool1.to_bytes().to_vec(), user.to_bytes().to_vec()],
445 instructions: vec![swap(0), swap(1), CompiledInstruction { data: fee(1), ..Default::default() }],
446 ..Default::default()
447 }),
448 ..Default::default()
449 });
450 let mut meta = TransactionStatusMeta {
451 inner_instructions: vec![
452 InnerInstructions { index: 0, instructions: vec![InnerInstruction { data: fee(1), stack_height: Some(2), ..Default::default() }] },
453 InnerInstructions { index: 1, instructions: vec![InnerInstruction { data: fee(0), stack_height: Some(2), ..Default::default() }] },
454 ],
455 ..Default::default()
456 };
457 let invokes = HashMap::from([
458 (crate::grpc::program_ids::PUMPSWAP_PROGRAM, vec![(0, -1), (1, -1)]),
459 (crate::grpc::program_ids::PUMPSWAP_FEES_PROGRAM, vec![(0, 0), (1, 0), (2, -1)]),
460 ]);
461 let mut flag = false;
462 set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool0, user, true, &mut flag);
463 assert!(flag);
464 set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool1, user, true, &mut flag);
465 assert!(!flag, "last standalone query must not override the other pool");
466 meta.inner_instructions.clear();
467 set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool0, user, true, &mut flag);
468 assert!(!flag, "standalone query cannot be attributed to this pool");
469 meta.inner_instructions = vec![InnerInstructions { index: 1, instructions: vec![InnerInstruction { data: fee(0), stack_height: None, ..Default::default() }] }];
470 flag = true;
471 set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool1, user, true, &mut flag);
472 assert!(flag, "preserve a log-derived flag without parentage evidence");
473 }
474}