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 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, .. } = &current 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        // Type aliases inherit methods from the underlying type.
394        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            // Follow only when the alias body names a different type. For
409            // opaque prelude natives (e.g. `type Map<K, V>`) the body points
410            // to itself — following would loop.
411            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
439/// Return the qualified name used to look up methods/fields for a given type.
440/// For `Type::Compound` and `Type::Simple`, this is the prelude-qualified name
441/// (e.g. `Type::Compound { Slice, .. }` → `"prelude.Slice"`).
442fn 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}