lust/typechecker/
mod.rs

1mod expr_checker;
2mod item_checker;
3mod stmt_checker;
4mod type_env;
5use crate::modules::{LoadedModule, ModuleImports};
6use crate::{
7    ast::*,
8    config::LustConfig,
9    error::{LustError, Result},
10};
11pub(super) use alloc::{
12    boxed::Box,
13    format,
14    string::{String, ToString},
15    vec,
16    vec::Vec,
17};
18use core::mem;
19use hashbrown::{HashMap, HashSet};
20pub use type_env::FunctionSignature;
21pub use type_env::TypeEnv;
22pub struct TypeChecker {
23    env: TypeEnv,
24    current_function_return_type: Option<Type>,
25    in_loop: bool,
26    pending_generic_instances: Option<HashMap<String, Type>>,
27    expected_lambda_signature: Option<(Vec<Type>, Option<Type>)>,
28    current_trait_bounds: HashMap<String, Vec<String>>,
29    current_module: Option<String>,
30    imports_by_module: HashMap<String, ModuleImports>,
31    expr_types_by_module: HashMap<String, HashMap<Span, Type>>,
32    variable_types_by_module: HashMap<String, HashMap<Span, Type>>,
33    short_circuit_info: HashMap<String, HashMap<Span, ShortCircuitInfo>>,
34}
35
36pub struct TypeCollection {
37    pub expr_types: HashMap<String, HashMap<Span, Type>>,
38    pub variable_types: HashMap<String, HashMap<Span, Type>>,
39}
40
41#[derive(Clone, Debug)]
42struct ShortCircuitInfo {
43    truthy: Option<Type>,
44    falsy: Option<Type>,
45    option_inner: Option<Type>,
46}
47
48impl TypeChecker {
49    pub fn new() -> Self {
50        Self::with_config(&LustConfig::default())
51    }
52
53    pub fn with_config(config: &LustConfig) -> Self {
54        Self {
55            env: TypeEnv::with_config(config),
56            current_function_return_type: None,
57            in_loop: false,
58            pending_generic_instances: None,
59            expected_lambda_signature: None,
60            current_trait_bounds: HashMap::new(),
61            current_module: None,
62            imports_by_module: HashMap::new(),
63            expr_types_by_module: HashMap::new(),
64            variable_types_by_module: HashMap::new(),
65            short_circuit_info: HashMap::new(),
66        }
67    }
68
69    fn dummy_span() -> Span {
70        Span::new(0, 0, 0, 0)
71    }
72
73    pub fn check_module(&mut self, items: &[Item]) -> Result<()> {
74        for item in items {
75            self.register_type_definition(item)?;
76        }
77
78        self.validate_struct_cycles()?;
79        self.env.push_scope();
80        self.register_module_init_locals(items)?;
81        for item in items {
82            self.check_item(item)?;
83        }
84
85        self.env.pop_scope();
86        Ok(())
87    }
88
89    pub fn check_program(&mut self, modules: &[LoadedModule]) -> Result<()> {
90        for m in modules {
91            self.current_module = Some(m.path.clone());
92            for item in &m.items {
93                self.register_type_definition(item)?;
94            }
95        }
96
97        self.validate_struct_cycles()?;
98        for m in modules {
99            self.current_module = Some(m.path.clone());
100            self.env.push_scope();
101            self.register_module_init_locals(&m.items)?;
102            for item in &m.items {
103                self.check_item(item)?;
104            }
105
106            self.env.pop_scope();
107        }
108
109        self.current_module = None;
110        Ok(())
111    }
112
113    fn validate_struct_cycles(&self) -> Result<()> {
114        use hashbrown::{HashMap, HashSet};
115        let struct_defs = self.env.struct_definitions();
116        if struct_defs.is_empty() {
117            return Ok(());
118        }
119
120        let mut simple_to_full: HashMap<String, Vec<String>> = HashMap::new();
121        for name in struct_defs.keys() {
122            let simple = name.rsplit('.').next().unwrap_or(name).to_string();
123            simple_to_full.entry(simple).or_default().push(name.clone());
124        }
125
126        let mut struct_has_weak: HashMap<String, bool> = HashMap::new();
127        for (name, def) in &struct_defs {
128            let has_weak = def
129                .fields
130                .iter()
131                .any(|field| matches!(field.ownership, FieldOwnership::Weak));
132            struct_has_weak.insert(name.clone(), has_weak);
133        }
134
135        let mut graph: HashMap<String, Vec<String>> = HashMap::new();
136        for (name, def) in &struct_defs {
137            let module_prefix = name.rsplit_once('.').map(|(module, _)| module.to_string());
138            let mut edges: HashSet<String> = HashSet::new();
139            for field in &def.fields {
140                if matches!(field.ownership, FieldOwnership::Weak) {
141                    let target = field.weak_target.as_ref().ok_or_else(|| {
142                        self.type_error(format!(
143                            "Field '{}.{}' is marked as 'ref' but has no target type",
144                            name, field.name
145                        ))
146                    })?;
147                    let target_name = if let TypeKind::Named(inner) = &target.kind {
148                        inner
149                    } else {
150                        return Err(self.type_error(format!(
151                            "Field '{}.{}' uses 'ref' but only struct types are supported",
152                            name, field.name
153                        )));
154                    };
155                    let resolved = self.resolve_struct_name_for_cycle(
156                        target_name.as_str(),
157                        module_prefix.as_deref(),
158                        &struct_defs,
159                        &simple_to_full,
160                    );
161                    if resolved.is_none() {
162                        return Err(self.type_error(format!(
163                            "Field '{}.{}' uses 'ref' but '{}' is not a known struct type",
164                            name, field.name, target_name
165                        )));
166                    }
167
168                    continue;
169                }
170
171                self.collect_strong_struct_targets(
172                    &field.ty,
173                    module_prefix.as_deref(),
174                    &struct_defs,
175                    &simple_to_full,
176                    &mut edges,
177                );
178            }
179
180            graph.insert(name.clone(), edges.into_iter().collect());
181        }
182
183        fn dfs(
184            node: &str,
185            graph: &HashMap<String, Vec<String>>,
186            visited: &mut HashSet<String>,
187            on_stack: &mut HashSet<String>,
188            stack: &mut Vec<String>,
189        ) -> Option<Vec<String>> {
190            visited.insert(node.to_string());
191            on_stack.insert(node.to_string());
192            stack.push(node.to_string());
193            if let Some(neighbors) = graph.get(node) {
194                for neighbor in neighbors {
195                    if !visited.contains(neighbor) {
196                        if let Some(cycle) = dfs(neighbor, graph, visited, on_stack, stack) {
197                            return Some(cycle);
198                        }
199                    } else if on_stack.contains(neighbor) {
200                        if let Some(pos) = stack.iter().position(|n| n == neighbor) {
201                            let mut cycle = stack[pos..].to_vec();
202                            cycle.push(neighbor.clone());
203                            return Some(cycle);
204                        }
205                    }
206                }
207            }
208
209            stack.pop();
210            on_stack.remove(node);
211            None
212        }
213
214        let mut visited: HashSet<String> = HashSet::new();
215        let mut on_stack: HashSet<String> = HashSet::new();
216        let mut stack: Vec<String> = Vec::new();
217        for name in struct_defs.keys() {
218            if !visited.contains(name) {
219                if let Some(cycle) = dfs(name, &graph, &mut visited, &mut on_stack, &mut stack) {
220                    let contains_weak = cycle
221                        .iter()
222                        .any(|node| struct_has_weak.get(node).copied().unwrap_or(false));
223                    if contains_weak {
224                        continue;
225                    }
226
227                    let description = cycle.join(" -> ");
228                    break;
229                    // return Err(self.type_error(format!(
230                    //     "Strong ownership cycle detected: {}. Mark at least one field as 'ref' to break the cycle.",
231                    //     description
232                    // )));
233                }
234            }
235        }
236
237        Ok(())
238    }
239
240    fn collect_strong_struct_targets(
241        &self,
242        ty: &Type,
243        parent_module: Option<&str>,
244        struct_defs: &HashMap<String, StructDef>,
245        simple_to_full: &HashMap<String, Vec<String>>,
246        out: &mut HashSet<String>,
247    ) {
248        match &ty.kind {
249            TypeKind::Named(name) => {
250                if let Some(resolved) = self.resolve_struct_name_for_cycle(
251                    name,
252                    parent_module,
253                    struct_defs,
254                    simple_to_full,
255                ) {
256                    out.insert(resolved);
257                }
258            }
259
260            TypeKind::Array(inner)
261            | TypeKind::Ref(inner)
262            | TypeKind::MutRef(inner)
263            | TypeKind::Option(inner) => {
264                self.collect_strong_struct_targets(
265                    inner,
266                    parent_module,
267                    struct_defs,
268                    simple_to_full,
269                    out,
270                );
271            }
272
273            TypeKind::Map(key, value) => {
274                self.collect_strong_struct_targets(
275                    key,
276                    parent_module,
277                    struct_defs,
278                    simple_to_full,
279                    out,
280                );
281                self.collect_strong_struct_targets(
282                    value,
283                    parent_module,
284                    struct_defs,
285                    simple_to_full,
286                    out,
287                );
288            }
289
290            TypeKind::Tuple(elements) | TypeKind::Union(elements) => {
291                for element in elements {
292                    self.collect_strong_struct_targets(
293                        element,
294                        parent_module,
295                        struct_defs,
296                        simple_to_full,
297                        out,
298                    );
299                }
300            }
301
302            TypeKind::Result(ok, err) => {
303                self.collect_strong_struct_targets(
304                    ok,
305                    parent_module,
306                    struct_defs,
307                    simple_to_full,
308                    out,
309                );
310                self.collect_strong_struct_targets(
311                    err,
312                    parent_module,
313                    struct_defs,
314                    simple_to_full,
315                    out,
316                );
317            }
318
319            TypeKind::GenericInstance { type_args, .. } => {
320                for arg in type_args {
321                    self.collect_strong_struct_targets(
322                        arg,
323                        parent_module,
324                        struct_defs,
325                        simple_to_full,
326                        out,
327                    );
328                }
329            }
330
331            _ => {}
332        }
333    }
334
335    fn resolve_struct_name_for_cycle(
336        &self,
337        name: &str,
338        parent_module: Option<&str>,
339        struct_defs: &HashMap<String, StructDef>,
340        simple_to_full: &HashMap<String, Vec<String>>,
341    ) -> Option<String> {
342        if struct_defs.contains_key(name) {
343            return Some(name.to_string());
344        }
345
346        if name.contains('.') {
347            return None;
348        }
349
350        if let Some(candidates) = simple_to_full.get(name) {
351            if candidates.len() == 1 {
352                return Some(candidates[0].clone());
353            }
354
355            if let Some(module) = parent_module {
356                for candidate in candidates {
357                    if let Some((candidate_module, _)) = candidate.rsplit_once('.') {
358                        if candidate_module == module {
359                            return Some(candidate.clone());
360                        }
361                    }
362                }
363            }
364        }
365
366        None
367    }
368
369    pub fn set_imports_by_module(&mut self, map: HashMap<String, ModuleImports>) {
370        self.imports_by_module = map;
371    }
372
373    pub fn take_type_info(&mut self) -> TypeCollection {
374        TypeCollection {
375            expr_types: mem::take(&mut self.expr_types_by_module),
376            variable_types: mem::take(&mut self.variable_types_by_module),
377        }
378    }
379
380    pub fn take_option_coercions(&mut self) -> HashMap<String, HashSet<Span>> {
381        let mut result: HashMap<String, HashSet<Span>> = HashMap::new();
382        let info = mem::take(&mut self.short_circuit_info);
383        for (module, entries) in info {
384            let mut spans: HashSet<Span> = HashSet::new();
385            for (span, entry) in entries {
386                if entry.option_inner.is_some() {
387                    spans.insert(span);
388                }
389            }
390            if !spans.is_empty() {
391                result.insert(module, spans);
392            }
393        }
394
395        result
396    }
397
398    pub fn function_signatures(&self) -> HashMap<String, type_env::FunctionSignature> {
399        self.env.function_signatures()
400    }
401
402    pub fn struct_definitions(&self) -> HashMap<String, StructDef> {
403        self.env.struct_definitions()
404    }
405
406    pub fn enum_definitions(&self) -> HashMap<String, EnumDef> {
407        self.env.enum_definitions()
408    }
409
410    fn register_module_init_locals(&mut self, items: &[Item]) -> Result<()> {
411        let module = match &self.current_module {
412            Some(m) => m.clone(),
413            None => return Ok(()),
414        };
415        let init_name = format!("__init@{}", module);
416        for item in items {
417            if let ItemKind::Function(func) = &item.kind {
418                if func.name == init_name {
419                    for stmt in &func.body {
420                        if let StmtKind::Local {
421                            bindings,
422                            ref mutable,
423                            initializer,
424                        } = &stmt.kind
425                        {
426                            self.check_local_stmt(
427                                bindings.as_slice(),
428                                *mutable,
429                                initializer.as_ref().map(|values| values.as_slice()),
430                            )?;
431                        }
432                    }
433                }
434            }
435        }
436
437        Ok(())
438    }
439
440    pub fn resolve_function_key(&self, name: &str) -> String {
441        if name.contains('.') || name.contains(':') {
442            return name.to_string();
443        }
444
445        if let Some(module) = &self.current_module {
446            if let Some(imports) = self.imports_by_module.get(module) {
447                if let Some(fq) = imports.function_aliases.get(name) {
448                    return fq.clone();
449                }
450            }
451
452            let qualified = format!("{}.{}", module, name);
453            if self.env.lookup_function(&qualified).is_some() {
454                return qualified;
455            }
456
457            if self.env.lookup_function(name).is_some() {
458                return name.to_string();
459            }
460
461            return qualified;
462        }
463
464        name.to_string()
465    }
466
467    pub fn resolve_module_alias(&self, alias: &str) -> Option<String> {
468        if let Some(module) = &self.current_module {
469            if let Some(imports) = self.imports_by_module.get(module) {
470                if let Some(m) = imports.module_aliases.get(alias) {
471                    return Some(m.clone());
472                }
473            }
474        }
475
476        None
477    }
478
479    pub fn register_external_struct(&mut self, mut def: StructDef) -> Result<()> {
480        def.name = self.resolve_type_key(&def.name);
481        for field in &mut def.fields {
482            field.ty = self.canonicalize_type(&field.ty);
483            if let Some(target) = &field.weak_target {
484                field.weak_target = Some(self.canonicalize_type(target));
485            }
486        }
487        self.env.register_struct(&def)
488    }
489
490    pub fn register_external_enum(&mut self, mut def: EnumDef) -> Result<()> {
491        def.name = self.resolve_type_key(&def.name);
492        for variant in &mut def.variants {
493            if let Some(fields) = &mut variant.fields {
494                for field in fields {
495                    *field = self.canonicalize_type(field);
496                }
497            }
498        }
499        self.env.register_enum(&def)
500    }
501
502    pub fn register_external_trait(&mut self, mut def: TraitDef) -> Result<()> {
503        def.name = self.resolve_type_key(&def.name);
504        for method in &mut def.methods {
505            for param in &mut method.params {
506                param.ty = self.canonicalize_type(&param.ty);
507            }
508            if let Some(ret) = method.return_type.clone() {
509                method.return_type = Some(self.canonicalize_type(&ret));
510            }
511        }
512        self.env.register_trait(&def)
513    }
514
515    pub fn register_external_function(
516        &mut self,
517        (name, mut signature): (String, FunctionSignature),
518    ) -> Result<()> {
519        signature.params = signature
520            .params
521            .into_iter()
522            .map(|ty| self.canonicalize_type(&ty))
523            .collect();
524        signature.return_type = self.canonicalize_type(&signature.return_type);
525        let canonical = self.resolve_type_key(&name);
526        self.env.register_or_update_function(canonical, signature)
527    }
528
529    pub fn register_external_impl(&mut self, mut impl_block: ImplBlock) -> Result<()> {
530        impl_block.target_type = self.canonicalize_type(&impl_block.target_type);
531        if let Some(trait_name) = &impl_block.trait_name {
532            impl_block.trait_name = Some(self.resolve_type_key(trait_name));
533        }
534        for method in &mut impl_block.methods {
535            for param in &mut method.params {
536                param.ty = self.canonicalize_type(&param.ty);
537            }
538            if let Some(ret) = method.return_type.clone() {
539                method.return_type = Some(self.canonicalize_type(&ret));
540            }
541        }
542
543        let type_name = match &impl_block.target_type.kind {
544            TypeKind::Named(name) => self.resolve_type_key(name),
545            TypeKind::GenericInstance { name, .. } => self.resolve_type_key(name),
546            _ => {
547                return Err(self.type_error(
548                    "Impl target must be a named type when registering from Rust".to_string(),
549                ))
550            }
551        };
552
553        self.env.register_impl(&impl_block);
554        for method in &impl_block.methods {
555            let params: Vec<Type> = method.params.iter().map(|p| p.ty.clone()).collect();
556            let return_type = method
557                .return_type
558                .clone()
559                .unwrap_or(Type::new(TypeKind::Unit, Span::dummy()));
560            let has_self = method.params.iter().any(|p| p.is_self);
561            let canonical_name = if method.name.contains(':') || method.name.contains('.') {
562                self.resolve_type_key(&method.name)
563            } else if has_self {
564                format!("{}:{}", type_name, method.name)
565            } else {
566                format!("{}.{}", type_name, method.name)
567            };
568            #[cfg(debug_assertions)]
569            eprintln!(
570                "register_external_impl canonical method {} (has_self={})",
571                canonical_name, has_self
572            );
573            let signature = FunctionSignature {
574                params,
575                return_type,
576                is_method: has_self,
577            };
578            self.env
579                .register_or_update_function(canonical_name, signature)?;
580        }
581
582        Ok(())
583    }
584
585    pub fn resolve_type_key(&self, name: &str) -> String {
586        if let Some((head, tail)) = name.split_once('.') {
587            if let Some(module) = &self.current_module {
588                if let Some(imports) = self.imports_by_module.get(module) {
589                    if let Some(real_module) = imports.module_aliases.get(head) {
590                        if tail.is_empty() {
591                            return real_module.clone();
592                        } else {
593                            return format!("{}.{}", real_module, tail);
594                        }
595                    }
596                }
597            }
598
599            return name.to_string();
600        }
601
602        if self.env.lookup_struct(name).is_some()
603            || self.env.lookup_enum(name).is_some()
604            || self.env.lookup_trait(name).is_some()
605        {
606            return name.to_string();
607        }
608
609        if self.env.is_builtin_type(name) {
610            return name.to_string();
611        }
612
613        if let Some(module) = &self.current_module {
614            if let Some(imports) = self.imports_by_module.get(module) {
615                if let Some(fq) = imports.type_aliases.get(name) {
616                    return fq.clone();
617                }
618            }
619
620            return format!("{}.{}", module, name);
621        }
622
623        name.to_string()
624    }
625
626    fn register_type_definition(&mut self, item: &Item) -> Result<()> {
627        match &item.kind {
628            ItemKind::Struct(s) => {
629                let mut s2 = s.clone();
630                if let Some(module) = &self.current_module {
631                    if !s2.name.contains('.') {
632                        s2.name = format!("{}.{}", module, s2.name);
633                    }
634                }
635
636                self.env.register_struct(&s2)?;
637            }
638
639            ItemKind::Enum(e) => {
640                let mut e2 = e.clone();
641                if let Some(module) = &self.current_module {
642                    if !e2.name.contains('.') {
643                        e2.name = format!("{}.{}", module, e2.name);
644                    }
645                }
646
647                self.env.register_enum(&e2)?;
648            }
649
650            ItemKind::Trait(t) => {
651                let mut t2 = t.clone();
652                if let Some(module) = &self.current_module {
653                    if !t2.name.contains('.') {
654                        t2.name = format!("{}.{}", module, t2.name);
655                    }
656                }
657
658                self.env.register_trait(&t2)?;
659            }
660
661            ItemKind::TypeAlias {
662                name,
663                type_params,
664                target,
665            } => {
666                let qname = if let Some(module) = &self.current_module {
667                    if name.contains('.') {
668                        name.clone()
669                    } else {
670                        format!("{}.{}", module, name)
671                    }
672                } else {
673                    name.clone()
674                };
675                self.env
676                    .register_type_alias(qname, type_params.clone(), target.clone())?;
677            }
678
679            _ => {}
680        }
681
682        Ok(())
683    }
684
685    fn type_error(&self, message: String) -> LustError {
686        LustError::TypeError { message }
687    }
688
689    fn type_error_at(&self, message: String, span: Span) -> LustError {
690        if span.start_line > 0 {
691            LustError::TypeErrorWithSpan {
692                message,
693                line: span.start_line,
694                column: span.start_col,
695                module: self.current_module.clone(),
696            }
697        } else {
698            LustError::TypeError { message }
699        }
700    }
701
702    fn types_equal(&self, t1: &Type, t2: &Type) -> bool {
703        t1.kind == t2.kind
704    }
705
706    pub fn canonicalize_type(&self, ty: &Type) -> Type {
707        use crate::ast::TypeKind as TK;
708        match &ty.kind {
709            TK::Named(name) => Type::new(TK::Named(self.resolve_type_key(name)), ty.span),
710            TK::Array(inner) => {
711                Type::new(TK::Array(Box::new(self.canonicalize_type(inner))), ty.span)
712            }
713
714            TK::Tuple(elements) => Type::new(
715                TK::Tuple(elements.iter().map(|t| self.canonicalize_type(t)).collect()),
716                ty.span,
717            ),
718            TK::Option(inner) => {
719                Type::new(TK::Option(Box::new(self.canonicalize_type(inner))), ty.span)
720            }
721
722            TK::Result(ok, err) => Type::new(
723                TK::Result(
724                    Box::new(self.canonicalize_type(ok)),
725                    Box::new(self.canonicalize_type(err)),
726                ),
727                ty.span,
728            ),
729            TK::Map(k, v) => Type::new(
730                TK::Map(
731                    Box::new(self.canonicalize_type(k)),
732                    Box::new(self.canonicalize_type(v)),
733                ),
734                ty.span,
735            ),
736            TK::Ref(inner) => Type::new(TK::Ref(Box::new(self.canonicalize_type(inner))), ty.span),
737            TK::MutRef(inner) => {
738                Type::new(TK::MutRef(Box::new(self.canonicalize_type(inner))), ty.span)
739            }
740
741            TK::Pointer { mutable, pointee } => Type::new(
742                TK::Pointer {
743                    mutable: *mutable,
744                    pointee: Box::new(self.canonicalize_type(pointee)),
745                },
746                ty.span,
747            ),
748            _ => ty.clone(),
749        }
750    }
751
752    fn unify(&self, expected: &Type, actual: &Type) -> Result<()> {
753        let span = if actual.span.start_line > 0 {
754            Some(actual.span)
755        } else if expected.span.start_line > 0 {
756            Some(expected.span)
757        } else {
758            None
759        };
760        self.unify_at(expected, actual, span)
761    }
762
763    fn unify_at(&self, expected: &Type, actual: &Type, span: Option<Span>) -> Result<()> {
764        if matches!(expected.kind, TypeKind::Unknown) || matches!(actual.kind, TypeKind::Unknown) {
765            return Ok(());
766        }
767
768        if matches!(expected.kind, TypeKind::Infer) || matches!(actual.kind, TypeKind::Infer) {
769            return Ok(());
770        }
771
772        match (&expected.kind, &actual.kind) {
773            (TypeKind::Union(expected_types), TypeKind::Union(actual_types)) => {
774                if expected_types.len() != actual_types.len() {
775                    return Err(self.type_error(format!(
776                        "Union types have different number of members: expected {}, got {}",
777                        expected_types.len(),
778                        actual_types.len()
779                    )));
780                }
781
782                for exp_type in expected_types {
783                    let mut found = false;
784                    for act_type in actual_types {
785                        if self.types_equal(exp_type, act_type) {
786                            found = true;
787                            break;
788                        }
789                    }
790
791                    if !found {
792                        return Err(match span {
793                            Some(s) => self.type_error_at(
794                                format!(
795                                    "Union type member '{}' not found in actual union",
796                                    exp_type
797                                ),
798                                s,
799                            ),
800                            None => self.type_error(format!(
801                                "Union type member '{}' not found in actual union",
802                                exp_type
803                            )),
804                        });
805                    }
806                }
807
808                return Ok(());
809            }
810
811            (TypeKind::Union(expected_types), _) => {
812                for union_member in expected_types {
813                    if self.unify(union_member, actual).is_ok() {
814                        return Ok(());
815                    }
816                }
817
818                return Err(match span {
819                    Some(s) => self.type_error_at(
820                        format!("Type '{}' is not compatible with union type", actual),
821                        s,
822                    ),
823                    None => self.type_error(format!(
824                        "Type '{}' is not compatible with union type",
825                        actual
826                    )),
827                });
828            }
829
830            (_, TypeKind::Union(actual_types)) => {
831                for union_member in actual_types {
832                    self.unify(expected, union_member)?;
833                }
834
835                return Ok(());
836            }
837
838            _ => {}
839        }
840
841        match (&expected.kind, &actual.kind) {
842            (TypeKind::Tuple(expected_elems), TypeKind::Tuple(actual_elems)) => {
843                if expected_elems.len() != actual_elems.len() {
844                    return Err(match span {
845                        Some(s) => self.type_error_at(
846                            format!(
847                                "Tuple length mismatch: expected {} element(s), got {}",
848                                expected_elems.len(),
849                                actual_elems.len()
850                            ),
851                            s,
852                        ),
853                        None => self.type_error(format!(
854                            "Tuple length mismatch: expected {} element(s), got {}",
855                            expected_elems.len(),
856                            actual_elems.len()
857                        )),
858                    });
859                }
860
861                for (exp_elem, act_elem) in expected_elems.iter().zip(actual_elems.iter()) {
862                    self.unify(exp_elem, act_elem)?;
863                }
864
865                return Ok(());
866            }
867
868            (TypeKind::Tuple(_), _) | (_, TypeKind::Tuple(_)) => {
869                return Err(match span {
870                    Some(s) => self.type_error_at(
871                        format!("Tuple type is not compatible with type '{}'", actual),
872                        s,
873                    ),
874                    None => self.type_error(format!(
875                        "Tuple type is not compatible with type '{}'",
876                        actual
877                    )),
878                })
879            }
880
881            (TypeKind::Named(name), TypeKind::Array(_))
882            | (TypeKind::Array(_), TypeKind::Named(name))
883                if name == "Array" =>
884            {
885                return Ok(());
886            }
887
888            (TypeKind::Array(exp_el), TypeKind::Array(act_el)) => {
889                if matches!(exp_el.kind, TypeKind::Unknown | TypeKind::Infer)
890                    || matches!(act_el.kind, TypeKind::Unknown | TypeKind::Infer)
891                {
892                    return Ok(());
893                } else {
894                    return self.unify(exp_el, act_el);
895                }
896            }
897
898            (TypeKind::Map(exp_key, exp_value), TypeKind::Map(act_key, act_value)) => {
899                self.unify(exp_key, act_key)?;
900                return self.unify(exp_value, act_value);
901            }
902
903            (TypeKind::Named(name), TypeKind::Option(_))
904            | (TypeKind::Option(_), TypeKind::Named(name))
905                if name == "Option" =>
906            {
907                return Ok(());
908            }
909
910            (TypeKind::Option(exp_inner), TypeKind::Option(act_inner)) => {
911                if matches!(exp_inner.kind, TypeKind::Unknown | TypeKind::Infer)
912                    || matches!(act_inner.kind, TypeKind::Unknown | TypeKind::Infer)
913                {
914                    return Ok(());
915                } else {
916                    return self.unify(exp_inner, act_inner);
917                }
918            }
919
920            (TypeKind::Named(name), TypeKind::Result(_, _))
921            | (TypeKind::Result(_, _), TypeKind::Named(name))
922                if name == "Result" =>
923            {
924                return Ok(());
925            }
926
927            (TypeKind::Result(exp_ok, exp_err), TypeKind::Result(act_ok, act_err)) => {
928                if matches!(exp_ok.kind, TypeKind::Unknown | TypeKind::Infer)
929                    || matches!(act_ok.kind, TypeKind::Unknown | TypeKind::Infer)
930                {
931                    if matches!(exp_err.kind, TypeKind::Unknown | TypeKind::Infer)
932                        || matches!(act_err.kind, TypeKind::Unknown | TypeKind::Infer)
933                    {
934                        return Ok(());
935                    } else {
936                        return self.unify(exp_err, act_err);
937                    }
938                } else {
939                    self.unify(exp_ok, act_ok)?;
940                    return self.unify(exp_err, act_err);
941                }
942            }
943
944            _ => {}
945        }
946
947        if self.types_equal(expected, actual) {
948            Ok(())
949        } else {
950            Err(match span {
951                Some(s) => self.type_error_at(
952                    format!("Type mismatch: expected '{}', got '{}'", expected, actual),
953                    s,
954                ),
955                None => self.type_error(format!(
956                    "Type mismatch: expected '{}', got '{}'",
957                    expected, actual
958                )),
959            })
960        }
961    }
962
963    fn types_compatible(&self, expected: &Type, actual: &Type) -> bool {
964        if matches!(expected.kind, TypeKind::Unknown) || matches!(actual.kind, TypeKind::Unknown) {
965            return true;
966        }
967
968        if matches!(expected.kind, TypeKind::Infer) || matches!(actual.kind, TypeKind::Infer) {
969            return true;
970        }
971
972        match (&expected.kind, &actual.kind) {
973            (TypeKind::Generic(_), TypeKind::Generic(_)) => return true,
974            (TypeKind::Generic(_), _) | (_, TypeKind::Generic(_)) => return true,
975            _ => {}
976        }
977
978        match (&expected.kind, &actual.kind) {
979            (TypeKind::Array(e1), TypeKind::Array(e2)) => {
980                return self.types_compatible(e1, e2);
981            }
982
983            (TypeKind::Named(name), TypeKind::Array(_))
984            | (TypeKind::Array(_), TypeKind::Named(name))
985                if name == "Array" =>
986            {
987                return true;
988            }
989
990            _ => {}
991        }
992
993        match (&expected.kind, &actual.kind) {
994            (TypeKind::Map(k1, v1), TypeKind::Map(k2, v2)) => {
995                return self.types_compatible(k1, k2) && self.types_compatible(v1, v2);
996            }
997
998            _ => {}
999        }
1000
1001        match (&expected.kind, &actual.kind) {
1002            (TypeKind::Option(t1), TypeKind::Option(t2)) => {
1003                return self.types_compatible(t1, t2);
1004            }
1005
1006            (TypeKind::Named(name), TypeKind::Option(_))
1007            | (TypeKind::Option(_), TypeKind::Named(name))
1008                if name == "Option" =>
1009            {
1010                return true;
1011            }
1012
1013            _ => {}
1014        }
1015
1016        match (&expected.kind, &actual.kind) {
1017            (TypeKind::Result(ok1, err1), TypeKind::Result(ok2, err2)) => {
1018                return self.types_compatible(ok1, ok2) && self.types_compatible(err1, err2);
1019            }
1020
1021            (TypeKind::Named(name), TypeKind::Result(_, _))
1022            | (TypeKind::Result(_, _), TypeKind::Named(name))
1023                if name == "Result" =>
1024            {
1025                return true;
1026            }
1027
1028            _ => {}
1029        }
1030
1031        match (&expected.kind, &actual.kind) {
1032            (
1033                TypeKind::Function {
1034                    params: p1,
1035                    return_type: r1,
1036                },
1037                TypeKind::Function {
1038                    params: p2,
1039                    return_type: r2,
1040                },
1041            ) => {
1042                if p1.len() != p2.len() {
1043                    return false;
1044                }
1045
1046                for (t1, t2) in p1.iter().zip(p2.iter()) {
1047                    if !self.types_compatible(t1, t2) {
1048                        return false;
1049                    }
1050                }
1051
1052                return self.types_compatible(r1, r2);
1053            }
1054
1055            _ => {}
1056        }
1057
1058        self.types_equal(expected, actual)
1059    }
1060
1061    fn unify_with_bounds(&self, expected: &Type, actual: &Type) -> Result<()> {
1062        if let TypeKind::Generic(type_param) = &expected.kind {
1063            if let Some(trait_names) = self.current_trait_bounds.get(type_param) {
1064                for trait_name in trait_names {
1065                    if !self.env.type_implements_trait(actual, trait_name) {
1066                        return Err(self.type_error(format!(
1067                            "Type '{}' does not implement required trait '{}'",
1068                            actual, trait_name
1069                        )));
1070                    }
1071                }
1072
1073                return Ok(());
1074            }
1075
1076            return Ok(());
1077        }
1078
1079        self.unify(expected, actual)
1080    }
1081
1082    fn record_short_circuit_info(&mut self, span: Span, info: &ShortCircuitInfo) {
1083        let truthy = info.truthy.as_ref().map(|ty| self.canonicalize_type(ty));
1084        let falsy = info.falsy.as_ref().map(|ty| self.canonicalize_type(ty));
1085        let option_inner = info
1086            .option_inner
1087            .as_ref()
1088            .map(|ty| self.canonicalize_type(ty));
1089        let module_key = self.current_module_key();
1090        self.short_circuit_info
1091            .entry(module_key)
1092            .or_default()
1093            .insert(
1094                span,
1095                ShortCircuitInfo {
1096                    truthy,
1097                    falsy,
1098                    option_inner,
1099                },
1100            );
1101    }
1102
1103    fn short_circuit_profile(&self, expr: &Expr, ty: &Type) -> ShortCircuitInfo {
1104        let module_key = self
1105            .current_module
1106            .as_ref()
1107            .map(String::as_str)
1108            .unwrap_or("");
1109        if let Some(module_map) = self.short_circuit_info.get(module_key) {
1110            if let Some(info) = module_map.get(&expr.span) {
1111                return info.clone();
1112            }
1113        }
1114
1115        ShortCircuitInfo {
1116            truthy: if self.type_can_be_truthy(ty) {
1117                Some(self.canonicalize_type(ty))
1118            } else {
1119                None
1120            },
1121            falsy: self.extract_falsy_type(ty),
1122            option_inner: None,
1123        }
1124    }
1125
1126    fn current_module_key(&self) -> String {
1127        self.current_module
1128            .as_ref()
1129            .cloned()
1130            .unwrap_or_else(|| "".to_string())
1131    }
1132
1133    fn clear_option_for_span(&mut self, span: Span) {
1134        let module_key = self.current_module_key();
1135        if let Some(module_map) = self.short_circuit_info.get_mut(&module_key) {
1136            if let Some(info) = module_map.get_mut(&span) {
1137                info.option_inner = None;
1138            }
1139        }
1140    }
1141
1142    fn type_can_be_truthy(&self, ty: &Type) -> bool {
1143        match &ty.kind {
1144            TypeKind::Union(members) => {
1145                members.iter().any(|member| self.type_can_be_truthy(member))
1146            }
1147            TypeKind::Bool => true,
1148            TypeKind::Unknown => true,
1149            _ => true,
1150        }
1151    }
1152
1153    fn type_can_be_falsy(&self, ty: &Type) -> bool {
1154        match &ty.kind {
1155            TypeKind::Union(members) => members.iter().any(|member| self.type_can_be_falsy(member)),
1156            TypeKind::Bool => true,
1157            TypeKind::Unknown => true,
1158            TypeKind::Option(_) => true,
1159            _ => false,
1160        }
1161    }
1162
1163    fn extract_falsy_type(&self, ty: &Type) -> Option<Type> {
1164        match &ty.kind {
1165            TypeKind::Bool => Some(Type::new(TypeKind::Bool, ty.span)),
1166            TypeKind::Unknown => Some(Type::new(TypeKind::Unknown, ty.span)),
1167            TypeKind::Option(inner) => Some(Type::new(
1168                TypeKind::Option(Box::new(self.canonicalize_type(inner))),
1169                ty.span,
1170            )),
1171            TypeKind::Union(members) => {
1172                let mut parts = Vec::new();
1173                for member in members {
1174                    if let Some(part) = self.extract_falsy_type(member) {
1175                        parts.push(part);
1176                    }
1177                }
1178                self.merge_optional_types(parts)
1179            }
1180            _ => None,
1181        }
1182    }
1183
1184    fn merge_optional_types(&self, types: Vec<Type>) -> Option<Type> {
1185        if types.is_empty() {
1186            return None;
1187        }
1188
1189        Some(self.make_union_from_types(types))
1190    }
1191
1192    fn make_union_from_types(&self, types: Vec<Type>) -> Type {
1193        let mut flat: Vec<Type> = Vec::new();
1194        for ty in types {
1195            let canonical = self.canonicalize_type(&ty);
1196            match &canonical.kind {
1197                TypeKind::Union(members) => {
1198                    for member in members {
1199                        self.push_unique_type(&mut flat, member.clone());
1200                    }
1201                }
1202                _ => self.push_unique_type(&mut flat, canonical),
1203            }
1204        }
1205
1206        match flat.len() {
1207            0 => Type::new(TypeKind::Unknown, Self::dummy_span()),
1208            1 => flat.into_iter().next().unwrap(),
1209            _ => Type::new(TypeKind::Union(flat), Self::dummy_span()),
1210        }
1211    }
1212
1213    fn push_unique_type(&self, list: &mut Vec<Type>, candidate: Type) {
1214        if !list
1215            .iter()
1216            .any(|existing| self.types_equal(existing, &candidate))
1217        {
1218            list.push(candidate);
1219        }
1220    }
1221
1222    fn combine_truthy_falsy(&self, truthy: Option<Type>, falsy: Option<Type>) -> Type {
1223        match (truthy, falsy) {
1224            (Some(t), Some(f)) => self.make_union_from_types(vec![t, f]),
1225            (Some(t), None) => t,
1226            (None, Some(f)) => f,
1227            (None, None) => Type::new(TypeKind::Unknown, Self::dummy_span()),
1228        }
1229    }
1230
1231    fn is_bool_like(&self, ty: &Type) -> bool {
1232        match &ty.kind {
1233            TypeKind::Bool => true,
1234            TypeKind::Union(members) => members.iter().all(|member| self.is_bool_like(member)),
1235            _ => false,
1236        }
1237    }
1238
1239    fn option_inner_type<'a>(&self, ty: &'a Type) -> Option<&'a Type> {
1240        match &ty.kind {
1241            TypeKind::Option(inner) => Some(inner.as_ref()),
1242            TypeKind::Union(members) => {
1243                for member in members {
1244                    if let Some(inner) = self.option_inner_type(member) {
1245                        return Some(inner);
1246                    }
1247                }
1248                None
1249            }
1250            _ => None,
1251        }
1252    }
1253
1254    fn should_optionize(&self, left: &Type, right: &Type) -> bool {
1255        self.is_bool_like(left)
1256            && !self.is_bool_like(right)
1257            && self.option_inner_type(right).is_none()
1258    }
1259}