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 {
322 params,
323 param_mutability,
324 bounds,
325 return_type,
326 } => Type::Function {
327 params: params.iter().map(|p| self.peel_alias_deep(p)).collect(),
328 param_mutability,
329 bounds,
330 return_type: Box::new(self.peel_alias_deep(&return_type)),
331 },
332 other => other,
333 }
334 }
335
336 pub fn get_own_methods(&self, qualified_name: &str) -> Option<&MethodSignatures> {
337 match &self.get_definition(qualified_name)?.body {
338 DefinitionBody::Struct { methods, .. } => Some(methods),
339 DefinitionBody::TypeAlias { methods, .. } => Some(methods),
340 DefinitionBody::Enum { methods, .. } => Some(methods),
341 DefinitionBody::ValueEnum { methods, .. } => Some(methods),
342 _ => None,
343 }
344 }
345
346 pub fn get_all_methods(
347 &self,
348 ty: &Type,
349 trait_bounds: &HashMap<Symbol, Vec<Type>>,
350 ) -> MethodSignatures {
351 let stripped = ty.strip_refs();
352 let Some(qualified_name) = method_lookup_key(&stripped) else {
353 return MethodSignatures::default();
354 };
355
356 if let Some(interface) = self.get_interface(&qualified_name) {
357 let mut all_interface_methods = MethodSignatures::default();
358
359 let type_args = ty.get_type_params().unwrap_or_default();
360 let map: SubstitutionMap = interface
361 .generics
362 .iter()
363 .map(|g| g.name.clone())
364 .zip(type_args.iter().cloned())
365 .collect();
366
367 for (name, method_ty) in &interface.methods {
368 let substituted = substitute(method_ty, &map);
369 all_interface_methods.insert(name.clone(), substituted.with_receiver_placeholder());
370 }
371
372 for parent in &interface.parents {
373 for (name, method_ty) in self.get_all_methods(parent, trait_bounds) {
374 all_interface_methods.insert(name, method_ty);
375 }
376 }
377
378 return all_interface_methods;
379 }
380
381 if let Some(bound_types) = trait_bounds.get(&qualified_name) {
382 return bound_types
383 .iter()
384 .flat_map(|interface_ty| self.get_all_methods(interface_ty, trait_bounds))
385 .collect();
386 }
387
388 let mut methods = self
389 .get_own_methods(&qualified_name)
390 .cloned()
391 .unwrap_or_default();
392
393 if let Some(definition) = self.get_definition(&qualified_name)
395 && matches!(definition.body, DefinitionBody::TypeAlias { .. })
396 {
397 let alias_ty = &definition.ty;
398 let underlying = match alias_ty {
399 Type::Forall { body, .. } => body.as_ref(),
400 other => other,
401 };
402 let underlying_key = match underlying {
403 Type::Nominal { id, .. } => Some(id.as_str().to_string()),
404 Type::Simple(kind) => Some(format!("prelude.{}", kind.leaf_name())),
405 Type::Compound { kind, .. } => Some(format!("prelude.{}", kind.leaf_name())),
406 _ => None,
407 };
408 if let Some(k) = underlying_key
412 && k != qualified_name.as_str()
413 {
414 let alias_ty = alias_ty.clone();
415 for (name, method_ty) in self.get_all_methods(&alias_ty, trait_bounds) {
416 methods.entry(name).or_insert(method_ty);
417 }
418 }
419 }
420
421 methods
422 }
423
424 pub fn get_methods_from_bounds(
425 &self,
426 qualified_name: &str,
427 trait_bounds: &HashMap<Symbol, Vec<Type>>,
428 ) -> MethodSignatures {
429 if let Some(bound_types) = trait_bounds.get(qualified_name) {
430 return bound_types
431 .iter()
432 .flat_map(|interface_ty| self.get_all_methods(interface_ty, trait_bounds))
433 .collect();
434 }
435 MethodSignatures::default()
436 }
437}
438
439fn method_lookup_key(ty: &Type) -> Option<Symbol> {
443 match ty {
444 Type::Nominal { id, .. } => Some(id.clone()),
445 Type::Compound { kind, .. } => Some(Symbol::from_parts("prelude", kind.leaf_name())),
446 Type::Simple(kind) => Some(Symbol::from_parts("prelude", kind.leaf_name())),
447 _ => None,
448 }
449}