rukalang/elab/
signatures.rs

1use super::helpers::{
2    has_meta_params, is_runtime_template, runtime_params, ty_to_type_expr, type_param_names,
3    variadic_param_name,
4};
5use super::prelude::*;
6
7impl Elaborator {
8    pub(super) fn new(program: &Program) -> Result<Self, ElabError> {
9        let mut templates = BTreeMap::new();
10        let mut enum_templates = BTreeMap::new();
11        let mut function_templates = BTreeMap::new();
12        for decl in &program.structs {
13            if templates.contains_key(&decl.name) {
14                return Err(ElabError::DuplicateStructDecl {
15                    name: decl.name.clone(),
16                });
17            }
18            let mut seen_params = BTreeSet::new();
19            for param in &decl.type_params {
20                if !seen_params.insert(param.clone()) {
21                    return Err(ElabError::DuplicateStructTypeParam {
22                        name: decl.name.clone(),
23                        param: param.clone(),
24                    });
25                }
26            }
27            let mut seen_fields = BTreeSet::new();
28            for field in &decl.fields {
29                if !seen_fields.insert(field.name.clone()) {
30                    return Err(ElabError::DuplicateStructField {
31                        name: decl.name.clone(),
32                        field: field.name.clone(),
33                    });
34                }
35            }
36            templates.insert(
37                decl.name.clone(),
38                StructTemplate {
39                    type_params: decl.type_params.clone(),
40                    fields: decl.fields.clone(),
41                },
42            );
43        }
44
45        for decl in &program.enums {
46            if templates.contains_key(&decl.name) || enum_templates.contains_key(&decl.name) {
47                return Err(ElabError::DuplicateStructDecl {
48                    name: decl.name.clone(),
49                });
50            }
51            let mut seen_params = BTreeSet::new();
52            for param in &decl.type_params {
53                if !seen_params.insert(param.clone()) {
54                    return Err(ElabError::DuplicateStructTypeParam {
55                        name: decl.name.clone(),
56                        param: param.clone(),
57                    });
58                }
59            }
60            enum_templates.insert(
61                decl.name.clone(),
62                EnumTemplate {
63                    type_params: decl.type_params.clone(),
64                    variants: decl
65                        .variants
66                        .iter()
67                        .map(|variant| (variant.name.clone(), variant.payload.clone()))
68                        .collect(),
69                },
70            );
71        }
72
73        for function in &program.functions {
74            let type_params = type_param_names(&function.params);
75            let variadic_param = variadic_param_name(&function.name, &function.params)?;
76            if !type_params.is_empty()
77                || has_meta_params(&function.params)
78                || variadic_param.is_some()
79            {
80                function_templates.insert(
81                    function.name.clone(),
82                    RuntimeFunctionTemplate {
83                        type_params,
84                        function: function.clone(),
85                    },
86                );
87            }
88        }
89
90        Ok(Self {
91            program: program.clone(),
92            templates,
93            enum_templates,
94            signatures: BTreeMap::new(),
95            function_templates,
96            specializations: BTreeMap::new(),
97            concrete_structs: BTreeMap::new(),
98            instantiating_structs: BTreeSet::new(),
99            instantiating_functions: BTreeSet::new(),
100            next_function_instance_id: 0,
101            pass_timings: Vec::new(),
102        })
103    }
104
105    pub(super) fn elaborate(&mut self) -> Result<Program, ElabError> {
106        self.resolve_signatures()?;
107
108        for module_index in 0..self.program.extern_modules.len() {
109            let function_count = self.program.extern_modules[module_index].functions.len();
110            for function_index in 0..function_count {
111                if is_runtime_template(
112                    &self.program.extern_modules[module_index].functions[function_index].params,
113                ) {
114                    continue;
115                }
116                let return_ty = self.program.extern_modules[module_index].functions[function_index]
117                    .return_type
118                    .ty
119                    .clone();
120                let resolved_return =
121                    self.resolve_type_expr(&return_ty, &BTreeMap::new(), &BTreeSet::new(), None)?;
122                self.ensure_ty_instantiated(&resolved_return)?;
123                self.program.extern_modules[module_index].functions[function_index]
124                    .return_type
125                    .ty = ty_to_type_expr(&resolved_return);
126
127                let param_count = self.program.extern_modules[module_index].functions
128                    [function_index]
129                    .params
130                    .len();
131                for param_index in 0..param_count {
132                    let param_ty = self.program.extern_modules[module_index].functions
133                        [function_index]
134                        .params[param_index]
135                        .ty
136                        .ty
137                        .clone();
138                    let resolved = self.resolve_type_expr(
139                        &param_ty,
140                        &BTreeMap::new(),
141                        &BTreeSet::new(),
142                        None,
143                    )?;
144                    self.ensure_ty_instantiated(&resolved)?;
145                    self.program.extern_modules[module_index].functions[function_index].params
146                        [param_index]
147                        .ty
148                        .ty = ty_to_type_expr(&resolved);
149                }
150            }
151        }
152
153        for function_index in 0..self.program.functions.len() {
154            if is_runtime_template(&self.program.functions[function_index].params) {
155                continue;
156            }
157            let return_ty = self.program.functions[function_index]
158                .return_type
159                .ty
160                .clone();
161            let resolved_return =
162                self.resolve_type_expr(&return_ty, &BTreeMap::new(), &BTreeSet::new(), None)?;
163            self.ensure_ty_instantiated(&resolved_return)?;
164            self.program.functions[function_index].return_type.ty =
165                ty_to_type_expr(&resolved_return);
166
167            let param_count = self.program.functions[function_index].params.len();
168            for param_index in 0..param_count {
169                let param_ty = self.program.functions[function_index].params[param_index]
170                    .ty
171                    .ty
172                    .clone();
173                let resolved =
174                    self.resolve_type_expr(&param_ty, &BTreeMap::new(), &BTreeSet::new(), None)?;
175                self.ensure_ty_instantiated(&resolved)?;
176                self.program.functions[function_index].params[param_index]
177                    .ty
178                    .ty = ty_to_type_expr(&resolved);
179            }
180            let mut function = self.program.functions[function_index].clone();
181            self.elaborate_function_body(&mut function)?;
182            self.program.functions[function_index] = function;
183        }
184
185        for module in &mut self.program.extern_modules {
186            module
187                .functions
188                .retain(|function| !is_runtime_template(&function.params));
189        }
190        self.program
191            .functions
192            .retain(|function| !is_runtime_template(&function.params));
193
194        self.program.structs = self.concrete_structs.values().cloned().collect();
195        Ok(self.program.clone())
196    }
197
198    pub(super) fn resolve_signatures(&mut self) -> Result<(), ElabError> {
199        let extern_modules = self.program.extern_modules.clone();
200        for module in &extern_modules {
201            for function in &module.functions {
202                if is_runtime_template(&function.params) {
203                    continue;
204                }
205                let name = format!("{}::{}", module.name, function.name);
206                let concrete_params = runtime_params(&function.params);
207                let mut params = Vec::with_capacity(concrete_params.len());
208                for param in &concrete_params {
209                    let ty = self.resolve_type_expr(
210                        &param.ty.ty,
211                        &BTreeMap::new(),
212                        &BTreeSet::new(),
213                        None,
214                    )?;
215                    self.ensure_ty_instantiated(&ty)?;
216                    params.push((param.ty.mode, ty));
217                }
218                let return_ty = self.resolve_type_expr(
219                    &function.return_type.ty,
220                    &BTreeMap::new(),
221                    &BTreeSet::new(),
222                    None,
223                )?;
224                self.ensure_ty_instantiated(&return_ty)?;
225                self.signatures
226                    .insert(name, FunctionSig { params, return_ty });
227            }
228        }
229
230        let functions = self.program.functions.clone();
231        for function in &functions {
232            if is_runtime_template(&function.params) {
233                continue;
234            }
235            let concrete_params = runtime_params(&function.params);
236            let mut params = Vec::with_capacity(concrete_params.len());
237            for param in &concrete_params {
238                let ty =
239                    self.resolve_type_expr(&param.ty.ty, &BTreeMap::new(), &BTreeSet::new(), None)?;
240                self.ensure_ty_instantiated(&ty)?;
241                params.push((param.ty.mode, ty));
242            }
243            let return_ty = self.resolve_type_expr(
244                &function.return_type.ty,
245                &BTreeMap::new(),
246                &BTreeSet::new(),
247                None,
248            )?;
249            self.ensure_ty_instantiated(&return_ty)?;
250            self.signatures
251                .insert(function.name.clone(), FunctionSig { params, return_ty });
252        }
253
254        let concrete_roots = self
255            .templates
256            .iter()
257            .filter_map(|(name, template)| {
258                if template.type_params.is_empty() {
259                    Some(name.clone())
260                } else {
261                    None
262                }
263            })
264            .collect::<Vec<_>>();
265        for name in concrete_roots {
266            self.ensure_ty_instantiated(&Ty::Struct {
267                name,
268                args: Vec::new(),
269            })?;
270        }
271
272        Ok(())
273    }
274}