Refactor: consolidate memo extraction for each message version (#22422)

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<Pubkey> = (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());
+}

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<T: ExtractMemos>(message: &T) -> Option<String> {
@@ -27,12 +33,10 @@ pub fn extract_and_fmt_memos<T: ExtractMemos>(message: &T) -> Option<String> {
     }
 }
 
-fn maybe_push_parsed_memo(memos: &mut Vec<String>, 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<String> {
-        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<String> {
-        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<Item = &'a Pubkey>,
+    instructions: &[CompiledInstruction],
+) -> Vec<String> {
+    let mut account_keys: Vec<KeyType> = 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);
     }
 }