Skip to main content

codama_plugin_core/
resolve_directives_visitor.rs

1use crate::DirectiveResolver;
2use codama_attributes::{
3    Attribute, CodamaDirective, Resolvable, SeedDirectiveType, TypeDirectiveNode,
4};
5use codama_errors::CodamaResult;
6use codama_korok_visitors::KorokVisitor;
7use codama_koroks::*;
8use codama_nodes::{InstructionInputValueNode, ValueNode};
9
10/// A visitor that resolves all `Resolvable::Unresolved` entries in the korok tree
11/// by delegating to the `DirectiveResolver`.
12///
13/// This visitor runs automatically as the first step in `resolve_plugins()`,
14/// before any lifecycle hooks fire.
15pub struct ResolveDirectivesVisitor<'a> {
16    resolver: &'a dyn DirectiveResolver,
17}
18
19impl<'a> ResolveDirectivesVisitor<'a> {
20    pub fn new(resolver: &'a dyn DirectiveResolver) -> Self {
21        Self { resolver }
22    }
23
24    fn resolve_attributes(
25        &self,
26        attributes: &mut codama_attributes::Attributes,
27    ) -> CodamaResult<()> {
28        for attr in attributes.iter_mut() {
29            let Attribute::Codama(codama_attr) = attr else {
30                continue;
31            };
32            self.resolve_directive(codama_attr.directive.as_mut())?;
33        }
34        Ok(())
35    }
36
37    fn resolve_directive(&self, directive: &mut CodamaDirective) -> CodamaResult<()> {
38        match directive {
39            CodamaDirective::Type(d) => {
40                self.resolve_type(&mut d.node)?;
41            }
42            CodamaDirective::DefaultValue(d) => {
43                self.resolve_instruction_input_value(&mut d.node)?;
44            }
45            CodamaDirective::Account(d) => {
46                if let Some(ref mut dv) = d.default_value {
47                    self.resolve_instruction_input_value(dv)?;
48                }
49            }
50            CodamaDirective::Field(d) => {
51                self.resolve_type_node(&mut d.r#type)?;
52                if let Some(ref mut dv) = d.default_value {
53                    self.resolve_value_node(dv)?;
54                }
55            }
56            CodamaDirective::Argument(d) => {
57                self.resolve_type_node(&mut d.r#type)?;
58                if let Some(ref mut dv) = d.default_value {
59                    self.resolve_instruction_input_value(dv)?;
60                }
61            }
62            CodamaDirective::Seed(d) => match &mut d.seed {
63                SeedDirectiveType::Variable { r#type, .. } => {
64                    self.resolve_type_node(r#type)?;
65                }
66                SeedDirectiveType::Constant { r#type, value } => {
67                    self.resolve_type_node(r#type)?;
68                    self.resolve_value_node(value)?;
69                }
70                SeedDirectiveType::Linked(_) => {}
71            },
72            // Other directives don't contain resolvable slots.
73            _ => {}
74        }
75        Ok(())
76    }
77
78    fn resolve_type(&self, resolvable: &mut Resolvable<TypeDirectiveNode>) -> CodamaResult<()> {
79        if let Resolvable::Unresolved(directive) = resolvable {
80            let resolved = self.resolver.resolve_type_directive(directive)?;
81            *resolvable = Resolvable::Resolved(resolved.into());
82        }
83        Ok(())
84    }
85
86    fn resolve_type_node(
87        &self,
88        resolvable: &mut Resolvable<codama_nodes::TypeNode>,
89    ) -> CodamaResult<()> {
90        if let Resolvable::Unresolved(directive) = resolvable {
91            let registered = self.resolver.resolve_type_directive(directive)?;
92            let type_node = codama_nodes::TypeNode::try_from(registered)?;
93            *resolvable = Resolvable::Resolved(type_node);
94        }
95        Ok(())
96    }
97
98    fn resolve_instruction_input_value(
99        &self,
100        resolvable: &mut Resolvable<InstructionInputValueNode>,
101    ) -> CodamaResult<()> {
102        if let Resolvable::Unresolved(directive) = resolvable {
103            let resolved = self.resolver.resolve_value_directive(directive)?;
104            *resolvable = Resolvable::Resolved(resolved);
105        }
106        Ok(())
107    }
108
109    fn resolve_value_node(&self, resolvable: &mut Resolvable<ValueNode>) -> CodamaResult<()> {
110        if let Resolvable::Unresolved(directive) = resolvable {
111            let instruction_input = self.resolver.resolve_value_directive(directive)?;
112            let value_node = ValueNode::try_from(instruction_input)?;
113            *resolvable = Resolvable::Resolved(value_node);
114        }
115        Ok(())
116    }
117}
118
119impl KorokVisitor for ResolveDirectivesVisitor<'_> {
120    fn visit_root(&mut self, korok: &mut RootKorok) -> CodamaResult<()> {
121        self.visit_children(korok)
122    }
123
124    fn visit_crate(&mut self, korok: &mut CrateKorok) -> CodamaResult<()> {
125        self.resolve_attributes(&mut korok.attributes)?;
126        self.visit_children(korok)
127    }
128
129    fn visit_file_module(&mut self, korok: &mut FileModuleKorok) -> CodamaResult<()> {
130        self.resolve_attributes(&mut korok.attributes)?;
131        self.visit_children(korok)
132    }
133
134    fn visit_module(&mut self, korok: &mut ModuleKorok) -> CodamaResult<()> {
135        self.resolve_attributes(&mut korok.attributes)?;
136        self.visit_children(korok)
137    }
138
139    fn visit_struct(&mut self, korok: &mut StructKorok) -> CodamaResult<()> {
140        self.resolve_attributes(&mut korok.attributes)?;
141        self.visit_children(korok)
142    }
143
144    fn visit_enum(&mut self, korok: &mut EnumKorok) -> CodamaResult<()> {
145        self.resolve_attributes(&mut korok.attributes)?;
146        self.visit_children(korok)
147    }
148
149    fn visit_enum_variant(&mut self, korok: &mut EnumVariantKorok) -> CodamaResult<()> {
150        self.resolve_attributes(&mut korok.attributes)?;
151        self.visit_children(korok)
152    }
153
154    fn visit_field(&mut self, korok: &mut FieldKorok) -> CodamaResult<()> {
155        self.resolve_attributes(&mut korok.attributes)
156    }
157
158    fn visit_unsupported_item(&mut self, korok: &mut UnsupportedItemKorok) -> CodamaResult<()> {
159        self.resolve_attributes(&mut korok.attributes)
160    }
161
162    fn visit_impl(&mut self, korok: &mut ImplKorok) -> CodamaResult<()> {
163        self.resolve_attributes(&mut korok.attributes)?;
164        self.visit_children(korok)
165    }
166
167    fn visit_const(&mut self, korok: &mut ConstKorok) -> CodamaResult<()> {
168        self.resolve_attributes(&mut korok.attributes)
169    }
170
171    fn visit_type_alias(&mut self, korok: &mut TypeAliasKorok) -> CodamaResult<()> {
172        self.resolve_attributes(&mut korok.attributes)
173    }
174
175    fn visit_unsupported_impl_item(
176        &mut self,
177        korok: &mut UnsupportedImplItemKorok,
178    ) -> CodamaResult<()> {
179        self.resolve_attributes(&mut korok.attributes)
180    }
181}