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

use {
    crate::parse_instruction::parse_memo_data,
    solana_sdk::{
        instruction::CompiledInstruction,
        message::{Message, SanitizedMessage},
        pubkey::Pubkey,
    },
};

// 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 {
    *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 {
    *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> {
    let memos = message.extract_memos();
    if memos.is_empty() {
        None
    } else {
        Some(memos.join("; "))
    }
}

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 {
    fn extract_memos(&self) -> Vec<String>;
}

impl ExtractMemos for Message {
    fn extract_memos(&self) -> Vec<String> {
        extract_memos_inner(self.account_keys.iter(), &self.instructions)
    }
}

impl ExtractMemos for SanitizedMessage {
    fn extract_memos(&self) -> Vec<String> {
        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::*;

    #[test]
    fn test_extract_memos_inner() {
        let fee_payer = Pubkey::new_unique();
        let another_program_id = Pubkey::new_unique();
        let memo0 = "Test memo";
        let memo1 = "🦖";
        let expected_memos = vec![
            format!("[{}] {}", memo0.len(), memo0),
            format!("[{}] {}", memo1.len(), memo1),
        ];
        let memo_instructions = vec![
            CompiledInstruction {
                program_id_index: 1,
                accounts: vec![],
                data: memo0.as_bytes().to_vec(),
            },
            CompiledInstruction {
                program_id_index: 2,
                accounts: vec![],
                data: memo1.as_bytes().to_vec(),
            },
            CompiledInstruction {
                program_id_index: 3,
                accounts: vec![],
                data: memo1.as_bytes().to_vec(),
            },
        ];
        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
        );
    }
}