Skip to main content

arete_interpreter/
program_sdk.rs

1use crate::ast::{
2    idl_type_snapshot_to_rust_string, AccountResolution, AmountDecimalsSource, IdlAccountSnapshot,
3    IdlArrayElementSnapshot, IdlArrayTypeSnapshot, IdlDefinedInnerSnapshot, IdlDefinedTypeSnapshot,
4    IdlEnumVariantFieldSnapshot, IdlEnumVariantSnapshot, IdlErrorSnapshot, IdlEventSnapshot,
5    IdlFieldSnapshot, IdlHashMapTypeSnapshot, IdlInstructionAccountSnapshot,
6    IdlInstructionSnapshot, IdlOptionTypeSnapshot, IdlSerializationSnapshot, IdlSnapshot,
7    IdlTypeDefKindSnapshot, IdlTypeDefSnapshot, IdlTypeSnapshot, IdlVecTypeSnapshot,
8    InstructionAccountDef, InstructionAmountHint, InstructionArgDef, InstructionDef, PdaDefinition,
9    PdaSeedDef, SerializableStackSpec, CURRENT_AST_VERSION,
10};
11use arete_idl as idl_parser;
12use std::collections::BTreeMap;
13
14fn sanitize_identifier_segment(name: &str) -> String {
15    let mut sanitized = String::new();
16
17    for ch in name.chars() {
18        if ch.is_ascii_alphanumeric() || ch == '_' {
19            sanitized.push(ch);
20        } else if !sanitized.ends_with('_') {
21            sanitized.push('_');
22        }
23    }
24
25    let sanitized = sanitized.trim_matches('_').to_string();
26    if sanitized.is_empty() {
27        return "value".to_string();
28    }
29    if sanitized
30        .chars()
31        .next()
32        .is_some_and(|ch| ch.is_ascii_digit())
33    {
34        return format!("_{}", sanitized);
35    }
36    sanitized
37}
38
39fn sanitize_identifier(name: &str) -> String {
40    sanitize_identifier_segment(name)
41}
42
43fn sanitize_seed_path(path: &str) -> String {
44    path.split('.')
45        .map(sanitize_identifier_segment)
46        .collect::<Vec<_>>()
47        .join(".")
48}
49
50pub fn build_program_only_stack_spec_from_idl(
51    idl: &idl_parser::IdlSpec,
52    stack_name: &str,
53) -> SerializableStackSpec {
54    let snapshot = convert_idl_to_snapshot(idl);
55    let program_id = snapshot.program_id.clone();
56    let pdas = extract_pdas_from_idl(idl);
57    let instructions = extract_instructions_from_idl(idl, &pdas);
58
59    let mut grouped_pdas = BTreeMap::new();
60    if !pdas.is_empty() {
61        grouped_pdas.insert(snapshot.name.clone(), pdas);
62    }
63
64    SerializableStackSpec {
65        ast_version: CURRENT_AST_VERSION.to_string(),
66        stack_name: stack_name.to_string(),
67        program_ids: program_id.into_iter().collect(),
68        idls: vec![snapshot],
69        entities: vec![],
70        pdas: grouped_pdas,
71        instructions,
72        content_hash: None,
73    }
74    .with_content_hash()
75}
76
77pub fn convert_idl_to_snapshot(idl: &idl_parser::IdlSpec) -> IdlSnapshot {
78    let mut types: Vec<IdlTypeDefSnapshot> = idl
79        .types
80        .iter()
81        .map(|typedef| IdlTypeDefSnapshot {
82            name: typedef.name.clone(),
83            docs: typedef.docs.clone(),
84            serialization: typedef.serialization.as_ref().map(|s| match s {
85                idl_parser::IdlSerialization::Borsh => IdlSerializationSnapshot::Borsh,
86                idl_parser::IdlSerialization::Bytemuck => IdlSerializationSnapshot::Bytemuck,
87                idl_parser::IdlSerialization::BytemuckUnsafe => {
88                    IdlSerializationSnapshot::BytemuckUnsafe
89                }
90            }),
91            type_def: match &typedef.type_def {
92                idl_parser::IdlTypeDefKind::Struct { kind, fields } => {
93                    IdlTypeDefKindSnapshot::Struct {
94                        kind: kind.clone(),
95                        fields: fields
96                            .iter()
97                            .map(|field| IdlFieldSnapshot {
98                                name: field.name.clone(),
99                                type_: convert_idl_type(&field.type_),
100                                amount_hint: field.amount_hint.clone(),
101                            })
102                            .collect(),
103                    }
104                }
105                idl_parser::IdlTypeDefKind::TupleStruct { kind, fields } => {
106                    IdlTypeDefKindSnapshot::TupleStruct {
107                        kind: kind.clone(),
108                        fields: fields.iter().map(convert_idl_type).collect(),
109                    }
110                }
111                idl_parser::IdlTypeDefKind::Enum { kind, variants } => {
112                    IdlTypeDefKindSnapshot::Enum {
113                        kind: kind.clone(),
114                        variants: variants.iter().map(convert_enum_variant).collect(),
115                    }
116                }
117            },
118        })
119        .collect();
120
121    for account in &idl.accounts {
122        if types.iter().any(|typedef| typedef.name == account.name) {
123            continue;
124        }
125
126        if let Some(type_def) = &account.type_def {
127            let snapshot_type = match type_def {
128                idl_parser::IdlTypeDefKind::Struct { kind, fields } => {
129                    IdlTypeDefKindSnapshot::Struct {
130                        kind: kind.clone(),
131                        fields: fields
132                            .iter()
133                            .map(|field| IdlFieldSnapshot {
134                                name: field.name.clone(),
135                                type_: convert_idl_type(&field.type_),
136                                amount_hint: field.amount_hint.clone(),
137                            })
138                            .collect(),
139                    }
140                }
141                idl_parser::IdlTypeDefKind::TupleStruct { kind, fields } => {
142                    IdlTypeDefKindSnapshot::TupleStruct {
143                        kind: kind.clone(),
144                        fields: fields.iter().map(convert_idl_type).collect(),
145                    }
146                }
147                idl_parser::IdlTypeDefKind::Enum { kind, variants } => {
148                    IdlTypeDefKindSnapshot::Enum {
149                        kind: kind.clone(),
150                        variants: variants.iter().map(convert_enum_variant).collect(),
151                    }
152                }
153            };
154            types.push(IdlTypeDefSnapshot {
155                name: account.name.clone(),
156                docs: account.docs.clone(),
157                serialization: None,
158                type_def: snapshot_type,
159            });
160        }
161    }
162
163    let uses_steel_discriminant = idl.instructions.iter().any(|instruction| {
164        instruction.discriminant.is_some() && instruction.discriminator.is_empty()
165    });
166    let discriminant_size = if uses_steel_discriminant { 1 } else { 8 };
167    let program_id = idl.address.clone().or_else(|| {
168        idl.metadata
169            .as_ref()
170            .and_then(|metadata| metadata.address.clone())
171    });
172
173    IdlSnapshot {
174        name: idl.get_name().to_string(),
175        program_id: program_id.clone(),
176        version: idl.get_version().to_string(),
177        accounts: idl
178            .accounts
179            .iter()
180            .map(|account| {
181                let serialization = idl
182                    .types
183                    .iter()
184                    .find(|typedef| typedef.name == account.name)
185                    .and_then(|typedef| typedef.serialization.as_ref())
186                    .map(|serialization| match serialization {
187                        idl_parser::IdlSerialization::Borsh => IdlSerializationSnapshot::Borsh,
188                        idl_parser::IdlSerialization::Bytemuck => {
189                            IdlSerializationSnapshot::Bytemuck
190                        }
191                        idl_parser::IdlSerialization::BytemuckUnsafe => {
192                            IdlSerializationSnapshot::BytemuckUnsafe
193                        }
194                    });
195
196                IdlAccountSnapshot {
197                    name: account.name.clone(),
198                    discriminator: account.get_discriminator(),
199                    docs: account.docs.clone(),
200                    serialization,
201                    fields: account.type_def.as_ref().map_or_else(Vec::new, |type_def| {
202                        match type_def {
203                            idl_parser::IdlTypeDefKind::Struct { fields, .. } => fields
204                                .iter()
205                                .map(|field| IdlFieldSnapshot {
206                                    name: field.name.clone(),
207                                    type_: convert_idl_type(&field.type_),
208                                    amount_hint: field.amount_hint.clone(),
209                                })
210                                .collect(),
211                            _ => Vec::new(),
212                        }
213                    }),
214                    type_def: None,
215                }
216            })
217            .collect(),
218        instructions: idl
219            .instructions
220            .iter()
221            .map(|instruction| IdlInstructionSnapshot {
222                name: instruction.name.clone(),
223                discriminator: instruction.get_discriminator(),
224                discriminant: instruction.discriminant.clone(),
225                docs: instruction.docs.clone(),
226                accounts: instruction
227                    .flattened_accounts()
228                    .iter()
229                    .map(|account| IdlInstructionAccountSnapshot {
230                        name: account.name.clone(),
231                        writable: account.is_mut,
232                        signer: account.is_signer,
233                        optional: account.optional,
234                        address: account.address.clone(),
235                        docs: account.docs.clone(),
236                    })
237                    .collect(),
238                args: instruction
239                    .args
240                    .iter()
241                    .map(|arg| IdlFieldSnapshot {
242                        name: arg.name.clone(),
243                        type_: convert_idl_type(&arg.type_),
244                        amount_hint: arg.amount_hint.clone(),
245                    })
246                    .collect(),
247            })
248            .collect(),
249        types,
250        events: idl
251            .events
252            .iter()
253            .map(|event| IdlEventSnapshot {
254                name: event.name.clone(),
255                discriminator: event.get_discriminator(),
256                docs: event.docs.clone(),
257                fields: event
258                    .fields
259                    .iter()
260                    .map(|field| IdlFieldSnapshot {
261                        name: field.name.clone(),
262                        type_: convert_idl_type(&field.type_),
263                        amount_hint: field.amount_hint.clone(),
264                    })
265                    .collect(),
266            })
267            .collect(),
268        errors: idl
269            .errors
270            .iter()
271            .map(|error| IdlErrorSnapshot {
272                code: error.code,
273                name: error.name.clone(),
274                msg: error.msg.clone(),
275            })
276            .collect(),
277        discriminant_size,
278    }
279}
280
281pub fn extract_pdas_from_idl(idl: &idl_parser::IdlSpec) -> BTreeMap<String, PdaDefinition> {
282    let mut pdas = BTreeMap::new();
283
284    for pda in &idl.pdas {
285        let pda_name = sanitize_identifier(&pda.name);
286        let pda_def = convert_idl_pda_to_def(&pda_name, &pda.seeds, pda.program.as_ref());
287        pdas.insert(pda_name, pda_def);
288    }
289
290    for instruction in &idl.instructions {
291        for account in instruction.flattened_accounts() {
292            if let Some(pda_info) = &account.pda {
293                let pda_name =
294                    sanitize_identifier(pda_info.name.as_deref().unwrap_or(&account.name));
295                let pda_def =
296                    convert_idl_pda_to_def(&pda_name, &pda_info.seeds, pda_info.program.as_ref());
297                pdas.entry(pda_name).or_insert(pda_def);
298            }
299        }
300    }
301
302    pdas
303}
304
305pub fn extract_instructions_from_idl(
306    idl: &idl_parser::IdlSpec,
307    pdas: &BTreeMap<String, PdaDefinition>,
308) -> Vec<InstructionDef> {
309    let program_id = idl.address.clone().or_else(|| {
310        idl.metadata
311            .as_ref()
312            .and_then(|metadata| metadata.address.clone())
313    });
314
315    let uses_steel = idl.instructions.iter().any(|instruction| {
316        instruction.discriminant.is_some() && instruction.discriminator.is_empty()
317    });
318    let discriminator_size = if uses_steel { 1 } else { 8 };
319
320    idl.instructions
321        .iter()
322        .map(|instruction| {
323            let accounts = instruction
324                .flattened_accounts()
325                .iter()
326                .map(|account| convert_account_to_def(account, pdas))
327                .collect();
328
329            let args = instruction
330                .args
331                .iter()
332                .map(|arg| InstructionArgDef {
333                    name: arg.name.clone(),
334                    arg_type: idl_type_snapshot_to_rust_string(&convert_idl_type(&arg.type_)),
335                    docs: vec![],
336                    amount_hint: arg.amount_hint.as_ref().map(convert_amount_hint),
337                })
338                .collect();
339
340            InstructionDef {
341                name: instruction.name.clone(),
342                discriminator: instruction.get_discriminator(),
343                discriminator_size,
344                accounts,
345                args,
346                errors: Vec::new(),
347                program_id: program_id.clone(),
348                docs: instruction.docs.clone(),
349            }
350        })
351        .collect()
352}
353
354pub fn convert_idl_type(idl_type: &idl_parser::IdlType) -> IdlTypeSnapshot {
355    match idl_type {
356        idl_parser::IdlType::Simple(simple) => IdlTypeSnapshot::Simple(simple.clone()),
357        idl_parser::IdlType::Array(array) => IdlTypeSnapshot::Array(IdlArrayTypeSnapshot {
358            array: array
359                .array
360                .iter()
361                .map(|element| match element {
362                    idl_parser::IdlTypeArrayElement::Nested(ty) => {
363                        IdlArrayElementSnapshot::Type(convert_idl_type(ty))
364                    }
365                    idl_parser::IdlTypeArrayElement::Type(type_name) => {
366                        IdlArrayElementSnapshot::TypeName(type_name.clone())
367                    }
368                    idl_parser::IdlTypeArrayElement::Size(size) => {
369                        IdlArrayElementSnapshot::Size(*size)
370                    }
371                })
372                .collect(),
373        }),
374        idl_parser::IdlType::Option(option) => IdlTypeSnapshot::Option(IdlOptionTypeSnapshot {
375            option: Box::new(convert_idl_type(&option.option)),
376        }),
377        idl_parser::IdlType::Vec(vec_type) => IdlTypeSnapshot::Vec(IdlVecTypeSnapshot {
378            vec: Box::new(convert_idl_type(&vec_type.vec)),
379        }),
380        idl_parser::IdlType::Defined(defined) => IdlTypeSnapshot::Defined(IdlDefinedTypeSnapshot {
381            defined: match &defined.defined {
382                idl_parser::IdlTypeDefinedInner::Named { name } => {
383                    IdlDefinedInnerSnapshot::Named { name: name.clone() }
384                }
385                idl_parser::IdlTypeDefinedInner::Simple(simple) => {
386                    IdlDefinedInnerSnapshot::Simple(simple.clone())
387                }
388            },
389        }),
390        idl_parser::IdlType::HashMap(hash_map) => {
391            IdlTypeSnapshot::HashMap(IdlHashMapTypeSnapshot {
392                hash_map: (
393                    Box::new(convert_idl_type(&hash_map.hash_map.0)),
394                    Box::new(convert_idl_type(&hash_map.hash_map.1)),
395                ),
396            })
397        }
398    }
399}
400
401fn convert_amount_hint(hint: &idl_parser::IdlAmountHint) -> InstructionAmountHint {
402    InstructionAmountHint {
403        decimals_source: match &hint.decimals_source {
404            idl_parser::IdlAmountDecimalsSource::ArgMint { arg_name } => {
405                AmountDecimalsSource::ArgMint {
406                    arg_name: arg_name.clone(),
407                }
408            }
409            idl_parser::IdlAmountDecimalsSource::ArgDecimals { arg_name } => {
410                AmountDecimalsSource::ArgDecimals {
411                    arg_name: arg_name.clone(),
412                }
413            }
414            idl_parser::IdlAmountDecimalsSource::KnownAccount { account_name } => {
415                AmountDecimalsSource::KnownAccount {
416                    account_name: account_name.clone(),
417                }
418            }
419            idl_parser::IdlAmountDecimalsSource::Constant { decimals } => {
420                AmountDecimalsSource::Constant {
421                    decimals: *decimals,
422                }
423            }
424        },
425    }
426}
427
428fn convert_enum_variant(variant: &idl_parser::IdlEnumVariant) -> IdlEnumVariantSnapshot {
429    IdlEnumVariantSnapshot {
430        name: variant.name.clone(),
431        fields: variant
432            .fields
433            .iter()
434            .map(|field| match field {
435                idl_parser::IdlEnumVariantField::Named(named) => {
436                    IdlEnumVariantFieldSnapshot::Named(IdlFieldSnapshot {
437                        name: named.name.clone(),
438                        type_: convert_idl_type(&named.type_),
439                        amount_hint: named.amount_hint.clone(),
440                    })
441                }
442                idl_parser::IdlEnumVariantField::Tuple(tuple) => {
443                    IdlEnumVariantFieldSnapshot::Tuple(convert_idl_type(tuple))
444                }
445            })
446            .collect(),
447    }
448}
449
450fn convert_idl_pda_to_def(
451    name: &str,
452    pda_seeds: &[idl_parser::IdlPdaSeed],
453    pda_program: Option<&idl_parser::IdlPdaProgram>,
454) -> PdaDefinition {
455    let seeds = pda_seeds
456        .iter()
457        .map(|seed| match seed {
458            idl_parser::IdlPdaSeed::Const { value } => {
459                if let Ok(string) = String::from_utf8(value.clone()) {
460                    PdaSeedDef::Literal { value: string }
461                } else {
462                    PdaSeedDef::Bytes {
463                        value: value.clone(),
464                    }
465                }
466            }
467            idl_parser::IdlPdaSeed::Account { path, .. } => PdaSeedDef::AccountRef {
468                account_name: sanitize_seed_path(path),
469            },
470            idl_parser::IdlPdaSeed::Arg { path, arg_type } => PdaSeedDef::ArgRef {
471                arg_name: sanitize_seed_path(path),
472                arg_type: arg_type.clone(),
473            },
474        })
475        .collect();
476
477    let program_id = pda_program.and_then(|program| match program {
478        idl_parser::IdlPdaProgram::Literal { value, .. } => Some(value.clone()),
479        idl_parser::IdlPdaProgram::Const { value, .. } => Some(bs58::encode(value).into_string()),
480        idl_parser::IdlPdaProgram::Account { .. } => None,
481    });
482
483    PdaDefinition {
484        name: name.to_string(),
485        seeds,
486        program_id,
487    }
488}
489
490fn convert_account_to_def(
491    account: &idl_parser::IdlAccountArg,
492    pdas: &BTreeMap<String, PdaDefinition>,
493) -> InstructionAccountDef {
494    let resolution = if account.is_signer && account.address.is_none() && account.pda.is_none() {
495        AccountResolution::Signer
496    } else if let Some(address) = &account.address {
497        AccountResolution::Known {
498            address: address.clone(),
499        }
500    } else if account.pda.is_some() {
501        let pda_name = sanitize_identifier(
502            account
503                .pda
504                .as_ref()
505                .and_then(|pda| pda.name.as_deref())
506                .unwrap_or(&account.name),
507        );
508        if pdas.contains_key(&pda_name) {
509            AccountResolution::PdaRef {
510                pda_name: pda_name.to_string(),
511            }
512        } else if let Some(pda_info) = &account.pda {
513            let pda_def =
514                convert_idl_pda_to_def(&pda_name, &pda_info.seeds, pda_info.program.as_ref());
515            AccountResolution::PdaInline {
516                seeds: pda_def.seeds,
517                program_id: pda_def.program_id,
518            }
519        } else {
520            AccountResolution::UserProvided
521        }
522    } else if pdas.contains_key(&sanitize_identifier(&account.name)) {
523        AccountResolution::PdaRef {
524            pda_name: sanitize_identifier(&account.name),
525        }
526    } else {
527        AccountResolution::UserProvided
528    };
529
530    InstructionAccountDef {
531        name: sanitize_identifier(&account.name),
532        is_signer: account.is_signer,
533        is_writable: account.is_mut,
534        resolution,
535        is_optional: account.optional,
536        docs: account.docs.clone(),
537    }
538}
539
540#[cfg(test)]
541mod tests {
542    use super::*;
543
544    #[test]
545    fn builds_program_only_stack_spec_from_raw_idl() {
546        let idl = arete_idl::parse::parse_idl_content(
547            r#"{
548              "address": "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
549              "version": "0.0.0",
550              "name": "token",
551              "instructions": [
552                {
553                  "name": "InitializeMint2",
554                  "accounts": [
555                    { "name": "mint", "isMut": true, "isSigner": false }
556                  ],
557                  "args": [
558                    { "name": "decimals", "type": "u8" },
559                    { "name": "mintAuthority", "type": "publicKey" }
560                  ],
561                  "discriminant": { "type": "u8", "value": 20 }
562                }
563              ],
564              "accounts": [],
565              "types": [],
566              "events": [],
567              "errors": []
568            }"#,
569        )
570        .expect("IDL should parse");
571
572        let spec = build_program_only_stack_spec_from_idl(&idl, "SplToken");
573        assert_eq!(spec.stack_name, "SplToken");
574        assert!(spec.entities.is_empty());
575        assert_eq!(spec.idls.len(), 1);
576        assert_eq!(spec.instructions.len(), 1);
577        assert_eq!(spec.instructions[0].name, "InitializeMint2");
578        assert_eq!(spec.instructions[0].discriminator, vec![20]);
579        assert_eq!(spec.instructions[0].discriminator_size, 1);
580        assert_eq!(
581            spec.instructions[0].program_id.as_deref(),
582            Some("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA")
583        );
584        assert_eq!(
585            spec.instructions[0].args[1].arg_type,
586            "solana_pubkey::Pubkey"
587        );
588        assert!(spec.content_hash.is_some());
589    }
590
591    #[test]
592    fn preserves_nested_seed_paths_when_building_program_only_specs() {
593        let idl = arete_idl::parse::parse_idl_content(
594            r#"{
595              "address": "Prog111111111111111111111111111111111111111",
596              "version": "0.0.0",
597              "name": "demo",
598              "instructions": [
599                {
600                  "name": "proposalCreate",
601                  "accounts": [
602                    {
603                      "name": "proposal",
604                      "isMut": true,
605                      "isSigner": false,
606                      "pda": {
607                        "name": "proposal",
608                        "seeds": [
609                          {
610                            "kind": "arg",
611                            "path": "args.transactionIndex",
612                            "type": "u64"
613                          }
614                        ]
615                      }
616                    }
617                  ],
618                  "args": [
619                    {
620                      "name": "args",
621                      "type": {
622                        "defined": {
623                          "name": "ProposalArgs"
624                        }
625                      }
626                    }
627                  ],
628                  "discriminant": { "type": "u8", "value": 3 }
629                }
630              ],
631              "accounts": [],
632              "types": [
633                {
634                  "name": "ProposalArgs",
635                  "type": {
636                    "kind": "struct",
637                    "fields": [
638                      { "name": "transactionIndex", "type": "u64" }
639                    ]
640                  }
641                }
642              ],
643              "events": [],
644              "errors": []
645            }"#,
646        )
647        .expect("IDL should parse");
648
649        let spec = build_program_only_stack_spec_from_idl(&idl, "Demo");
650        let pda = spec
651            .pdas
652            .get("demo")
653            .and_then(|program| program.get("proposal"))
654            .expect("proposal PDA should be present");
655        assert_eq!(
656            pda.seeds,
657            vec![PdaSeedDef::ArgRef {
658                arg_name: "args.transactionIndex".to_string(),
659                arg_type: Some("u64".to_string()),
660            }]
661        );
662    }
663
664    #[test]
665    fn preserves_amount_hints_from_idl_args() {
666        let idl = arete_idl::parse::parse_idl_content(
667            r#"{
668              "address": "Prog111111111111111111111111111111111111111",
669              "version": "0.0.0",
670              "name": "demo",
671              "instructions": [
672                {
673                  "name": "deposit",
674                  "accounts": [],
675                  "args": [
676                    {
677                      "name": "amount",
678                      "type": "u64",
679                      "amountHint": {
680                        "decimalsSource": {
681                          "kind": "argMint",
682                          "argName": "mint"
683                        }
684                      }
685                    },
686                    {
687                      "name": "mint",
688                      "type": "publicKey"
689                    }
690                  ],
691                  "discriminant": { "type": "u8", "value": 7 }
692                }
693              ],
694              "accounts": [],
695              "types": [],
696              "events": [],
697              "errors": []
698            }"#,
699        )
700        .expect("IDL should parse");
701
702        let spec = build_program_only_stack_spec_from_idl(&idl, "Demo");
703        assert_eq!(
704            spec.instructions[0].args[0].amount_hint,
705            Some(InstructionAmountHint {
706                decimals_source: AmountDecimalsSource::ArgMint {
707                    arg_name: "mint".to_string(),
708                },
709            })
710        );
711        assert!(spec.idls[0].instructions[0].args[0].amount_hint.is_some());
712    }
713}