1use std::collections::HashSet;
6use std::sync::Arc;
7
8use crate::{instr::read_pubkey_fast, DexEvent};
9use serde::{Deserialize, Serialize};
10use solana_sdk::pubkey::Pubkey;
11use solana_sdk::signature::Signature;
12use yellowstone_grpc_proto::prelude::{TokenBalance, Transaction, TransactionStatusMeta};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
19#[serde(rename_all = "lowercase")]
20pub enum YellowstoneMessageVersion {
21 Legacy,
22 V0,
23 V1,
24}
25
26#[inline]
30pub fn yellowstone_message_version(
31 message: &yellowstone_grpc_proto::prelude::Message,
32) -> YellowstoneMessageVersion {
33 if message.config.is_some() {
34 YellowstoneMessageVersion::V1
35 } else if message.versioned {
36 YellowstoneMessageVersion::V0
37 } else {
38 YellowstoneMessageVersion::Legacy
39 }
40}
41
42#[inline]
47pub(crate) fn fill_transaction_context(events: &mut [DexEvent], transaction: &Option<Transaction>) {
48 let Some(message) = transaction.as_ref().and_then(|tx| tx.message.as_ref()) else {
49 return;
50 };
51
52 let version = yellowstone_message_version(message);
53 for event in events.iter_mut() {
54 let alt_addresses = match event {
55 DexEvent::PumpFunCreate(create) => {
56 create.transaction_version = Some(version);
57 &mut create.alt_addresses
58 }
59 DexEvent::PumpFunCreateV2(create) => {
60 create.transaction_version = Some(version);
61 &mut create.alt_addresses
62 }
63 _ => continue,
64 };
65 *alt_addresses = creation_alt_addresses(message);
66 }
67 let blockhash = &message.recent_blockhash;
68 if blockhash.is_empty() {
69 return;
70 }
71 let mut metadata = events.iter_mut().filter_map(DexEvent::metadata_mut);
72 let Some(first) = metadata.next() else { return };
73 let encoded = bs58::encode(blockhash).into_string();
74 for item in metadata {
75 item.recent_blockhash = Some(encoded.clone());
76 }
77 first.recent_blockhash = Some(encoded);
78}
79
80fn creation_alt_addresses(
81 message: &yellowstone_grpc_proto::prelude::Message,
82) -> Option<Vec<Pubkey>> {
83 let lookups = &message.address_table_lookups;
84 if lookups.is_empty() {
85 return None;
86 }
87 let mut addresses = Vec::with_capacity(lookups.len());
88 for lookup in lookups {
89 addresses.push(Pubkey::try_from(lookup.account_key.as_slice()).ok()?);
91 }
92 Some(addresses)
93}
94
95#[inline]
97pub fn pubkey_bytes_to_bs58(bytes: &[u8]) -> Option<String> {
98 let a: [u8; 32] = bytes.try_into().ok()?;
99 Some(solana_sdk::pubkey::Pubkey::from(a).to_string())
100}
101
102pub fn collect_account_keys_bs58(
105 tx: &Transaction,
106 meta: &TransactionStatusMeta,
107) -> Option<Vec<String>> {
108 let msg = tx.message.as_ref()?;
109 let mut keys: Vec<String> =
110 msg.account_keys.iter().filter_map(|b| pubkey_bytes_to_bs58(b.as_slice())).collect();
111 for b in &meta.loaded_writable_addresses {
112 keys.push(pubkey_bytes_to_bs58(b)?);
113 }
114 for b in &meta.loaded_readonly_addresses {
115 keys.push(pubkey_bytes_to_bs58(b)?);
116 }
117 Some(keys)
118}
119
120#[inline]
122pub fn lamport_balance_deltas(meta: &TransactionStatusMeta) -> Vec<i128> {
123 meta.pre_balances
124 .iter()
125 .zip(meta.post_balances.iter())
126 .map(|(pre, post)| *post as i128 - *pre as i128)
127 .collect()
128}
129
130pub fn heuristic_sol_counterparties_for_watched_keys(
133 account_keys_bs58: &[String],
134 lamport_deltas: &[i128],
135 watched_bs58: &HashSet<&str>,
136 min_outflow_lamports: u64,
137) -> Vec<(String, String)> {
138 let min_l = min_outflow_lamports as i128;
139 let mut pairs = Vec::new();
140 for (i, key) in account_keys_bs58.iter().enumerate() {
141 if !watched_bs58.contains(key.as_str()) {
142 continue;
143 }
144 let d = lamport_deltas.get(i).copied().unwrap_or(0);
145 if d >= -min_l {
146 continue;
147 }
148 for (j, dj) in lamport_deltas.iter().enumerate() {
149 if i == j || *dj <= min_l / 2 {
150 continue;
151 }
152 pairs.push((key.clone(), account_keys_bs58[j].clone()));
153 }
154 }
155 pairs
156}
157
158pub fn collect_watch_transfer_counterparty_pairs(
162 tx: &Transaction,
163 meta: &TransactionStatusMeta,
164 watched_bs58: &[String],
165 min_native_outflow_lamports: u64,
166 spl_min_watch_decrease_raw: u64,
167) -> Option<Vec<(String, String)>> {
168 let keys = collect_account_keys_bs58(tx, meta)?;
169 let n = keys.len();
170 if meta.pre_balances.len() != n || meta.post_balances.len() != n {
171 return None;
172 }
173 let deltas = lamport_balance_deltas(meta);
174 let watched_h: HashSet<&str> = watched_bs58.iter().map(|s| s.as_str()).collect();
175
176 let mut pairs = heuristic_sol_counterparties_for_watched_keys(
177 &keys,
178 &deltas,
179 &watched_h,
180 min_native_outflow_lamports,
181 );
182 for w in watched_bs58 {
183 pairs.extend(spl_token_counterparty_by_owner(meta, w, spl_min_watch_decrease_raw));
184 }
185 pairs.sort_by(|a, b| a.1.cmp(&b.1));
186 pairs.dedup_by(|a, b| a.0 == b.0 && a.1 == b.1);
187 Some(pairs)
188}
189
190#[inline]
192pub fn token_balance_raw_amount(t: &TokenBalance) -> u64 {
193 try_token_balance_raw_amount(t).unwrap_or(0)
194}
195
196#[inline]
198pub fn try_token_balance_raw_amount(t: &TokenBalance) -> Option<u64> {
199 t.ui_token_amount.as_ref()?.amount.parse().ok()
200}
201
202pub fn spl_token_counterparty_by_owner(
207 meta: &TransactionStatusMeta,
208 watch_owner_bs58: &str,
209 min_watch_decrease_raw: u64,
210) -> Vec<(String, String)> {
211 use std::collections::{HashMap, HashSet};
212
213 let pre = meta.pre_token_balances.as_slice();
214 let post = meta.post_token_balances.as_slice();
215
216 let mut pre_m: HashMap<(String, String), u64> = HashMap::new();
217 for b in pre {
218 if b.owner.is_empty() {
219 continue;
220 }
221 let k = (b.mint.clone(), b.owner.clone());
222 *pre_m.entry(k).or_insert(0) += token_balance_raw_amount(b);
223 }
224 let mut post_m: HashMap<(String, String), u64> = HashMap::new();
225 for b in post {
226 if b.owner.is_empty() {
227 continue;
228 }
229 let k = (b.mint.clone(), b.owner.clone());
230 *post_m.entry(k).or_insert(0) += token_balance_raw_amount(b);
231 }
232
233 let mut mints = HashSet::new();
234 for (m, o) in pre_m.keys() {
235 if o == watch_owner_bs58 {
236 mints.insert(m.clone());
237 }
238 }
239 for (m, o) in post_m.keys() {
240 if o == watch_owner_bs58 {
241 mints.insert(m.clone());
242 }
243 }
244
245 let mut out = Vec::new();
246 let min_l = min_watch_decrease_raw;
247 for mint in mints {
248 let w_pre = pre_m.get(&(mint.clone(), watch_owner_bs58.to_string())).copied().unwrap_or(0);
249 let w_post =
250 post_m.get(&(mint.clone(), watch_owner_bs58.to_string())).copied().unwrap_or(0);
251 let lost = w_pre.saturating_sub(w_post);
252 if lost < min_l.max(1) {
253 continue;
254 }
255 for ((m, owner), po) in &post_m {
256 if m != &mint || owner == watch_owner_bs58 {
257 continue;
258 }
259 let pr = pre_m.get(&(mint.clone(), owner.clone())).copied().unwrap_or(0);
260 if *po > pr {
261 out.push((watch_owner_bs58.to_string(), owner.clone()));
262 }
263 }
264 }
265 out.sort_by(|a, b| a.1.cmp(&b.1));
266 out.dedup_by(|a, b| a.0 == b.0 && a.1 == b.1);
267 out
268}
269
270#[inline]
273pub fn yellowstone_static_account_keys_arc(tx: &Option<Transaction>) -> Arc<[Pubkey]> {
274 let Some(t) = tx.as_ref() else {
275 return Arc::from(Vec::<Pubkey>::new().into_boxed_slice());
276 };
277 let Some(msg) = t.message.as_ref() else {
278 return Arc::from(Vec::<Pubkey>::new().into_boxed_slice());
279 };
280 let keys: Vec<Pubkey> =
281 msg.account_keys.iter().map(|bytes| read_pubkey_fast(bytes.as_slice())).collect();
282 Arc::from(keys.into_boxed_slice())
283}
284
285#[inline]
287pub fn try_yellowstone_signature(sig: &[u8]) -> Option<Signature> {
288 if sig.len() != 64 {
289 return None;
290 }
291 let a: [u8; 64] = sig.try_into().ok()?;
292 Some(Signature::from(a))
293}
294
295#[cfg(test)]
296mod tests {
297 use super::*;
298 use yellowstone_grpc_proto::prelude::{Message, TransactionConfig};
299
300 #[test]
301 fn creation_alt_addresses_keep_order_and_reject_partial_provider_lists() {
302 let first = Pubkey::new_unique();
303 let second = Pubkey::new_unique();
304 for (keys, expected) in [
305 (vec![], None),
306 (
307 vec![
308 second.to_bytes().to_vec(),
309 first.to_bytes().to_vec(),
310 second.to_bytes().to_vec(),
311 ],
312 Some(vec![second, first, second]),
313 ),
314 (vec![first.to_bytes().to_vec(), vec![0; 31]], None),
315 ] {
316 let tx = Some(Transaction {
317 message: Some(Message {
318 address_table_lookups: keys
319 .into_iter()
320 .map(|account_key| {
321 yellowstone_grpc_proto::prelude::MessageAddressTableLookup {
322 account_key,
323 ..Default::default()
324 }
325 })
326 .collect(),
327 ..Default::default()
328 }),
329 ..Default::default()
330 });
331 let mut events = [
332 DexEvent::PumpFunCreate(Default::default()),
333 DexEvent::PumpFunCreateV2(Default::default()),
334 ];
335 fill_transaction_context(&mut events, &tx);
336 for event in events {
337 let actual = match event {
338 DexEvent::PumpFunCreate(c) => c.alt_addresses,
339 DexEvent::PumpFunCreateV2(c) => c.alt_addresses,
340 _ => unreachable!(),
341 };
342 assert_eq!(actual, expected);
343 }
344 }
345 }
346
347 #[test]
348 fn message_version_uses_config_before_versioned_flag() {
349 let legacy = Message::default();
350 assert_eq!(yellowstone_message_version(&legacy), YellowstoneMessageVersion::Legacy);
351
352 let v0 = Message { versioned: true, ..Message::default() };
353 assert_eq!(yellowstone_message_version(&v0), YellowstoneMessageVersion::V0);
354
355 let v1 = Message {
356 versioned: true,
357 config: Some(TransactionConfig::default()),
358 ..Message::default()
359 };
360 assert_eq!(yellowstone_message_version(&v1), YellowstoneMessageVersion::V1);
361
362 let v1_with_inconsistent_legacy_flag = Message {
363 versioned: false,
364 config: Some(TransactionConfig::default()),
365 ..Message::default()
366 };
367 assert_eq!(
368 yellowstone_message_version(&v1_with_inconsistent_legacy_flag),
369 YellowstoneMessageVersion::V1
370 );
371
372 for (message, version, json_version) in [
374 (legacy, YellowstoneMessageVersion::Legacy, "legacy"),
375 (v0, YellowstoneMessageVersion::V0, "v0"),
376 (v1, YellowstoneMessageVersion::V1, "v1"),
377 (v1_with_inconsistent_legacy_flag, YellowstoneMessageVersion::V1, "v1"),
378 ] {
379 let mut events = [
380 DexEvent::PumpFunCreate(Default::default()),
381 DexEvent::PumpFunCreateV2(Default::default()),
382 ];
383 fill_transaction_context(&mut events, &None);
384 fill_transaction_context(&mut events, &Some(Transaction::default()));
385 for event in &events {
386 let json = serde_json::to_value(event).unwrap();
387 assert!(json.as_object().unwrap().values().next().unwrap()["transaction_version"]
388 .is_null());
389 }
390 fill_transaction_context(
391 &mut events,
392 &Some(Transaction { message: Some(message), ..Default::default() }),
393 );
394 for event in &events {
395 let json = serde_json::to_value(event).unwrap();
396 assert_eq!(
397 json.as_object().unwrap().values().next().unwrap()["transaction_version"],
398 json_version
399 );
400 assert!(event.metadata().recent_blockhash.is_none());
401 }
402 assert_eq!(
403 serde_json::from_str::<YellowstoneMessageVersion>(&format!("\"{json_version}\""))
404 .unwrap(),
405 version
406 );
407 }
408 }
409}