Skip to main content

hyperstack_idl/analysis/
connect.rs

1use crate::analysis::relations::build_account_index;
2use crate::search::suggest_similar;
3use crate::types::IdlSpec;
4use std::collections::HashSet;
5
6#[derive(Debug, Clone)]
7pub struct AccountRole {
8    pub writable: bool,
9    pub signer: bool,
10    pub pda: bool,
11}
12
13#[derive(Debug, Clone)]
14pub struct InstructionContext {
15    pub instruction_name: String,
16    pub from_role: AccountRole,
17    pub to_role: AccountRole,
18    pub all_accounts: Vec<String>,
19}
20
21#[derive(Debug, Clone)]
22pub struct DirectConnection {
23    pub from: String,
24    pub to: String,
25    pub instructions: Vec<InstructionContext>,
26}
27
28#[derive(Debug, Clone)]
29pub struct TransitiveConnection {
30    pub from: String,
31    pub intermediary: String,
32    pub to: String,
33    pub hop1_instruction: String,
34    pub hop2_instruction: String,
35}
36
37#[derive(Debug, Clone)]
38pub struct ConnectionReport {
39    pub new_account: String,
40    pub direct: Vec<DirectConnection>,
41    pub transitive: Vec<TransitiveConnection>,
42    pub invalid_existing: Vec<(String, Vec<String>)>,
43}
44
45const INFRASTRUCTURE_ACCOUNTS: &[&str] = &[
46    "system_program",
47    "token_program",
48    "rent",
49    "event_authority",
50    "program",
51    "associated_token_program",
52    "memo_program",
53    "token_2022_program",
54    "clock",
55    "instructions",
56    "sysvar_instructions",
57];
58
59pub fn find_connections(idl: &IdlSpec, new_account: &str, existing: &[&str]) -> ConnectionReport {
60    let all_account_names: Vec<&str> = idl
61        .instructions
62        .iter()
63        .flat_map(|ix| ix.accounts.iter().map(|a| a.name.as_str()))
64        .collect::<HashSet<_>>()
65        .into_iter()
66        .collect();
67
68    let new_account_exists = all_account_names.contains(&new_account);
69
70    let mut invalid_existing = Vec::new();
71    let mut valid_existing = Vec::new();
72
73    for &account in existing {
74        if all_account_names.contains(&account) {
75            valid_existing.push(account);
76        } else {
77            let suggestions = suggest_similar(account, &all_account_names, 3);
78            let suggestion_names: Vec<String> =
79                suggestions.iter().map(|s| s.candidate.clone()).collect();
80            invalid_existing.push((account.to_string(), suggestion_names));
81        }
82    }
83
84    if !new_account_exists {
85        return ConnectionReport {
86            new_account: new_account.to_string(),
87            direct: Vec::new(),
88            transitive: Vec::new(),
89            invalid_existing,
90        };
91    }
92
93    let mut direct = Vec::new();
94    for &existing_account in &valid_existing {
95        let mut instructions = Vec::new();
96
97        for instruction in &idl.instructions {
98            let account_names: Vec<&str> = instruction
99                .accounts
100                .iter()
101                .map(|account| account.name.as_str())
102                .collect();
103
104            if account_names.contains(&new_account) && account_names.contains(&existing_account) {
105                let from_account = instruction
106                    .accounts
107                    .iter()
108                    .find(|account| account.name == new_account);
109                let to_account = instruction
110                    .accounts
111                    .iter()
112                    .find(|account| account.name == existing_account);
113
114                if let (Some(from_account), Some(to_account)) = (from_account, to_account) {
115                    instructions.push(InstructionContext {
116                        instruction_name: instruction.name.clone(),
117                        from_role: AccountRole {
118                            writable: from_account.is_mut,
119                            signer: from_account.is_signer,
120                            pda: from_account.pda.is_some(),
121                        },
122                        to_role: AccountRole {
123                            writable: to_account.is_mut,
124                            signer: to_account.is_signer,
125                            pda: to_account.pda.is_some(),
126                        },
127                        all_accounts: account_names.iter().map(|name| name.to_string()).collect(),
128                    });
129                }
130            }
131        }
132
133        if !instructions.is_empty() {
134            direct.push(DirectConnection {
135                from: new_account.to_string(),
136                to: existing_account.to_string(),
137                instructions,
138            });
139        }
140    }
141
142    let mut transitive = Vec::new();
143    let directly_connected: HashSet<&str> = direct
144        .iter()
145        .map(|connection| connection.to.as_str())
146        .collect();
147    let unconnected: Vec<&str> = valid_existing
148        .iter()
149        .filter(|&&account| !directly_connected.contains(account))
150        .copied()
151        .collect();
152
153    if !unconnected.is_empty() {
154        let index = build_account_index(idl);
155        let new_account_instructions: HashSet<&str> = index
156            .get(new_account)
157            .map(|usage| {
158                usage
159                    .instructions
160                    .iter()
161                    .map(|instruction| instruction.name.as_str())
162                    .collect()
163            })
164            .unwrap_or_default();
165
166        for &target in &unconnected {
167            let target_instructions: HashSet<&str> = index
168                .get(target)
169                .map(|usage| {
170                    usage
171                        .instructions
172                        .iter()
173                        .map(|instruction| instruction.name.as_str())
174                        .collect()
175                })
176                .unwrap_or_default();
177
178            for (intermediary, usage) in &index {
179                if intermediary == new_account || intermediary == target {
180                    continue;
181                }
182                if INFRASTRUCTURE_ACCOUNTS.contains(&intermediary.as_str()) {
183                    continue;
184                }
185
186                let intermediary_instructions: HashSet<&str> = usage
187                    .instructions
188                    .iter()
189                    .map(|instruction| instruction.name.as_str())
190                    .collect();
191
192                let hop1 = new_account_instructions
193                    .iter()
194                    .find(|instruction| intermediary_instructions.contains(**instruction));
195                let hop2 = target_instructions
196                    .iter()
197                    .find(|instruction| intermediary_instructions.contains(**instruction));
198
199                if let (Some(hop1), Some(hop2)) = (hop1, hop2) {
200                    transitive.push(TransitiveConnection {
201                        from: new_account.to_string(),
202                        intermediary: intermediary.clone(),
203                        to: target.to_string(),
204                        hop1_instruction: (*hop1).to_string(),
205                        hop2_instruction: (*hop2).to_string(),
206                    });
207                    break;
208                }
209            }
210        }
211    }
212
213    ConnectionReport {
214        new_account: new_account.to_string(),
215        direct,
216        transitive,
217        invalid_existing,
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use crate::parse::parse_idl_file;
225    use std::path::PathBuf;
226
227    fn meteora_fixture() -> IdlSpec {
228        let path =
229            PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/meteora_dlmm.json");
230        parse_idl_file(&path).expect("should parse meteora_dlmm.json")
231    }
232
233    fn ore_fixture() -> IdlSpec {
234        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/ore.json");
235        parse_idl_file(&path).expect("should parse ore.json")
236    }
237
238    #[test]
239    fn test_connect_reward_vault() {
240        let idl = meteora_fixture();
241        let report = find_connections(&idl, "reward_vault", &["lb_pair", "position"]);
242        assert!(
243            !report.direct.is_empty(),
244            "reward_vault should have direct connections"
245        );
246
247        let lb_pair_connection = report
248            .direct
249            .iter()
250            .find(|connection| connection.to == "lb_pair");
251        assert!(
252            lb_pair_connection.is_some(),
253            "reward_vault should connect to lb_pair"
254        );
255        assert!(!lb_pair_connection
256            .expect("connection should exist")
257            .instructions
258            .is_empty());
259    }
260
261    #[test]
262    fn test_connect_invalid_name() {
263        let idl = meteora_fixture();
264        let report = find_connections(&idl, "lb_pair", &["bogus_account_xyz"]);
265        assert!(
266            !report.invalid_existing.is_empty(),
267            "bogus_account_xyz should be invalid"
268        );
269
270        let (name, suggestions) = &report.invalid_existing[0];
271        assert_eq!(name, "bogus_account_xyz");
272        let _ = suggestions;
273    }
274
275    #[test]
276    fn test_connect_ore_entropyvar() {
277        let idl = ore_fixture();
278        let report = find_connections(&idl, "entropyVar", &["round"]);
279        assert!(
280            !report.direct.is_empty() || !report.transitive.is_empty(),
281            "entropyVar should connect to round somehow"
282        );
283    }
284}