Skip to main content

lisette_semantics/
store.rs

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    /// file ID -> module ID
19    pub files: HashMap<u32, String>,
20    /// Go module ID -> Go package name, from the typedef `// Package:` directive.
21    /// Present only when the package name differs from the final path segment.
22    pub go_package_names: HashMap<String, String>,
23    /// File ID -> on-disk path of the `.d.lis` typedef. Lets the LSP map go: typedef
24    /// file IDs to the actual cache path so go-to-definition can navigate there.
25    pub typedef_paths: HashMap<u32, PathBuf>,
26    visited_modules: HashSet<String>,
27    /// File ID counter. Starts at 2 because 0 is reserved for entry, 1 for prelude.
28    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), // 0 = entrypoint, 1 = prelude
59        }
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    /// Initializes the entry module with reserved file ID 0.
75    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    /// Stores a file in the module and registers the file_id -> module_id mapping.
110    /// .d.lis files go to `typedefs`, .lis files go to `files`.
111    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, .. } = &current 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        // Cyclic embeddings survive registration as an error with parents intact; guard the walk.
372        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        // Type aliases inherit methods from the underlying type.
418        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            // Follow only when the alias body names a different type. For
433            // opaque prelude natives (e.g. `type Map<K, V>`) the body points
434            // to itself — following would loop.
435            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
465/// Return the qualified name used to look up methods/fields for a given type.
466/// For `Type::Compound` and `Type::Simple`, this is the prelude-qualified name
467/// (e.g. `Type::Compound { Slice, .. }` → `"prelude.Slice"`).
468fn 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}