1use std::path::PathBuf;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU32, Ordering};
4
5use ecow::EcoString;
6use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
7
8use syntax::ast::{EnumVariant, Expression, Literal, StructFieldDefinition};
9use syntax::program::{
10 Definition, DefinitionBody, File, Interface, MethodSignatures, Module, ModuleId,
11};
12use syntax::types::{SimpleKind, SubstitutionMap, Symbol, Type, substitute};
13
14pub const ENTRY_MODULE_ID: &str = "_entry_";
15pub const ENTRY_FILE_ID: u32 = 0;
16
17#[derive(Debug, Clone)]
18pub struct ClosedMember {
19 pub display_name: EcoString,
21 pub literal: Literal,
23 pub value: DomainValue,
25}
26
27#[derive(Debug, Clone)]
29pub struct ClosedDomain {
30 pub base: SimpleKind,
31 pub type_display: EcoString,
32 pub members: Vec<ClosedMember>,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
38pub enum DomainValue {
39 Int(i128),
40 Str(String),
41}
42
43impl DomainValue {
44 pub fn from_literal(literal: &Literal, base: SimpleKind) -> Option<DomainValue> {
45 match base {
49 SimpleKind::Rune => match literal {
50 Literal::Char(text) => char_codepoint(text).map(|cp| DomainValue::Int(cp as i128)),
51 Literal::Integer { value, .. } => Some(DomainValue::Int(*value as i64 as i128)),
52 _ => None,
53 },
54 SimpleKind::String => match literal {
55 Literal::String { value, .. } => Some(DomainValue::Str(value.clone())),
56 _ => None,
57 },
58 _ if is_unsigned_base(base) => match literal {
59 Literal::Integer { value, .. } => Some(DomainValue::Int(*value as i128)),
60 _ => None,
61 },
62 _ if base.is_signed_int() => match literal {
63 Literal::Integer { value, .. } => Some(DomainValue::Int(*value as i64 as i128)),
64 _ => None,
65 },
66 _ => None,
67 }
68 }
69}
70
71pub fn is_unsigned_base(base: SimpleKind) -> bool {
74 base.is_unsigned_int() || base == SimpleKind::Uintptr
75}
76
77fn char_codepoint(text: &str) -> Option<u64> {
80 let Some(rest) = text.strip_prefix('\\') else {
81 return text.chars().next().map(|c| c as u64);
82 };
83 match rest.as_bytes().first()? {
84 b'a' => Some(7),
85 b'b' => Some(8),
86 b'f' => Some(12),
87 b'n' => Some(10),
88 b'r' => Some(13),
89 b't' => Some(9),
90 b'v' => Some(11),
91 b'\\' => Some(92),
92 b'\'' => Some(39),
93 b'x' => u64::from_str_radix(&rest[1..], 16).ok(),
94 b'0'..=b'7' => u64::from_str_radix(rest, 8).ok(),
95 _ => None,
96 }
97}
98
99pub struct Store {
100 pub modules: HashMap<String, Arc<Module>>,
103 pub module_ids: Vec<ModuleId>,
104 pub files: HashMap<u32, String>,
106 pub go_package_names: HashMap<String, String>,
109 pub typedef_paths: HashMap<u32, PathBuf>,
112 visited_modules: HashSet<String>,
113 next_file_id: AtomicU32,
115 pub closed_domains: HashMap<Symbol, ClosedDomain>,
118}
119
120impl Default for Store {
121 fn default() -> Self {
122 Self::new()
123 }
124}
125
126impl Store {
127 pub fn new() -> Self {
128 let prelude_module = Module::new("prelude");
129 let nominal_module = Module::nominal();
130
131 let modules = vec![
132 (prelude_module.id.clone(), Arc::new(prelude_module)),
133 (nominal_module.id.clone(), Arc::new(nominal_module)),
134 ]
135 .into_iter()
136 .collect();
137
138 let module_ids = vec!["prelude".to_string()];
139
140 Self {
141 files: Default::default(),
142 modules,
143 module_ids,
144 go_package_names: Default::default(),
145 typedef_paths: Default::default(),
146 visited_modules: Default::default(),
147 next_file_id: AtomicU32::new(2), closed_domains: Default::default(),
149 }
150 }
151
152 pub fn new_file_id(&self) -> u32 {
153 self.next_file_id.fetch_add(1, Ordering::Relaxed)
154 }
155
156 pub fn register_file(&mut self, file_id: u32, module_id: &str) {
157 self.files.insert(file_id, module_id.to_string());
158 }
159
160 pub fn entry_module_id(&self) -> &'static str {
161 ENTRY_MODULE_ID
162 }
163
164 pub fn init_entry_module(&mut self) {
166 self.add_module(ENTRY_MODULE_ID);
167 self.register_file(ENTRY_FILE_ID, ENTRY_MODULE_ID);
168 }
169
170 pub fn store_entry_file(
171 &mut self,
172 filename: &str,
173 display_path: &str,
174 source: &str,
175 ast: Vec<Expression>,
176 ) {
177 self.store_file(
178 ENTRY_MODULE_ID,
179 File {
180 id: ENTRY_FILE_ID,
181 module_id: ENTRY_MODULE_ID.to_string(),
182 name: filename.to_string(),
183 display_path: display_path.to_string(),
184 source: source.to_string(),
185 items: ast,
186 },
187 );
188 }
189
190 pub fn store_module(&mut self, module_id: &str, files: Vec<File>) {
191 self.mark_visited(module_id);
192 self.add_module(module_id);
193
194 for file in files {
195 self.store_file(module_id, file);
196 }
197 }
198
199 pub fn store_file(&mut self, module_id: &str, file: File) {
202 self.files.insert(file.id, module_id.to_string());
203
204 let module = self
205 .get_module_mut(module_id)
206 .expect("module must exist to store file");
207
208 if file.is_d_lis() {
209 module.typedefs.insert(file.id, file);
210 } else {
211 module.files.insert(file.id, file);
212 }
213 }
214
215 pub fn get_file(&self, file_id: u32) -> Option<&File> {
216 let module_id = self.files.get(&file_id)?;
217 let module = self.get_module(module_id)?;
218 module
219 .get_file(file_id)
220 .or_else(|| module.get_typedef_by_id(file_id))
221 }
222
223 pub fn get_file_mut(&mut self, file_id: u32) -> Option<&mut File> {
224 let module_id = self.files.get(&file_id)?.clone();
225 let module = Arc::make_mut(self.modules.get_mut(&module_id)?);
226 module
227 .files
228 .get_mut(&file_id)
229 .or_else(|| module.typedefs.get_mut(&file_id))
230 }
231
232 pub fn get_module(&self, module_id: &str) -> Option<&Module> {
233 self.modules.get(module_id).map(Arc::as_ref)
234 }
235
236 pub fn has(&self, module_id: &str) -> bool {
237 self.modules.contains_key(module_id)
238 }
239
240 pub fn add_module(&mut self, module_id: &str) {
241 if self.modules.contains_key(module_id) {
242 return;
243 }
244
245 self.modules
246 .insert(module_id.to_string(), Arc::new(Module::new(module_id)));
247 self.module_ids.push(module_id.to_string());
248 }
249
250 pub fn get_module_mut(&mut self, module_id: &str) -> Option<&mut Module> {
251 self.modules.get_mut(module_id).map(Arc::make_mut)
252 }
253
254 pub(crate) fn registration_view(&self) -> Store {
257 Store {
258 modules: self.modules.clone(),
259 module_ids: self.module_ids.clone(),
260 files: self.files.clone(),
261 go_package_names: self.go_package_names.clone(),
262 typedef_paths: HashMap::default(),
263 visited_modules: HashSet::default(),
264 next_file_id: AtomicU32::new(self.next_file_id.load(Ordering::Relaxed)),
265 closed_domains: HashMap::default(),
266 }
267 }
268
269 pub fn is_visited(&self, module_id: &str) -> bool {
270 self.visited_modules.contains(module_id)
271 }
272
273 pub fn mark_visited(&mut self, module_id: &str) {
274 self.visited_modules.insert(module_id.to_string());
275 }
276
277 pub fn get_definition(&self, qualified_name: &str) -> Option<&Definition> {
278 let module_name = self.module_for_qualified_name(qualified_name)?;
279
280 self.get_module(module_name)?
281 .definitions
282 .get(qualified_name)
283 }
284
285 pub fn module_for_qualified_name<'a>(&'a self, qualified_name: &'a str) -> Option<&'a str> {
286 syntax::types::module_for_qualified_name(
287 qualified_name,
288 self.modules.keys().map(String::as_str),
289 )
290 }
291
292 pub fn variants_of(&self, qualified_name: &str) -> Option<&[EnumVariant]> {
293 match &self.get_definition(qualified_name)?.body {
294 DefinitionBody::Enum { variants, .. } => Some(variants),
295 _ => None,
296 }
297 }
298
299 pub fn variant_of(&self, enum_qualified: &str, variant_name: &str) -> Option<&EnumVariant> {
300 self.variants_of(enum_qualified)?
301 .iter()
302 .find(|v| v.name == variant_name)
303 }
304
305 pub fn is_nominal_defined_type(&self, qualified_name: &str) -> bool {
306 match self.get_definition(qualified_name) {
307 Some(def) => def.is_newtype(),
308 None => false,
309 }
310 }
311
312 pub fn build_closed_domains(&mut self) {
313 let mut bases: HashMap<Symbol, (SimpleKind, String)> = HashMap::default();
315 for module in self.modules.values() {
316 for (qualified_name, definition) in &module.definitions {
317 if definition.is_closed_domain()
320 && let Some(base) = definition.ty().underlying_simple_kind()
321 && !base.is_float()
322 {
323 bases.insert(qualified_name.clone(), (base, module.id.clone()));
324 }
325 }
326 }
327
328 if bases.is_empty() {
329 return;
330 }
331
332 let mut members: HashMap<Symbol, Vec<ClosedMember>> = HashMap::default();
333 for module in self.modules.values() {
334 for (qualified_name, definition) in &module.definitions {
335 let Some(const_literal) = definition.const_value() else {
336 continue;
337 };
338 let Type::Nominal { id, .. } = definition.ty() else {
339 continue;
340 };
341 let Some((base, declaring_module)) = bases.get(id) else {
342 continue;
343 };
344 if module.id != *declaring_module {
347 continue;
348 }
349 let Some(value) = DomainValue::from_literal(const_literal, *base) else {
350 continue;
351 };
352 members.entry(id.clone()).or_default().push(ClosedMember {
353 display_name: domain_display_name(qualified_name.as_str()).into(),
354 literal: const_literal.clone(),
355 value,
356 });
357 }
358 }
359
360 let mut domains: HashMap<Symbol, ClosedDomain> = HashMap::default();
361 for (type_id, (base, _)) in bases {
362 let Some(mut domain_members) = members.remove(&type_id) else {
363 continue;
364 };
365 domain_members.sort_by(|a, b| a.value.cmp(&b.value));
366 domains.insert(
367 type_id.clone(),
368 ClosedDomain {
369 base,
370 type_display: domain_display_name(type_id.as_str()).into(),
371 members: domain_members,
372 },
373 );
374 }
375
376 self.closed_domains = domains;
377 }
378
379 pub fn fields_of(&self, qualified_name: &str) -> Option<&[StructFieldDefinition]> {
380 match &self.get_definition(qualified_name)?.body {
381 DefinitionBody::Struct { fields, .. } => Some(fields),
382 _ => None,
383 }
384 }
385
386 pub fn struct_kind(&self, qualified_name: &str) -> Option<syntax::ast::StructKind> {
387 match &self.get_definition(qualified_name)?.body {
388 DefinitionBody::Struct { kind, .. } => Some(*kind),
389 _ => None,
390 }
391 }
392
393 pub fn struct_constructor(&self, qualified_name: &str) -> Option<&Type> {
394 match &self.get_definition(qualified_name)?.body {
395 DefinitionBody::Struct { constructor, .. } => constructor.as_ref(),
396 _ => None,
397 }
398 }
399
400 pub fn parent_interfaces_of(&self, qualified_name: &str) -> Option<&[Type]> {
401 match &self.get_definition(qualified_name)?.body {
402 DefinitionBody::Interface { definition, .. } => Some(&definition.parents),
403 _ => None,
404 }
405 }
406
407 pub fn get_type(&self, qualified_name: &str) -> Option<&Type> {
408 self.get_definition(qualified_name)
409 .map(|definition| definition.ty())
410 }
411
412 pub fn get_interface(&self, qualified_name: &str) -> Option<&Interface> {
413 match &self.get_definition(qualified_name)?.body {
414 DefinitionBody::Interface { definition, .. } => Some(definition),
415 _ => None,
416 }
417 }
418
419 pub fn is_interface(&self, ty: &Type) -> bool {
420 matches!(ty, Type::Nominal { id, .. } if self.get_interface(id.as_str()).is_some())
421 }
422
423 pub fn is_nilable_go_type(&self, ty: &Type) -> bool {
424 if ty.is_ref() || matches!(ty, Type::Function(_)) {
425 return true;
426 }
427 let Type::Nominal { id, .. } = ty else {
428 return false;
429 };
430 if self.get_definition(id.as_str()).is_none() {
431 return false;
432 }
433 if self.get_interface(id.as_str()).is_some() {
434 return true;
435 }
436 match ty.get_underlying() {
437 Some(Type::Function(_)) => true,
438 Some(u) if u.is_ref() => true,
439 _ => false,
440 }
441 }
442
443 pub fn peel_alias(&self, ty: &Type) -> Type {
444 syntax::types::peel_alias(ty, |id| {
445 self.get_definition(id)
446 .is_some_and(Definition::is_type_alias)
447 })
448 }
449
450 pub fn deep_resolve_alias(&self, ty: &Type) -> Type {
451 let mut current = ty.clone();
452 let mut seen: HashSet<Symbol> = HashSet::default();
453 loop {
454 let Type::Nominal { id, params, .. } = ¤t else {
455 return current;
456 };
457 if !seen.insert(id.clone()) {
458 return current;
459 }
460 let Some(def) = self.get_definition(id.as_str()) else {
461 return current;
462 };
463 if !matches!(def.body, DefinitionBody::TypeAlias { .. }) {
464 return current;
465 }
466 let def_ty = &def.ty;
467 let (vars, body) = match def_ty {
468 Type::Forall { vars, body } => (vars.clone(), body.as_ref().clone()),
469 other => (vec![], other.clone()),
470 };
471 let map: SubstitutionMap = vars.iter().cloned().zip(params.iter().cloned()).collect();
472 current = substitute(&body, &map);
473 }
474 }
475
476 pub fn peel_alias_deep(&self, ty: &Type) -> Type {
477 match self.peel_alias(ty) {
478 Type::Compound { kind, args } => Type::Compound {
479 kind,
480 args: args.iter().map(|a| self.peel_alias_deep(a)).collect(),
481 },
482 Type::Tuple(elements) => {
483 Type::Tuple(elements.iter().map(|e| self.peel_alias_deep(e)).collect())
484 }
485 Type::Nominal {
486 id,
487 params,
488 underlying_ty,
489 } => Type::Nominal {
490 id,
491 params: params.iter().map(|p| self.peel_alias_deep(p)).collect(),
492 underlying_ty,
493 },
494 Type::Function(f) => {
495 let f = std::sync::Arc::try_unwrap(f).unwrap_or_else(|arc| (*arc).clone());
496 Type::function(
497 f.params.iter().map(|p| self.peel_alias_deep(p)).collect(),
498 f.param_mutability,
499 f.bounds,
500 Box::new(self.peel_alias_deep(&f.return_type)),
501 )
502 }
503 other => other,
504 }
505 }
506
507 pub fn get_own_methods(&self, qualified_name: &str) -> Option<&MethodSignatures> {
508 match &self.get_definition(qualified_name)?.body {
509 DefinitionBody::Struct { methods, .. } => Some(methods),
510 DefinitionBody::TypeAlias { methods, .. } => Some(methods),
511 DefinitionBody::Enum { methods, .. } => Some(methods),
512 _ => None,
513 }
514 }
515
516 pub fn get_all_methods(
517 &self,
518 ty: &Type,
519 trait_bounds: &HashMap<Symbol, Vec<Type>>,
520 ) -> MethodSignatures {
521 let mut visited = HashSet::default();
522 self.get_all_methods_recursive(ty, trait_bounds, &mut visited)
523 }
524
525 fn get_all_methods_recursive(
526 &self,
527 ty: &Type,
528 trait_bounds: &HashMap<Symbol, Vec<Type>>,
529 visited: &mut HashSet<String>,
530 ) -> MethodSignatures {
531 let stripped = ty.strip_refs();
532 let Some(qualified_name) = method_lookup_key(&stripped) else {
533 return MethodSignatures::default();
534 };
535
536 if !visited.insert(qualified_name.as_str().to_string()) {
538 return MethodSignatures::default();
539 }
540
541 if let Some(interface) = self.get_interface(&qualified_name) {
542 let mut all_interface_methods = MethodSignatures::default();
543
544 let type_args = ty.get_type_params().unwrap_or_default();
545 let map: SubstitutionMap = interface
546 .generics
547 .iter()
548 .map(|g| g.name.clone())
549 .zip(type_args.iter().cloned())
550 .collect();
551
552 for (name, method_ty) in &interface.methods {
553 let substituted = substitute(method_ty, &map);
554 all_interface_methods.insert(name.clone(), substituted.with_receiver_placeholder());
555 }
556
557 for parent in &interface.parents {
558 for (name, method_ty) in
559 self.get_all_methods_recursive(parent, trait_bounds, visited)
560 {
561 all_interface_methods.insert(name, method_ty);
562 }
563 }
564
565 return all_interface_methods;
566 }
567
568 if let Some(bound_types) = trait_bounds.get(&qualified_name) {
569 return bound_types
570 .iter()
571 .flat_map(|interface_ty| {
572 self.get_all_methods_recursive(interface_ty, trait_bounds, visited)
573 })
574 .collect();
575 }
576
577 let mut methods = self
578 .get_own_methods(&qualified_name)
579 .cloned()
580 .unwrap_or_default();
581
582 if let Some(definition) = self.get_definition(&qualified_name)
584 && matches!(definition.body, DefinitionBody::TypeAlias { .. })
585 {
586 let alias_ty = &definition.ty;
587 let underlying = match alias_ty {
588 Type::Forall { body, .. } => body.as_ref(),
589 other => other,
590 };
591 let underlying_key = match underlying {
592 Type::Nominal { id, .. } => Some(id.as_str().to_string()),
593 Type::Simple(kind) => Some(format!("prelude.{}", kind.leaf_name())),
594 Type::Compound { kind, .. } => Some(format!("prelude.{}", kind.leaf_name())),
595 _ => None,
596 };
597 if let Some(k) = underlying_key
601 && k != qualified_name.as_str()
602 {
603 let alias_ty = alias_ty.clone();
604 for (name, method_ty) in
605 self.get_all_methods_recursive(&alias_ty, trait_bounds, visited)
606 {
607 methods.entry(name).or_insert(method_ty);
608 }
609 }
610 }
611
612 methods
613 }
614
615 pub fn get_methods_from_bounds(
616 &self,
617 qualified_name: &str,
618 trait_bounds: &HashMap<Symbol, Vec<Type>>,
619 ) -> MethodSignatures {
620 if let Some(bound_types) = trait_bounds.get(qualified_name) {
621 return bound_types
622 .iter()
623 .flat_map(|interface_ty| self.get_all_methods(interface_ty, trait_bounds))
624 .collect();
625 }
626 MethodSignatures::default()
627 }
628}
629
630fn domain_display_name(qualified: &str) -> String {
631 let Some((module, name)) = qualified.rsplit_once('.') else {
632 return qualified.to_string();
633 };
634 match module.strip_prefix("go:") {
635 Some(go_module) => {
636 let package = go_module.rsplit('/').next().unwrap_or(go_module);
637 format!("{package}.{name}")
638 }
639 None => name.to_string(),
640 }
641}
642
643fn method_lookup_key(ty: &Type) -> Option<Symbol> {
647 match ty {
648 Type::Nominal { id, .. } => Some(id.clone()),
649 Type::Compound { kind, .. } => Some(Symbol::from_parts("prelude", kind.leaf_name())),
650 Type::Simple(kind) => Some(Symbol::from_parts("prelude", kind.leaf_name())),
651 _ => None,
652 }
653}
654
655#[cfg(test)]
656mod closed_domain_tests {
657 use super::*;
658 use syntax::ast::StructKind;
659 use syntax::program::{Attributes, TypeAttribute, Visibility};
660
661 fn nominal_int(id: &str) -> Type {
662 Type::Nominal {
663 id: Symbol::from_raw(id),
664 params: vec![],
665 underlying_ty: Some(Box::new(Type::Simple(SimpleKind::Int))),
666 }
667 }
668
669 fn struct_def(ty: Type, closed_domain: bool) -> Definition {
670 let mut attributes = Attributes::default();
671 if closed_domain {
672 attributes.insert(TypeAttribute::ClosedDomain, ());
673 }
674 Definition {
675 visibility: Visibility::Public,
676 ty,
677 name: None,
678 name_span: None,
679 doc: None,
680 body: DefinitionBody::Struct {
681 generics: vec![],
682 fields: vec![],
683 kind: StructKind::Tuple,
684 methods: Default::default(),
685 constructor: None,
686 attributes,
687 },
688 }
689 }
690
691 fn int_const(ty: Type, value: u64) -> Definition {
692 Definition {
693 visibility: Visibility::Public,
694 ty,
695 name: None,
696 name_span: None,
697 doc: None,
698 body: DefinitionBody::Value {
699 allowed_lints: vec![],
700 go_hints: vec![],
701 go_name: None,
702 const_value: Some(Literal::Integer { value, text: None }),
703 },
704 }
705 }
706
707 fn insert(store: &mut Store, module: &str, name: &str, def: Definition) {
708 store.add_module(module);
709 store
710 .get_module_mut(module)
711 .unwrap()
712 .definitions
713 .insert(Symbol::from_raw(name), def);
714 }
715
716 #[test]
717 fn tagged_type_with_members_is_indexed_and_sorted() {
718 let mut store = Store::new();
719 let ty = nominal_int("m.Weekday");
720 insert(&mut store, "m", "m.Weekday", struct_def(ty.clone(), true));
721 insert(&mut store, "m", "m.Saturday", int_const(ty.clone(), 6));
722 insert(&mut store, "m", "m.Sunday", int_const(ty.clone(), 0));
723
724 store.build_closed_domains();
725
726 let domain = store
727 .closed_domains
728 .get("m.Weekday")
729 .expect("tagged type with members should be indexed");
730 assert_eq!(domain.base, SimpleKind::Int);
731 assert_eq!(domain.type_display.as_str(), "Weekday");
732 let names: Vec<&str> = domain
733 .members
734 .iter()
735 .map(|m| m.display_name.as_str())
736 .collect();
737 assert_eq!(names, vec!["Sunday", "Saturday"]);
738 }
739
740 #[test]
741 fn untagged_type_is_absent() {
742 let mut store = Store::new();
743 let ty = nominal_int("m.Plain");
744 insert(&mut store, "m", "m.Plain", struct_def(ty.clone(), false));
745 insert(&mut store, "m", "m.One", int_const(ty, 1));
746
747 store.build_closed_domains();
748
749 assert!(store.closed_domains.is_empty());
750 }
751
752 #[test]
753 fn tagged_type_without_members_records_no_domain() {
754 let mut store = Store::new();
755 insert(
756 &mut store,
757 "m",
758 "m.Empty",
759 struct_def(nominal_int("m.Empty"), true),
760 );
761
762 store.build_closed_domains();
763
764 assert!(!store.closed_domains.contains_key("m.Empty"));
765 }
766
767 #[test]
768 fn const_in_other_module_does_not_widen_domain() {
769 let mut store = Store::new();
770 let ty = nominal_int("lib.Weekday");
771 insert(
772 &mut store,
773 "lib",
774 "lib.Weekday",
775 struct_def(ty.clone(), true),
776 );
777 insert(&mut store, "lib", "lib.Sunday", int_const(ty.clone(), 0));
778 insert(&mut store, "user", "user.Bad", int_const(ty, 99));
779
780 store.build_closed_domains();
781
782 let domain = store.closed_domains.get("lib.Weekday").unwrap();
783 let names: Vec<&str> = domain
784 .members
785 .iter()
786 .map(|m| m.display_name.as_str())
787 .collect();
788 assert_eq!(names, vec!["Sunday"]);
789 }
790}