1use std::path::PathBuf;
2use std::sync::atomic::{AtomicU32, Ordering};
3
4use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
5
6use syntax::ast::{EnumVariant, Expression, StructFieldDefinition};
7use syntax::program::{
8 Definition, DefinitionBody, File, Interface, MethodSignatures, Module, ModuleId,
9};
10use syntax::types::{SubstitutionMap, Symbol, Type, substitute};
11
12pub const ENTRY_MODULE_ID: &str = "_entry_";
13pub const ENTRY_FILE_ID: u32 = 0;
14
15pub struct Store {
16 pub modules: HashMap<String, Module>,
17 pub module_ids: Vec<ModuleId>,
18 pub files: HashMap<u32, String>,
20 pub go_package_names: HashMap<String, String>,
23 pub typedef_paths: HashMap<u32, PathBuf>,
26 visited_modules: HashSet<String>,
27 next_file_id: AtomicU32,
29}
30
31impl Default for Store {
32 fn default() -> Self {
33 Self::new()
34 }
35}
36
37impl Store {
38 pub fn new() -> Self {
39 let prelude_module = Module::new("prelude");
40 let nominal_module = Module::nominal();
41
42 let modules = vec![
43 (prelude_module.id.clone(), prelude_module),
44 (nominal_module.id.clone(), nominal_module),
45 ]
46 .into_iter()
47 .collect();
48
49 let module_ids = vec!["prelude".to_string()];
50
51 Self {
52 files: Default::default(),
53 modules,
54 module_ids,
55 go_package_names: Default::default(),
56 typedef_paths: Default::default(),
57 visited_modules: Default::default(),
58 next_file_id: AtomicU32::new(2), }
60 }
61
62 pub fn new_file_id(&self) -> u32 {
63 self.next_file_id.fetch_add(1, Ordering::Relaxed)
64 }
65
66 pub fn register_file(&mut self, file_id: u32, module_id: &str) {
67 self.files.insert(file_id, module_id.to_string());
68 }
69
70 pub fn entry_module_id(&self) -> &'static str {
71 ENTRY_MODULE_ID
72 }
73
74 pub fn init_entry_module(&mut self) {
76 self.add_module(ENTRY_MODULE_ID);
77 self.register_file(ENTRY_FILE_ID, ENTRY_MODULE_ID);
78 }
79
80 pub fn store_entry_file(
81 &mut self,
82 filename: &str,
83 display_path: &str,
84 source: &str,
85 ast: Vec<Expression>,
86 ) {
87 self.store_file(
88 ENTRY_MODULE_ID,
89 File {
90 id: ENTRY_FILE_ID,
91 module_id: ENTRY_MODULE_ID.to_string(),
92 name: filename.to_string(),
93 display_path: display_path.to_string(),
94 source: source.to_string(),
95 items: ast,
96 },
97 );
98 }
99
100 pub fn store_module(&mut self, module_id: &str, files: Vec<File>) {
101 self.mark_visited(module_id);
102 self.add_module(module_id);
103
104 for file in files {
105 self.store_file(module_id, file);
106 }
107 }
108
109 pub fn store_file(&mut self, module_id: &str, file: File) {
112 self.files.insert(file.id, module_id.to_string());
113
114 let module = self
115 .get_module_mut(module_id)
116 .expect("module must exist to store file");
117
118 if file.is_d_lis() {
119 module.typedefs.insert(file.id, file);
120 } else {
121 module.files.insert(file.id, file);
122 }
123 }
124
125 pub fn get_file(&self, file_id: u32) -> Option<&File> {
126 let module_id = self.files.get(&file_id)?;
127 let module = self.get_module(module_id)?;
128 module
129 .get_file(file_id)
130 .or_else(|| module.get_typedef_by_id(file_id))
131 }
132
133 pub fn get_file_mut(&mut self, file_id: u32) -> Option<&mut File> {
134 let module_id = self.files.get(&file_id)?.clone();
135 let module = self.modules.get_mut(&module_id)?;
136 module
137 .files
138 .get_mut(&file_id)
139 .or_else(|| module.typedefs.get_mut(&file_id))
140 }
141
142 pub fn get_module(&self, module_id: &str) -> Option<&Module> {
143 self.modules.get(module_id)
144 }
145
146 pub fn has(&self, module_id: &str) -> bool {
147 self.modules.contains_key(module_id)
148 }
149
150 pub fn add_module(&mut self, module_id: &str) {
151 if self.modules.contains_key(module_id) {
152 return;
153 }
154
155 self.modules
156 .insert(module_id.to_string(), Module::new(module_id));
157 self.module_ids.push(module_id.to_string());
158 }
159
160 pub fn get_module_mut(&mut self, module_id: &str) -> Option<&mut Module> {
161 self.modules.get_mut(module_id)
162 }
163
164 pub fn is_visited(&self, module_id: &str) -> bool {
165 self.visited_modules.contains(module_id)
166 }
167
168 pub fn mark_visited(&mut self, module_id: &str) {
169 self.visited_modules.insert(module_id.to_string());
170 }
171
172 pub fn get_definition(&self, qualified_name: &str) -> Option<&Definition> {
173 let module_name = self.module_for_qualified_name(qualified_name)?;
174
175 self.get_module(module_name)?
176 .definitions
177 .get(qualified_name)
178 }
179
180 pub fn module_for_qualified_name<'a>(&'a self, qualified_name: &'a str) -> Option<&'a str> {
181 syntax::types::module_for_qualified_name(
182 qualified_name,
183 self.modules.keys().map(String::as_str),
184 )
185 }
186
187 pub fn variants_of(&self, qualified_name: &str) -> Option<&[EnumVariant]> {
188 match &self.get_definition(qualified_name)?.body {
189 DefinitionBody::Enum { variants, .. } => Some(variants),
190 _ => None,
191 }
192 }
193
194 pub fn variant_of(&self, enum_qualified: &str, variant_name: &str) -> Option<&EnumVariant> {
195 self.variants_of(enum_qualified)?
196 .iter()
197 .find(|v| v.name == variant_name)
198 }
199
200 pub fn value_variants_of(
201 &self,
202 qualified_name: &str,
203 ) -> Option<&[syntax::ast::ValueEnumVariant]> {
204 match &self.get_definition(qualified_name)?.body {
205 DefinitionBody::ValueEnum { variants, .. } => Some(variants),
206 _ => None,
207 }
208 }
209
210 pub fn fields_of(&self, qualified_name: &str) -> Option<&[StructFieldDefinition]> {
211 match &self.get_definition(qualified_name)?.body {
212 DefinitionBody::Struct { fields, .. } => Some(fields),
213 _ => None,
214 }
215 }
216
217 pub fn struct_kind(&self, qualified_name: &str) -> Option<syntax::ast::StructKind> {
218 match &self.get_definition(qualified_name)?.body {
219 DefinitionBody::Struct { kind, .. } => Some(*kind),
220 _ => None,
221 }
222 }
223
224 pub fn struct_constructor(&self, qualified_name: &str) -> Option<&Type> {
225 match &self.get_definition(qualified_name)?.body {
226 DefinitionBody::Struct { constructor, .. } => constructor.as_ref(),
227 _ => None,
228 }
229 }
230
231 pub fn parent_interfaces_of(&self, qualified_name: &str) -> Option<&[Type]> {
232 match &self.get_definition(qualified_name)?.body {
233 DefinitionBody::Interface { definition, .. } => Some(&definition.parents),
234 _ => None,
235 }
236 }
237
238 pub fn get_type(&self, qualified_name: &str) -> Option<&Type> {
239 self.get_definition(qualified_name)
240 .map(|definition| definition.ty())
241 }
242
243 pub fn get_interface(&self, qualified_name: &str) -> Option<&Interface> {
244 match &self.get_definition(qualified_name)?.body {
245 DefinitionBody::Interface { definition, .. } => Some(definition),
246 _ => None,
247 }
248 }
249
250 pub fn is_nilable_go_type(&self, ty: &Type) -> bool {
251 if ty.is_ref() || matches!(ty, Type::Function(_)) {
252 return true;
253 }
254 let Type::Nominal { id, .. } = ty else {
255 return false;
256 };
257 if self.get_definition(id.as_str()).is_none() {
258 return false;
259 }
260 if self.get_interface(id.as_str()).is_some() {
261 return true;
262 }
263 match ty.get_underlying() {
264 Some(Type::Function(_)) => true,
265 Some(u) if u.is_ref() => true,
266 _ => false,
267 }
268 }
269
270 pub fn peel_alias(&self, ty: &Type) -> Type {
271 syntax::types::peel_alias(ty, |id| {
272 self.get_definition(id)
273 .is_some_and(Definition::is_type_alias)
274 })
275 }
276
277 pub fn deep_resolve_alias(&self, ty: &Type) -> Type {
278 let mut current = ty.clone();
279 let mut seen: HashSet<Symbol> = HashSet::default();
280 loop {
281 let Type::Nominal { id, params, .. } = ¤t else {
282 return current;
283 };
284 if !seen.insert(id.clone()) {
285 return current;
286 }
287 let Some(def) = self.get_definition(id.as_str()) else {
288 return current;
289 };
290 if !matches!(def.body, DefinitionBody::TypeAlias { .. }) {
291 return current;
292 }
293 let def_ty = &def.ty;
294 let (vars, body) = match def_ty {
295 Type::Forall { vars, body } => (vars.clone(), body.as_ref().clone()),
296 other => (vec![], other.clone()),
297 };
298 let map: SubstitutionMap = vars.iter().cloned().zip(params.iter().cloned()).collect();
299 current = substitute(&body, &map);
300 }
301 }
302
303 pub fn peel_alias_deep(&self, ty: &Type) -> Type {
304 match self.peel_alias(ty) {
305 Type::Compound { kind, args } => Type::Compound {
306 kind,
307 args: args.iter().map(|a| self.peel_alias_deep(a)).collect(),
308 },
309 Type::Tuple(elements) => {
310 Type::Tuple(elements.iter().map(|e| self.peel_alias_deep(e)).collect())
311 }
312 Type::Nominal {
313 id,
314 params,
315 underlying_ty,
316 } => Type::Nominal {
317 id,
318 params: params.iter().map(|p| self.peel_alias_deep(p)).collect(),
319 underlying_ty,
320 },
321 Type::Function(f) => {
322 let f = *f;
323 Type::function(
324 f.params.iter().map(|p| self.peel_alias_deep(p)).collect(),
325 f.param_mutability,
326 f.bounds,
327 Box::new(self.peel_alias_deep(&f.return_type)),
328 )
329 }
330 other => other,
331 }
332 }
333
334 pub fn get_own_methods(&self, qualified_name: &str) -> Option<&MethodSignatures> {
335 match &self.get_definition(qualified_name)?.body {
336 DefinitionBody::Struct { methods, .. } => Some(methods),
337 DefinitionBody::TypeAlias { methods, .. } => Some(methods),
338 DefinitionBody::Enum { methods, .. } => Some(methods),
339 DefinitionBody::ValueEnum { methods, .. } => Some(methods),
340 _ => None,
341 }
342 }
343
344 pub fn get_all_methods(
345 &self,
346 ty: &Type,
347 trait_bounds: &HashMap<Symbol, Vec<Type>>,
348 ) -> MethodSignatures {
349 let mut visited = HashSet::default();
350 self.get_all_methods_recursive(ty, trait_bounds, &mut visited)
351 }
352
353 fn get_all_methods_recursive(
354 &self,
355 ty: &Type,
356 trait_bounds: &HashMap<Symbol, Vec<Type>>,
357 visited: &mut HashSet<String>,
358 ) -> MethodSignatures {
359 let stripped = ty.strip_refs();
360 let Some(qualified_name) = method_lookup_key(&stripped) else {
361 return MethodSignatures::default();
362 };
363
364 if !visited.insert(qualified_name.as_str().to_string()) {
366 return MethodSignatures::default();
367 }
368
369 if let Some(interface) = self.get_interface(&qualified_name) {
370 let mut all_interface_methods = MethodSignatures::default();
371
372 let type_args = ty.get_type_params().unwrap_or_default();
373 let map: SubstitutionMap = interface
374 .generics
375 .iter()
376 .map(|g| g.name.clone())
377 .zip(type_args.iter().cloned())
378 .collect();
379
380 for (name, method_ty) in &interface.methods {
381 let substituted = substitute(method_ty, &map);
382 all_interface_methods.insert(name.clone(), substituted.with_receiver_placeholder());
383 }
384
385 for parent in &interface.parents {
386 for (name, method_ty) in
387 self.get_all_methods_recursive(parent, trait_bounds, visited)
388 {
389 all_interface_methods.insert(name, method_ty);
390 }
391 }
392
393 return all_interface_methods;
394 }
395
396 if let Some(bound_types) = trait_bounds.get(&qualified_name) {
397 return bound_types
398 .iter()
399 .flat_map(|interface_ty| {
400 self.get_all_methods_recursive(interface_ty, trait_bounds, visited)
401 })
402 .collect();
403 }
404
405 let mut methods = self
406 .get_own_methods(&qualified_name)
407 .cloned()
408 .unwrap_or_default();
409
410 if let Some(definition) = self.get_definition(&qualified_name)
412 && matches!(definition.body, DefinitionBody::TypeAlias { .. })
413 {
414 let alias_ty = &definition.ty;
415 let underlying = match alias_ty {
416 Type::Forall { body, .. } => body.as_ref(),
417 other => other,
418 };
419 let underlying_key = match underlying {
420 Type::Nominal { id, .. } => Some(id.as_str().to_string()),
421 Type::Simple(kind) => Some(format!("prelude.{}", kind.leaf_name())),
422 Type::Compound { kind, .. } => Some(format!("prelude.{}", kind.leaf_name())),
423 _ => None,
424 };
425 if let Some(k) = underlying_key
429 && k != qualified_name.as_str()
430 {
431 let alias_ty = alias_ty.clone();
432 for (name, method_ty) in
433 self.get_all_methods_recursive(&alias_ty, trait_bounds, visited)
434 {
435 methods.entry(name).or_insert(method_ty);
436 }
437 }
438 }
439
440 methods
441 }
442
443 pub fn get_methods_from_bounds(
444 &self,
445 qualified_name: &str,
446 trait_bounds: &HashMap<Symbol, Vec<Type>>,
447 ) -> MethodSignatures {
448 if let Some(bound_types) = trait_bounds.get(qualified_name) {
449 return bound_types
450 .iter()
451 .flat_map(|interface_ty| self.get_all_methods(interface_ty, trait_bounds))
452 .collect();
453 }
454 MethodSignatures::default()
455 }
456}
457
458fn method_lookup_key(ty: &Type) -> Option<Symbol> {
462 match ty {
463 Type::Nominal { id, .. } => Some(id.clone()),
464 Type::Compound { kind, .. } => Some(Symbol::from_parts("prelude", kind.leaf_name())),
465 Type::Simple(kind) => Some(Symbol::from_parts("prelude", kind.leaf_name())),
466 _ => None,
467 }
468}