lisette_semantics/
store.rs

1use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
2use std::sync::atomic::{AtomicU32, Ordering};
3
4use syntax::ast::{EnumVariant, Expression, StructFieldDefinition};
5use syntax::program::{Definition, File, Interface, MethodSignatures, Module, ModuleId};
6use syntax::types::{SubstitutionMap, Symbol, Type, substitute};
7
8pub const ENTRY_MODULE_ID: &str = "_entry_";
9pub const ENTRY_FILE_ID: u32 = 0;
10
11pub struct Store {
12    pub modules: HashMap<String, Module>,
13    pub module_ids: Vec<ModuleId>,
14    /// file ID -> module ID
15    pub files: HashMap<u32, String>,
16    /// Go module ID -> Go package name, from the typedef `// Package:` directive.
17    /// Present only when the package name differs from the final path segment.
18    pub go_package_names: HashMap<String, String>,
19    visited_modules: HashSet<String>,
20    /// File ID counter. Starts at 2 because 0 is reserved for entry, 1 for prelude.
21    next_file_id: AtomicU32,
22}
23
24impl Default for Store {
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl Store {
31    pub fn new() -> Self {
32        let prelude_module = Module::new("prelude");
33        let nominal_module = Module::nominal();
34
35        let modules = vec![
36            (prelude_module.id.clone(), prelude_module),
37            (nominal_module.id.clone(), nominal_module),
38        ]
39        .into_iter()
40        .collect();
41
42        let module_ids = vec!["prelude".to_string()];
43
44        Self {
45            files: Default::default(),
46            modules,
47            module_ids,
48            go_package_names: Default::default(),
49            visited_modules: Default::default(),
50            next_file_id: AtomicU32::new(2), // 0 = entrypoint, 1 = prelude
51        }
52    }
53
54    pub fn new_file_id(&self) -> u32 {
55        self.next_file_id.fetch_add(1, Ordering::Relaxed)
56    }
57
58    pub fn register_file(&mut self, file_id: u32, module_id: &str) {
59        self.files.insert(file_id, module_id.to_string());
60    }
61
62    pub fn entry_module_id(&self) -> &'static str {
63        ENTRY_MODULE_ID
64    }
65
66    /// Initializes the entry module with reserved file ID 0.
67    pub fn init_entry_module(&mut self) {
68        self.add_module(ENTRY_MODULE_ID);
69        self.register_file(ENTRY_FILE_ID, ENTRY_MODULE_ID);
70    }
71
72    pub fn store_entry_file(&mut self, filename: &str, source: &str, ast: Vec<Expression>) {
73        self.store_file(
74            ENTRY_MODULE_ID,
75            File {
76                id: ENTRY_FILE_ID,
77                module_id: ENTRY_MODULE_ID.to_string(),
78                name: filename.to_string(),
79                source: source.to_string(),
80                items: ast,
81            },
82        );
83    }
84
85    pub fn store_module(&mut self, module_id: &str, files: Vec<File>) {
86        self.mark_visited(module_id);
87        self.add_module(module_id);
88
89        for file in files {
90            self.store_file(module_id, file);
91        }
92    }
93
94    /// Stores a file in the module and registers the file_id -> module_id mapping.
95    /// .d.lis files go to `typedefs`, .lis files go to `files`.
96    pub fn store_file(&mut self, module_id: &str, file: File) {
97        self.files.insert(file.id, module_id.to_string());
98
99        let module = self
100            .get_module_mut(module_id)
101            .expect("module must exist to store file");
102
103        if file.is_d_lis() {
104            module.typedefs.insert(file.id, file);
105        } else {
106            module.files.insert(file.id, file);
107        }
108    }
109
110    pub fn get_file(&self, file_id: u32) -> Option<&File> {
111        let module_id = self.files.get(&file_id)?;
112        let module = self.get_module(module_id)?;
113        module
114            .get_file(file_id)
115            .or_else(|| module.get_typedef_by_id(file_id))
116    }
117
118    pub fn get_file_mut(&mut self, file_id: u32) -> Option<&mut File> {
119        let module_id = self.files.get(&file_id)?.clone();
120        let module = self.modules.get_mut(&module_id)?;
121        module
122            .files
123            .get_mut(&file_id)
124            .or_else(|| module.typedefs.get_mut(&file_id))
125    }
126
127    pub fn get_module(&self, module_id: &str) -> Option<&Module> {
128        self.modules.get(module_id)
129    }
130
131    pub fn has(&self, module_id: &str) -> bool {
132        self.modules.contains_key(module_id)
133    }
134
135    pub fn add_module(&mut self, module_id: &str) {
136        if self.modules.contains_key(module_id) {
137            return;
138        }
139
140        self.modules
141            .insert(module_id.to_string(), Module::new(module_id));
142        self.module_ids.push(module_id.to_string());
143    }
144
145    pub fn get_module_mut(&mut self, module_id: &str) -> Option<&mut Module> {
146        self.modules.get_mut(module_id)
147    }
148
149    pub fn is_visited(&self, module_id: &str) -> bool {
150        self.visited_modules.contains(module_id)
151    }
152
153    pub fn mark_visited(&mut self, module_id: &str) {
154        self.visited_modules.insert(module_id.to_string());
155    }
156
157    pub fn get_definition(&self, qualified_name: &str) -> Option<&Definition> {
158        let module_name = self.module_for_qualified_name(qualified_name)?;
159
160        self.get_module(module_name)?
161            .definitions
162            .get(qualified_name)
163    }
164
165    pub fn module_for_qualified_name<'a>(&'a self, qualified_name: &'a str) -> Option<&'a str> {
166        if !qualified_name.starts_with("go:") || !qualified_name.contains('/') {
167            let (module_name, _) = qualified_name.split_once('.')?;
168            return Some(module_name);
169        }
170
171        let mut best: Option<&str> = None;
172        for module_id in self.modules.keys() {
173            if qualified_name.starts_with(module_id.as_str())
174                && qualified_name.as_bytes().get(module_id.len()) == Some(&b'.')
175                && best
176                    .as_ref()
177                    .is_none_or(|prev| module_id.len() > prev.len())
178            {
179                best = Some(module_id.as_str());
180            }
181        }
182        best
183    }
184
185    pub fn variants_of(&self, qualified_name: &str) -> Option<&[EnumVariant]> {
186        match self.get_definition(qualified_name)? {
187            Definition::Enum { variants, .. } => Some(variants),
188            _ => None,
189        }
190    }
191
192    pub fn variant_of(&self, enum_qualified: &str, variant_name: &str) -> Option<&EnumVariant> {
193        self.variants_of(enum_qualified)?
194            .iter()
195            .find(|v| v.name == variant_name)
196    }
197
198    pub fn value_variants_of(
199        &self,
200        qualified_name: &str,
201    ) -> Option<&[syntax::ast::ValueEnumVariant]> {
202        match self.get_definition(qualified_name)? {
203            Definition::ValueEnum { variants, .. } => Some(variants),
204            _ => None,
205        }
206    }
207
208    pub fn fields_of(&self, qualified_name: &str) -> Option<&[StructFieldDefinition]> {
209        match self.get_definition(qualified_name)? {
210            Definition::Struct { fields, .. } => Some(fields),
211            _ => None,
212        }
213    }
214
215    pub fn struct_kind(&self, qualified_name: &str) -> Option<syntax::ast::StructKind> {
216        match self.get_definition(qualified_name)? {
217            Definition::Struct { kind, .. } => Some(*kind),
218            _ => None,
219        }
220    }
221
222    pub fn struct_constructor(&self, qualified_name: &str) -> Option<&Type> {
223        match self.get_definition(qualified_name)? {
224            Definition::Struct { constructor, .. } => constructor.as_ref(),
225            _ => None,
226        }
227    }
228
229    pub fn parent_interfaces_of(&self, qualified_name: &str) -> Option<&[Type]> {
230        match self.get_definition(qualified_name)? {
231            Definition::Interface { definition, .. } => Some(&definition.parents),
232            _ => None,
233        }
234    }
235
236    pub fn get_type(&self, qualified_name: &str) -> Option<&Type> {
237        self.get_definition(qualified_name)
238            .map(|definition| definition.ty())
239    }
240
241    pub fn get_interface(&self, qualified_name: &str) -> Option<&Interface> {
242        match self.get_definition(qualified_name)? {
243            Definition::Interface { definition, .. } => Some(definition),
244            _ => None,
245        }
246    }
247
248    pub fn peel_alias(&self, ty: &Type) -> Type {
249        let mut current = ty.clone();
250        while let Type::Nominal {
251            id,
252            underlying_ty: Some(u),
253            ..
254        } = &current
255        {
256            if !self
257                .get_definition(id)
258                .is_some_and(|d| matches!(d, Definition::TypeAlias { .. }))
259            {
260                break;
261            }
262            current = *u.clone();
263        }
264        current
265    }
266
267    pub fn get_own_methods(&self, qualified_name: &str) -> Option<&MethodSignatures> {
268        match self.get_definition(qualified_name)? {
269            Definition::Struct { methods, .. } => Some(methods),
270            Definition::TypeAlias { methods, .. } => Some(methods),
271            Definition::Enum { methods, .. } => Some(methods),
272            Definition::ValueEnum { methods, .. } => Some(methods),
273            _ => None,
274        }
275    }
276
277    pub fn get_all_methods(
278        &self,
279        ty: &Type,
280        trait_bounds: &HashMap<Symbol, Vec<Type>>,
281    ) -> MethodSignatures {
282        let stripped = ty.strip_refs();
283        let Some(qualified_name) = method_lookup_key(&stripped) else {
284            return MethodSignatures::default();
285        };
286
287        if let Some(interface) = self.get_interface(&qualified_name) {
288            let mut all_interface_methods = MethodSignatures::default();
289
290            let type_args = ty.get_type_params().unwrap_or_default();
291            let map: SubstitutionMap = interface
292                .generics
293                .iter()
294                .map(|g| g.name.clone())
295                .zip(type_args.iter().cloned())
296                .collect();
297
298            for (name, method_ty) in &interface.methods {
299                let substituted = substitute(method_ty, &map);
300                all_interface_methods.insert(name.clone(), substituted.with_receiver_placeholder());
301            }
302
303            for parent in &interface.parents {
304                for (name, method_ty) in self.get_all_methods(parent, trait_bounds) {
305                    all_interface_methods.insert(name, method_ty);
306                }
307            }
308
309            return all_interface_methods;
310        }
311
312        if let Some(bound_types) = trait_bounds.get(&qualified_name) {
313            return bound_types
314                .iter()
315                .flat_map(|interface_ty| self.get_all_methods(interface_ty, trait_bounds))
316                .collect();
317        }
318
319        let mut methods = self
320            .get_own_methods(&qualified_name)
321            .cloned()
322            .unwrap_or_default();
323
324        // Type aliases inherit methods from the underlying type.
325        if let Some(Definition::TypeAlias { ty: alias_ty, .. }) =
326            self.get_definition(&qualified_name)
327        {
328            let underlying = match &alias_ty {
329                Type::Forall { body, .. } => body.as_ref(),
330                other => other,
331            };
332            let underlying_key = match underlying {
333                Type::Nominal { id, .. } => Some(id.as_str().to_string()),
334                Type::Simple(kind) => Some(format!("prelude.{}", kind.leaf_name())),
335                Type::Compound { kind, .. } => Some(format!("prelude.{}", kind.leaf_name())),
336                _ => None,
337            };
338            // Follow only when the alias body names a different type. For
339            // opaque prelude natives (e.g. `type Map<K, V>`) the body points
340            // to itself — following would loop.
341            if let Some(k) = underlying_key
342                && k != qualified_name.as_str()
343            {
344                let alias_ty = alias_ty.clone();
345                for (name, method_ty) in self.get_all_methods(&alias_ty, trait_bounds) {
346                    methods.entry(name).or_insert(method_ty);
347                }
348            }
349        }
350
351        methods
352    }
353
354    pub fn get_methods_from_bounds(
355        &self,
356        qualified_name: &str,
357        trait_bounds: &HashMap<Symbol, Vec<Type>>,
358    ) -> MethodSignatures {
359        if let Some(bound_types) = trait_bounds.get(qualified_name) {
360            return bound_types
361                .iter()
362                .flat_map(|interface_ty| self.get_all_methods(interface_ty, trait_bounds))
363                .collect();
364        }
365        MethodSignatures::default()
366    }
367}
368
369/// Return the qualified name used to look up methods/fields for a given type.
370/// For `Type::Compound` and `Type::Simple`, this is the prelude-qualified name
371/// (e.g. `Type::Compound { Slice, .. }` → `"prelude.Slice"`).
372fn method_lookup_key(ty: &Type) -> Option<Symbol> {
373    match ty {
374        Type::Nominal { id, .. } => Some(id.clone()),
375        Type::Compound { kind, .. } => Some(Symbol::from_parts("prelude", kind.leaf_name())),
376        Type::Simple(kind) => Some(Symbol::from_parts("prelude", kind.leaf_name())),
377        _ => None,
378    }
379}
lisette_semantics/store.rs

lisette_semantics/
store.rs
