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 is_nominal_defined_type(&self, qualified_name: &str) -> bool {
211 match self.get_definition(qualified_name) {
212 Some(def) => matches!(def.body, DefinitionBody::ValueEnum { .. }) || def.is_newtype(),
213 None => false,
214 }
215 }
216
217 pub fn fields_of(&self, qualified_name: &str) -> Option<&[StructFieldDefinition]> {
218 match &self.get_definition(qualified_name)?.body {
219 DefinitionBody::Struct { fields, .. } => Some(fields),
220 _ => None,
221 }
222 }
223
224 pub fn struct_kind(&self, qualified_name: &str) -> Option<syntax::ast::StructKind> {
225 match &self.get_definition(qualified_name)?.body {
226 DefinitionBody::Struct { kind, .. } => Some(*kind),
227 _ => None,
228 }
229 }
230
231 pub fn struct_constructor(&self, qualified_name: &str) -> Option<&Type> {
232 match &self.get_definition(qualified_name)?.body {
233 DefinitionBody::Struct { constructor, .. } => constructor.as_ref(),
234 _ => None,
235 }
236 }
237
238 pub fn parent_interfaces_of(&self, qualified_name: &str) -> Option<&[Type]> {
239 match &self.get_definition(qualified_name)?.body {
240 DefinitionBody::Interface { definition, .. } => Some(&definition.parents),
241 _ => None,
242 }
243 }
244
245 pub fn get_type(&self, qualified_name: &str) -> Option<&Type> {
246 self.get_definition(qualified_name)
247 .map(|definition| definition.ty())
248 }
249
250 pub fn get_interface(&self, qualified_name: &str) -> Option<&Interface> {
251 match &self.get_definition(qualified_name)?.body {
252 DefinitionBody::Interface { definition, .. } => Some(definition),
253 _ => None,
254 }
255 }
256
257 pub fn is_nilable_go_type(&self, ty: &Type) -> bool {
258 if ty.is_ref() || matches!(ty, Type::Function(_)) {
259 return true;
260 }
261 let Type::Nominal { id, .. } = ty else {
262 return false;
263 };
264 if self.get_definition(id.as_str()).is_none() {
265 return false;
266 }
267 if self.get_interface(id.as_str()).is_some() {
268 return true;
269 }
270 match ty.get_underlying() {
271 Some(Type::Function(_)) => true,
272 Some(u) if u.is_ref() => true,
273 _ => false,
274 }
275 }
276
277 pub fn peel_alias(&self, ty: &Type) -> Type {
278 syntax::types::peel_alias(ty, |id| {
279 self.get_definition(id)
280 .is_some_and(Definition::is_type_alias)
281 })
282 }
283
284 pub fn deep_resolve_alias(&self, ty: &Type) -> Type {
285 let mut current = ty.clone();
286 let mut seen: HashSet<Symbol> = HashSet::default();
287 loop {
288 let Type::Nominal { id, params, .. } = ¤t else {
289 return current;
290 };
291 if !seen.insert(id.clone()) {
292 return current;
293 }
294 let Some(def) = self.get_definition(id.as_str()) else {
295 return current;
296 };
297 if !matches!(def.body, DefinitionBody::TypeAlias { .. }) {
298 return current;
299 }
300 let def_ty = &def.ty;
301 let (vars, body) = match def_ty {
302 Type::Forall { vars, body } => (vars.clone(), body.as_ref().clone()),
303 other => (vec![], other.clone()),
304 };
305 let map: SubstitutionMap = vars.iter().cloned().zip(params.iter().cloned()).collect();
306 current = substitute(&body, &map);
307 }
308 }
309
310 pub fn peel_alias_deep(&self, ty: &Type) -> Type {
311 match self.peel_alias(ty) {
312 Type::Compound { kind, args } => Type::Compound {
313 kind,
314 args: args.iter().map(|a| self.peel_alias_deep(a)).collect(),
315 },
316 Type::Tuple(elements) => {
317 Type::Tuple(elements.iter().map(|e| self.peel_alias_deep(e)).collect())
318 }
319 Type::Nominal {
320 id,
321 params,
322 underlying_ty,
323 } => Type::Nominal {
324 id,
325 params: params.iter().map(|p| self.peel_alias_deep(p)).collect(),
326 underlying_ty,
327 },
328 Type::Function(f) => {
329 let f = *f;
330 Type::function(
331 f.params.iter().map(|p| self.peel_alias_deep(p)).collect(),
332 f.param_mutability,
333 f.bounds,
334 Box::new(self.peel_alias_deep(&f.return_type)),
335 )
336 }
337 other => other,
338 }
339 }
340
341 pub fn get_own_methods(&self, qualified_name: &str) -> Option<&MethodSignatures> {
342 match &self.get_definition(qualified_name)?.body {
343 DefinitionBody::Struct { methods, .. } => Some(methods),
344 DefinitionBody::TypeAlias { methods, .. } => Some(methods),
345 DefinitionBody::Enum { methods, .. } => Some(methods),
346 DefinitionBody::ValueEnum { methods, .. } => Some(methods),
347 _ => None,
348 }
349 }
350
351 pub fn get_all_methods(
352 &self,
353 ty: &Type,
354 trait_bounds: &HashMap<Symbol, Vec<Type>>,
355 ) -> MethodSignatures {
356 let mut visited = HashSet::default();
357 self.get_all_methods_recursive(ty, trait_bounds, &mut visited)
358 }
359
360 fn get_all_methods_recursive(
361 &self,
362 ty: &Type,
363 trait_bounds: &HashMap<Symbol, Vec<Type>>,
364 visited: &mut HashSet<String>,
365 ) -> MethodSignatures {
366 let stripped = ty.strip_refs();
367 let Some(qualified_name) = method_lookup_key(&stripped) else {
368 return MethodSignatures::default();
369 };
370
371 if !visited.insert(qualified_name.as_str().to_string()) {
373 return MethodSignatures::default();
374 }
375
376 if let Some(interface) = self.get_interface(&qualified_name) {
377 let mut all_interface_methods = MethodSignatures::default();
378
379 let type_args = ty.get_type_params().unwrap_or_default();
380 let map: SubstitutionMap = interface
381 .generics
382 .iter()
383 .map(|g| g.name.clone())
384 .zip(type_args.iter().cloned())
385 .collect();
386
387 for (name, method_ty) in &interface.methods {
388 let substituted = substitute(method_ty, &map);
389 all_interface_methods.insert(name.clone(), substituted.with_receiver_placeholder());
390 }
391
392 for parent in &interface.parents {
393 for (name, method_ty) in
394 self.get_all_methods_recursive(parent, trait_bounds, visited)
395 {
396 all_interface_methods.insert(name, method_ty);
397 }
398 }
399
400 return all_interface_methods;
401 }
402
403 if let Some(bound_types) = trait_bounds.get(&qualified_name) {
404 return bound_types
405 .iter()
406 .flat_map(|interface_ty| {
407 self.get_all_methods_recursive(interface_ty, trait_bounds, visited)
408 })
409 .collect();
410 }
411
412 let mut methods = self
413 .get_own_methods(&qualified_name)
414 .cloned()
415 .unwrap_or_default();
416
417 if let Some(definition) = self.get_definition(&qualified_name)
419 && matches!(definition.body, DefinitionBody::TypeAlias { .. })
420 {
421 let alias_ty = &definition.ty;
422 let underlying = match alias_ty {
423 Type::Forall { body, .. } => body.as_ref(),
424 other => other,
425 };
426 let underlying_key = match underlying {
427 Type::Nominal { id, .. } => Some(id.as_str().to_string()),
428 Type::Simple(kind) => Some(format!("prelude.{}", kind.leaf_name())),
429 Type::Compound { kind, .. } => Some(format!("prelude.{}", kind.leaf_name())),
430 _ => None,
431 };
432 if let Some(k) = underlying_key
436 && k != qualified_name.as_str()
437 {
438 let alias_ty = alias_ty.clone();
439 for (name, method_ty) in
440 self.get_all_methods_recursive(&alias_ty, trait_bounds, visited)
441 {
442 methods.entry(name).or_insert(method_ty);
443 }
444 }
445 }
446
447 methods
448 }
449
450 pub fn get_methods_from_bounds(
451 &self,
452 qualified_name: &str,
453 trait_bounds: &HashMap<Symbol, Vec<Type>>,
454 ) -> MethodSignatures {
455 if let Some(bound_types) = trait_bounds.get(qualified_name) {
456 return bound_types
457 .iter()
458 .flat_map(|interface_ty| self.get_all_methods(interface_ty, trait_bounds))
459 .collect();
460 }
461 MethodSignatures::default()
462 }
463}
464
465fn method_lookup_key(ty: &Type) -> Option<Symbol> {
469 match ty {
470 Type::Nominal { id, .. } => Some(id.clone()),
471 Type::Compound { kind, .. } => Some(Symbol::from_parts("prelude", kind.leaf_name())),
472 Type::Simple(kind) => Some(Symbol::from_parts("prelude", kind.leaf_name())),
473 _ => None,
474 }
475}