From f4d15773379f82d662e25a066202a4de1897fa44 Mon Sep 17 00:00:00 2001 From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com> Date: Wed, 12 Jan 2022 00:23:03 +0000 Subject: [PATCH] Refactor: consolidate memo extraction for each message version (#22422) (#22435) (cherry picked from commit 35a5dd9c457502b3ab98119aa6cc0cd59c5fbb05) Co-authored-by: Justin Starry --- transaction-status/benches/extract_memos.rs | 33 +++++ transaction-status/src/extract_memos.rs | 134 +++++++++----------- 2 files changed, 93 insertions(+), 74 deletions(-) create mode 100644 transaction-status/benches/extract_memos.rs diff --git a/transaction-status/benches/extract_memos.rs b/transaction-status/benches/extract_memos.rs new file mode 100644 index 0000000000..4579a35f75 --- /dev/null +++ b/transaction-status/benches/extract_memos.rs @@ -0,0 +1,33 @@ +#![feature(test)] + +extern crate test; + +use { + solana_sdk::{instruction::CompiledInstruction, message::Message, pubkey::Pubkey}, + solana_transaction_status::extract_memos::{spl_memo_id_v1, spl_memo_id_v3, ExtractMemos}, + test::Bencher, +}; + +#[bench] +fn bench_extract_memos(b: &mut Bencher) { + let mut account_keys: Vec = (0..64).map(|_| Pubkey::new_unique()).collect(); + account_keys[62] = spl_memo_id_v1(); + account_keys[63] = spl_memo_id_v3(); + let memo = "Test memo"; + + let instructions: Vec<_> = (0..20) + .map(|i| CompiledInstruction { + program_id_index: 62 + (i % 2), + accounts: vec![], + data: memo.as_bytes().to_vec(), + }) + .collect(); + + let message = Message { + account_keys, + instructions, + ..Message::default() + }; + + b.iter(|| message.extract_memos()); +} diff --git a/transaction-status/src/extract_memos.rs b/transaction-status/src/extract_memos.rs index 0fb0564805..b83547b8bf 100644 --- a/transaction-status/src/extract_memos.rs +++ b/transaction-status/src/extract_memos.rs @@ -1,6 +1,7 @@ use { crate::parse_instruction::parse_memo_data, solana_sdk::{ + instruction::CompiledInstruction, message::{Message, SanitizedMessage}, pubkey::Pubkey, }, @@ -9,13 +10,18 @@ use { // A helper function to convert spl_memo::v1::id() as spl_sdk::pubkey::Pubkey to // solana_sdk::pubkey::Pubkey pub fn spl_memo_id_v1() -> Pubkey { - Pubkey::new_from_array(spl_memo::v1::id().to_bytes()) + *MEMO_PROGRAM_ID_V1 } // A helper function to convert spl_memo::id() as spl_sdk::pubkey::Pubkey to // solana_sdk::pubkey::Pubkey pub fn spl_memo_id_v3() -> Pubkey { - Pubkey::new_from_array(spl_memo::id().to_bytes()) + *MEMO_PROGRAM_ID_V3 +} + +lazy_static! { + static ref MEMO_PROGRAM_ID_V1: Pubkey = Pubkey::new_from_array(spl_memo::v1::id().to_bytes()); + static ref MEMO_PROGRAM_ID_V3: Pubkey = Pubkey::new_from_array(spl_memo::id().to_bytes()); } pub fn extract_and_fmt_memos(message: &T) -> Option { @@ -27,12 +33,10 @@ pub fn extract_and_fmt_memos(message: &T) -> Option { } } -fn maybe_push_parsed_memo(memos: &mut Vec, program_id: Pubkey, data: &[u8]) { - if program_id == spl_memo_id_v1() || program_id == spl_memo_id_v3() { - let memo_len = data.len(); - let parsed_memo = parse_memo_data(data).unwrap_or_else(|_| "(unparseable)".to_string()); - memos.push(format!("[{}] {}", memo_len, parsed_memo)); - } +fn extract_and_fmt_memo_data(data: &[u8]) -> String { + let memo_len = data.len(); + let parsed_memo = parse_memo_data(data).unwrap_or_else(|_| "(unparseable)".to_string()); + format!("[{}] {}", memo_len, parsed_memo) } pub trait ExtractMemos { @@ -41,50 +45,56 @@ pub trait ExtractMemos { impl ExtractMemos for Message { fn extract_memos(&self) -> Vec { - let mut memos = vec![]; - if self.account_keys.contains(&spl_memo_id_v1()) - || self.account_keys.contains(&spl_memo_id_v3()) - { - for instruction in &self.instructions { - let program_id = self.account_keys[instruction.program_id_index as usize]; - maybe_push_parsed_memo(&mut memos, program_id, &instruction.data); - } - } - memos + extract_memos_inner(self.account_keys.iter(), &self.instructions) } } impl ExtractMemos for SanitizedMessage { fn extract_memos(&self) -> Vec { - let mut memos = vec![]; - if self - .account_keys_iter() - .any(|&pubkey| pubkey == spl_memo_id_v1() || pubkey == spl_memo_id_v3()) - { - for (program_id, instruction) in self.program_instructions_iter() { - maybe_push_parsed_memo(&mut memos, *program_id, &instruction.data); - } - } - memos + extract_memos_inner(self.account_keys_iter(), self.instructions()) } } +enum KeyType<'a> { + MemoProgram, + OtherProgram, + Unknown(&'a Pubkey), +} + +fn extract_memos_inner<'a>( + account_keys: impl Iterator, + instructions: &[CompiledInstruction], +) -> Vec { + let mut account_keys: Vec = account_keys.map(KeyType::Unknown).collect(); + instructions + .iter() + .filter_map(|ix| { + let index = ix.program_id_index as usize; + let key_type = account_keys.get(index)?; + let memo_data = match key_type { + KeyType::MemoProgram => Some(&ix.data), + KeyType::OtherProgram => None, + KeyType::Unknown(program_id) => { + if **program_id == *MEMO_PROGRAM_ID_V1 || **program_id == *MEMO_PROGRAM_ID_V3 { + account_keys[index] = KeyType::MemoProgram; + Some(&ix.data) + } else { + account_keys[index] = KeyType::OtherProgram; + None + } + } + }?; + Some(extract_and_fmt_memo_data(memo_data)) + }) + .collect() +} + #[cfg(test)] mod test { - use { - super::*, - solana_sdk::{ - hash::Hash, - instruction::CompiledInstruction, - message::{ - v0::{self, LoadedAddresses}, - MessageHeader, - }, - }, - }; + use super::*; #[test] - fn test_extract_memos() { + fn test_extract_memos_inner() { let fee_payer = Pubkey::new_unique(); let another_program_id = Pubkey::new_unique(); let memo0 = "Test memo"; @@ -110,40 +120,16 @@ mod test { data: memo1.as_bytes().to_vec(), }, ]; - let message = Message::new_with_compiled_instructions( - 1, - 0, - 3, - vec![ - fee_payer, - spl_memo_id_v1(), - another_program_id, - spl_memo_id_v3(), - ], - Hash::default(), - memo_instructions.clone(), + let account_keys = vec![ + fee_payer, + spl_memo_id_v1(), + another_program_id, + spl_memo_id_v3(), + ]; + + assert_eq!( + extract_memos_inner(account_keys.iter(), &memo_instructions), + expected_memos ); - assert_eq!(message.extract_memos(), expected_memos); - - let sanitized_message = SanitizedMessage::Legacy(message); - assert_eq!(sanitized_message.extract_memos(), expected_memos); - - let sanitized_message = SanitizedMessage::V0(v0::LoadedMessage { - message: v0::Message { - header: MessageHeader { - num_required_signatures: 1, - num_readonly_signed_accounts: 0, - num_readonly_unsigned_accounts: 3, - }, - account_keys: vec![fee_payer], - instructions: memo_instructions, - ..v0::Message::default() - }, - loaded_addresses: LoadedAddresses { - writable: vec![], - readonly: vec![spl_memo_id_v1(), another_program_id, spl_memo_id_v3()], - }, - }); - assert_eq!(sanitized_message.extract_memos(), expected_memos); } }