Skip to main content

tsz_solver/contextual/
mod.rs

1//! Contextual typing (reverse inference).
2//!
3//! Contextual typing allows type information to flow "backwards" from
4//! an expected type to an expression. This is used for:
5//! - Arrow function parameters: `const f: (x: string) => void = (x) => ...`
6//! - Array literals: `const arr: number[] = [1, 2, 3]`
7//! - Object literals: `const obj: {x: number} = {x: 1}`
8//!
9//! The key insight is that when we have an expected type, we can use it
10//! to infer types for parts of the expression that would otherwise be unknown.
11//!
12//! The visitor-based type extractors used by [`ContextualTypeContext`] are in
13//! the [`extractors`] submodule.
14
15pub(crate) mod extractors;
16
17use crate::TypeDatabase;
18use crate::contextual::extractors::{
19    ApplicationArgExtractor, ArrayElementExtractor, ParameterExtractor, ParameterForCallExtractor,
20    PropertyExtractor, ReturnTypeExtractor, ThisTypeExtractor, ThisTypeMarkerExtractor,
21    TupleElementExtractor, collect_single_or_union,
22};
23#[cfg(test)]
24use crate::types::*;
25use crate::types::{TypeData, TypeId};
26
27/// Context for contextual typing.
28/// Holds the expected type and provides methods to extract type information.
29pub struct ContextualTypeContext<'a> {
30    interner: &'a dyn TypeDatabase,
31    /// The expected type (contextual type)
32    expected: Option<TypeId>,
33    /// Whether noImplicitAny is enabled (affects contextual typing for multi-signature functions)
34    no_implicit_any: bool,
35}
36
37impl<'a> ContextualTypeContext<'a> {
38    /// Create a new contextual type context.
39    /// Defaults to `no_implicit_any: false` for compatibility.
40    pub fn new(interner: &'a dyn TypeDatabase) -> Self {
41        ContextualTypeContext {
42            interner,
43            expected: None,
44            no_implicit_any: false,
45        }
46    }
47
48    /// Create a context with an expected type.
49    /// Defaults to `no_implicit_any: false` for compatibility.
50    pub fn with_expected(interner: &'a dyn TypeDatabase, expected: TypeId) -> Self {
51        ContextualTypeContext {
52            interner,
53            expected: Some(expected),
54            no_implicit_any: false,
55        }
56    }
57
58    /// Create a context with an expected type and explicit noImplicitAny setting.
59    pub fn with_expected_and_options(
60        interner: &'a dyn TypeDatabase,
61        expected: TypeId,
62        no_implicit_any: bool,
63    ) -> Self {
64        ContextualTypeContext {
65            interner,
66            expected: Some(expected),
67            no_implicit_any,
68        }
69    }
70
71    /// Get the expected type.
72    pub const fn expected(&self) -> Option<TypeId> {
73        self.expected
74    }
75
76    /// Check if we have a contextual type.
77    pub const fn has_context(&self) -> bool {
78        self.expected.is_some()
79    }
80
81    /// Get the contextual type for a function parameter at the given index.
82    ///
83    /// Example:
84    /// ```typescript
85    /// type Handler = (e: string, i: number) => void;
86    /// const h: Handler = (x, y) => {};  // x: string, y: number from context
87    /// ```
88    pub fn get_parameter_type(&self, index: usize) -> Option<TypeId> {
89        let expected = self.expected?;
90
91        // Handle Union explicitly - collect parameter types from all members
92        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
93            let members = self.interner.type_list(members);
94            let param_types: Vec<TypeId> = members
95                .iter()
96                .filter_map(|&m| {
97                    let ctx = ContextualTypeContext::with_expected_and_options(
98                        self.interner,
99                        m,
100                        self.no_implicit_any,
101                    );
102                    ctx.get_parameter_type(index)
103                })
104                .collect();
105
106            return collect_single_or_union(self.interner, param_types);
107        }
108
109        // Handle Application explicitly - unwrap to base type
110        if let Some(TypeData::Application(app_id)) = self.interner.lookup(expected) {
111            let app = self.interner.type_application(app_id);
112            let ctx = ContextualTypeContext::with_expected_and_options(
113                self.interner,
114                app.base,
115                self.no_implicit_any,
116            );
117            return ctx.get_parameter_type(index);
118        }
119
120        // Handle Intersection explicitly - pick the first callable member's parameter type
121        if let Some(TypeData::Intersection(members)) = self.interner.lookup(expected) {
122            let members = self.interner.type_list(members);
123            for &m in members.iter() {
124                let ctx = ContextualTypeContext::with_expected_and_options(
125                    self.interner,
126                    m,
127                    self.no_implicit_any,
128                );
129                if let Some(param_type) = ctx.get_parameter_type(index) {
130                    return Some(param_type);
131                }
132            }
133            return None;
134        }
135
136        // Handle Mapped, Conditional, and Lazy types by evaluating them first
137        if let Some(TypeData::Mapped(_) | TypeData::Conditional(_) | TypeData::Lazy(_)) =
138            self.interner.lookup(expected)
139        {
140            let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
141            if evaluated != expected {
142                let ctx = ContextualTypeContext::with_expected_and_options(
143                    self.interner,
144                    evaluated,
145                    self.no_implicit_any,
146                );
147                return ctx.get_parameter_type(index);
148            }
149        }
150
151        // Use visitor for Function/Callable types
152        let mut extractor = ParameterExtractor::new(self.interner, index, self.no_implicit_any);
153        extractor.extract(expected)
154    }
155
156    /// Get the contextual type for a call argument at the given index and arity.
157    pub fn get_parameter_type_for_call(&self, index: usize, arg_count: usize) -> Option<TypeId> {
158        let expected = self.expected?;
159
160        // Handle Union explicitly - collect parameter types from all members
161        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
162            let members = self.interner.type_list(members);
163            let param_types: Vec<TypeId> = members
164                .iter()
165                .filter_map(|&m| {
166                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
167                    ctx.get_parameter_type_for_call(index, arg_count)
168                })
169                .collect();
170
171            return collect_single_or_union(self.interner, param_types);
172        }
173
174        // Handle Application explicitly - unwrap to base type
175        if let Some(TypeData::Application(app_id)) = self.interner.lookup(expected) {
176            let app = self.interner.type_application(app_id);
177            let ctx = ContextualTypeContext::with_expected(self.interner, app.base);
178            return ctx.get_parameter_type_for_call(index, arg_count);
179        }
180
181        // Handle Intersection explicitly - pick the first callable member's parameter type
182        if let Some(TypeData::Intersection(members)) = self.interner.lookup(expected) {
183            let members = self.interner.type_list(members);
184            for &m in members.iter() {
185                let ctx = ContextualTypeContext::with_expected(self.interner, m);
186                if let Some(param_type) = ctx.get_parameter_type_for_call(index, arg_count) {
187                    return Some(param_type);
188                }
189            }
190            return None;
191        }
192
193        // Use visitor for Function/Callable types
194        let mut extractor = ParameterForCallExtractor::new(self.interner, index, arg_count);
195        extractor.extract(expected)
196    }
197
198    /// Get the contextual type for a `this` parameter, if present on the expected type.
199    pub fn get_this_type(&self) -> Option<TypeId> {
200        let expected = self.expected?;
201
202        // Handle Union explicitly - collect this types from all members
203        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
204            let members = self.interner.type_list(members);
205            let this_types: Vec<TypeId> = members
206                .iter()
207                .filter_map(|&m| {
208                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
209                    ctx.get_this_type()
210                })
211                .collect();
212
213            return collect_single_or_union(self.interner, this_types);
214        }
215
216        // Handle Application explicitly - unwrap to base type
217        if let Some(TypeData::Application(app_id)) = self.interner.lookup(expected) {
218            let app = self.interner.type_application(app_id);
219            let ctx = ContextualTypeContext::with_expected(self.interner, app.base);
220            return ctx.get_this_type();
221        }
222
223        // Use visitor for Function/Callable types
224        let mut extractor = ThisTypeExtractor::new(self.interner);
225        extractor.extract(expected)
226    }
227
228    /// Get the type T from a `ThisType`<T> marker in the contextual type.
229    ///
230    /// This is used for the Vue 2 / Options API pattern where object literal
231    /// methods have their `this` type overridden by contextual markers.
232    ///
233    /// Example:
234    /// ```typescript
235    /// type ObjectDescriptor<D, M> = {
236    ///     methods?: M & ThisType<D & M>;
237    /// };
238    /// const obj: ObjectDescriptor<{x: number}, {greet(): void}> = {
239    ///     methods: {
240    ///         greet() { console.log(this.x); } // this is D & M
241    ///     }
242    /// };
243    /// ```
244    pub fn get_this_type_from_marker(&self) -> Option<TypeId> {
245        let expected = self.expected?;
246        let mut extractor = ThisTypeMarkerExtractor::new(self.interner);
247        extractor.extract(expected)
248    }
249
250    /// Get the contextual return type for a function.
251    pub fn get_return_type(&self) -> Option<TypeId> {
252        let expected = self.expected?;
253
254        // Handle Union explicitly - collect return types from all members
255        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
256            let members = self.interner.type_list(members);
257            let return_types: Vec<TypeId> = members
258                .iter()
259                .filter_map(|&m| {
260                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
261                    ctx.get_return_type()
262                })
263                .collect();
264
265            return collect_single_or_union(self.interner, return_types);
266        }
267
268        // Handle Application explicitly - unwrap to base type
269        if let Some(TypeData::Application(app_id)) = self.interner.lookup(expected) {
270            let app = self.interner.type_application(app_id);
271            let ctx = ContextualTypeContext::with_expected(self.interner, app.base);
272            return ctx.get_return_type();
273        }
274
275        // Handle Lazy, Mapped, and Conditional types by evaluating first
276        if let Some(TypeData::Lazy(_) | TypeData::Mapped(_) | TypeData::Conditional(_)) =
277            self.interner.lookup(expected)
278        {
279            let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
280            if evaluated != expected {
281                let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
282                return ctx.get_return_type();
283            }
284        }
285
286        // Use visitor for Function/Callable types
287        let mut extractor = ReturnTypeExtractor::new(self.interner);
288        extractor.extract(expected)
289    }
290
291    /// Get the contextual element type for an array.
292    ///
293    /// Example:
294    /// ```typescript
295    /// const arr: number[] = [1, 2, 3];  // elements are contextually typed as number
296    /// ```
297    pub fn get_array_element_type(&self) -> Option<TypeId> {
298        let expected = self.expected?;
299
300        // Handle Union explicitly - collect element types from all array members
301        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
302            let members = self.interner.type_list(members);
303            let elem_types: Vec<TypeId> = members
304                .iter()
305                .filter_map(|&m| {
306                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
307                    ctx.get_array_element_type()
308                })
309                .collect();
310            return collect_single_or_union(self.interner, elem_types);
311        }
312
313        // Handle Application explicitly - evaluate to resolve type aliases.
314        // For generic iterable-like types (Iterable<T>, ReadonlyArray<T>, ArrayLike<T>, etc.),
315        // try to extract the element type from the first type argument, since these types
316        // all use T as their element type.
317        if let Some(TypeData::Application(app_id)) = self.interner.lookup(expected) {
318            let app = self.interner.type_application(app_id);
319            // First try evaluating to see if it resolves to an array type
320            let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
321            if evaluated != expected {
322                let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
323                if let Some(elem) = ctx.get_array_element_type() {
324                    return Some(elem);
325                }
326                // Check if the evaluated type has iterable-like structure
327                if !app.args.is_empty() && self.is_iterable_like_object(evaluated) {
328                    return Some(app.args[0]);
329                }
330            }
331            // If evaluation didn't change the type (e.g., Lazy(DefId) base that can't
332            // be resolved without TypeEnvironment), use the first type argument as
333            // the element type. This is a reasonable heuristic because assigning an
334            // array literal to a generic type like Iterable<T> or ArrayLike<T> means
335            // T is the expected element type.
336            if !app.args.is_empty() && evaluated == expected {
337                return Some(app.args[0]);
338            }
339        }
340
341        // Handle Mapped/Conditional/Lazy types
342        if let Some(TypeData::Mapped(_) | TypeData::Conditional(_) | TypeData::Lazy(_)) =
343            self.interner.lookup(expected)
344        {
345            let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
346            if evaluated != expected {
347                let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
348                return ctx.get_array_element_type();
349            }
350        }
351
352        // Handle TypeParameter - use its constraint for element extraction
353        if let Some(constraint) =
354            crate::type_queries::get_type_parameter_constraint(self.interner, expected)
355        {
356            let ctx = ContextualTypeContext::with_expected(self.interner, constraint);
357            return ctx.get_array_element_type();
358        }
359
360        // Handle Intersection - pick the first array member's element type
361        if let Some(TypeData::Intersection(members)) = self.interner.lookup(expected) {
362            let members = self.interner.type_list(members);
363            for &m in members.iter() {
364                let ctx = ContextualTypeContext::with_expected(self.interner, m);
365                if let Some(elem_type) = ctx.get_array_element_type() {
366                    return Some(elem_type);
367                }
368            }
369            return None;
370        }
371
372        let mut extractor = ArrayElementExtractor::new(self.interner);
373        extractor.extract(expected)
374    }
375
376    /// Check if a type looks like an iterable or array-like object type.
377    /// This is used as a heuristic to determine whether the first type argument
378    /// of an Application is the element type (for contextual typing of array literals).
379    fn is_iterable_like_object(&self, type_id: TypeId) -> bool {
380        use crate::types::TypeData;
381
382        // Check if type is an object with properties suggesting iterable/array-like behavior
383        match self.interner.lookup(type_id) {
384            Some(TypeData::Object(shape_id)) => {
385                let shape = self.interner.object_shape(shape_id);
386                // Has number index → array-like (ArrayLike<T>, ReadonlyArray<T>)
387                if shape.number_index.is_some() {
388                    return true;
389                }
390                // Has Symbol.iterator property → iterable (Iterable<T>)
391                for prop in &shape.properties {
392                    let name = self.interner.resolve_atom(prop.name);
393                    if name == "__@iterator" || name == "[Symbol.iterator]" {
394                        return true;
395                    }
396                }
397                false
398            }
399            Some(TypeData::Intersection(members)) => {
400                let members = self.interner.type_list(members);
401                members.iter().any(|&m| self.is_iterable_like_object(m))
402            }
403            _ => false,
404        }
405    }
406
407    /// Get the contextual type for a specific tuple element.
408    pub fn get_tuple_element_type(&self, index: usize) -> Option<TypeId> {
409        let expected = self.expected?;
410
411        // Handle Union explicitly - collect tuple element types from all members
412        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
413            let members = self.interner.type_list(members);
414            let elem_types: Vec<TypeId> = members
415                .iter()
416                .filter_map(|&m| {
417                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
418                    ctx.get_tuple_element_type(index)
419                })
420                .collect();
421            return collect_single_or_union(self.interner, elem_types);
422        }
423
424        // Handle Application explicitly - evaluate to resolve type aliases
425        if let Some(TypeData::Application(_)) = self.interner.lookup(expected) {
426            let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
427            if evaluated != expected {
428                let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
429                return ctx.get_tuple_element_type(index);
430            }
431        }
432
433        // Handle TypeParameter - use its constraint
434        if let Some(constraint) =
435            crate::type_queries::get_type_parameter_constraint(self.interner, expected)
436        {
437            let ctx = ContextualTypeContext::with_expected(self.interner, constraint);
438            return ctx.get_tuple_element_type(index);
439        }
440
441        // Handle Mapped, Conditional, and Lazy types by evaluating them first
442        if let Some(TypeData::Mapped(_) | TypeData::Conditional(_) | TypeData::Lazy(_)) =
443            self.interner.lookup(expected)
444        {
445            let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
446            if evaluated != expected {
447                let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
448                return ctx.get_tuple_element_type(index);
449            }
450        }
451
452        let mut extractor = TupleElementExtractor::new(self.interner, index);
453        extractor.extract(expected)
454    }
455
456    /// Get the contextual type for an object property.
457    ///
458    /// Example:
459    /// ```typescript
460    /// const obj: {x: number, y: string} = {x: 1, y: "hi"};
461    /// ```
462    pub fn get_property_type(&self, name: &str) -> Option<TypeId> {
463        let expected = self.expected?;
464
465        // Handle Union explicitly - collect property types from all members
466        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
467            let members = self.interner.type_list(members);
468            let prop_types: Vec<TypeId> = members
469                .iter()
470                .filter_map(|&m| {
471                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
472                    ctx.get_property_type(name)
473                })
474                .collect();
475
476            return if prop_types.is_empty() {
477                None
478            } else if prop_types.len() == 1 {
479                Some(prop_types[0])
480            } else {
481                // CRITICAL: Use union_preserve_members to keep literal types intact
482                // For discriminated unions like `{ success: false } | { success: true }`,
483                // the property type should be `false | true`, NOT widened to `boolean`.
484                // This preserves literal types for contextual typing.
485                Some(self.interner.union_preserve_members(prop_types))
486            };
487        }
488
489        // Handle Mapped, Conditional, and Application types.
490        // These complex types need to be resolved to concrete object types before
491        // property extraction can work.
492        match self.interner.lookup(expected) {
493            Some(TypeData::Mapped(mapped_id)) => {
494                // First try evaluating the mapped type directly
495                let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
496                if evaluated != expected {
497                    let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
498                    return ctx.get_property_type(name);
499                }
500                // If evaluation deferred (e.g. { [K in keyof T]: TakeString } where T is a type
501                // parameter), use the mapped type's template as the contextual property type
502                // IF the template doesn't reference the mapped type's bound parameter.
503                // For example, { [P in keyof T]: TakeString } has template=TakeString which
504                // is independent of P, so it's safe to use directly. But { [P in K]: T[P] }
505                // has template=T[P] which depends on P, so we can't use it.
506                let mapped = self.interner.mapped_type(mapped_id);
507                if mapped.template != TypeId::ANY
508                    && mapped.template != TypeId::ERROR
509                    && mapped.template != TypeId::NEVER
510                    && !crate::visitor::contains_type_matching(
511                        self.interner,
512                        mapped.template,
513                        |key| matches!(key, TypeData::BoundParameter(_)),
514                    )
515                {
516                    return Some(mapped.template);
517                }
518                // Fall back to the constraint of the mapped type's source.
519                // For `keyof P` where `P extends Props`, use `Props` as the contextual type.
520                if let Some(TypeData::KeyOf(operand)) = self.interner.lookup(mapped.constraint) {
521                    // The operand may be a Lazy type wrapping a type parameter — resolve it
522                    let resolved_operand =
523                        crate::evaluation::evaluate::evaluate_type(self.interner, operand);
524                    if let Some(constraint) = crate::type_queries::get_type_parameter_constraint(
525                        self.interner,
526                        resolved_operand,
527                    ) {
528                        let ctx = ContextualTypeContext::with_expected(self.interner, constraint);
529                        return ctx.get_property_type(name);
530                    }
531                    // Also try the original operand (may already be a TypeParameter)
532                    if let Some(constraint) =
533                        crate::type_queries::get_type_parameter_constraint(self.interner, operand)
534                    {
535                        let ctx = ContextualTypeContext::with_expected(self.interner, constraint);
536                        return ctx.get_property_type(name);
537                    }
538                }
539            }
540            Some(TypeData::Application(app_id)) => {
541                let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
542                if evaluated != expected {
543                    let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
544                    return ctx.get_property_type(name);
545                }
546                // Fallback for unevaluated Application types (e.g. Readonly<T>, Partial<T>).
547                // When evaluation fails (e.g. due to RefCell borrow conflicts during contextual
548                // typing), try to extract the property from the type argument directly.
549                // This is correct for homomorphic mapped types where property types are preserved.
550                let app = self.interner.type_application(app_id);
551                if !app.args.is_empty() {
552                    let ctx = ContextualTypeContext::with_expected(self.interner, app.args[0]);
553                    if let Some(prop) = ctx.get_property_type(name) {
554                        return Some(prop);
555                    }
556                }
557            }
558            Some(TypeData::Conditional(_) | TypeData::Lazy(_)) => {
559                let evaluated = crate::evaluation::evaluate::evaluate_type(self.interner, expected);
560                if evaluated != expected {
561                    let ctx = ContextualTypeContext::with_expected(self.interner, evaluated);
562                    return ctx.get_property_type(name);
563                }
564            }
565            _ => {}
566        }
567
568        // Handle TypeParameter - use its constraint for property extraction
569        // Example: Actions extends ActionsObject<State>
570        // When getting property from Actions, use ActionsObject<State> instead
571        if let Some(constraint) =
572            crate::type_queries::get_type_parameter_constraint(self.interner, expected)
573        {
574            let ctx = ContextualTypeContext::with_expected(self.interner, constraint);
575            return ctx.get_property_type(name);
576        }
577
578        // Use visitor for Object types
579        let mut extractor = PropertyExtractor::new(self.interner, name);
580        extractor.extract(expected)
581    }
582
583    /// Create a child context for a nested expression.
584    /// This is used when checking nested structures with contextual types.
585    pub fn for_property(&self, name: &str) -> Self {
586        match self.get_property_type(name) {
587            Some(ty) => ContextualTypeContext::with_expected(self.interner, ty),
588            None => ContextualTypeContext::new(self.interner),
589        }
590    }
591
592    /// Create a child context for an array element.
593    pub fn for_array_element(&self) -> Self {
594        match self.get_array_element_type() {
595            Some(ty) => ContextualTypeContext::with_expected(self.interner, ty),
596            None => ContextualTypeContext::new(self.interner),
597        }
598    }
599
600    /// Create a child context for a tuple element at the given index.
601    pub fn for_tuple_element(&self, index: usize) -> Self {
602        match self.get_tuple_element_type(index) {
603            Some(ty) => ContextualTypeContext::with_expected(self.interner, ty),
604            None => ContextualTypeContext::new(self.interner),
605        }
606    }
607
608    /// Create a child context for a function parameter at the given index.
609    pub fn for_parameter(&self, index: usize) -> Self {
610        match self.get_parameter_type(index) {
611            Some(ty) => ContextualTypeContext::with_expected(self.interner, ty),
612            None => ContextualTypeContext::new(self.interner),
613        }
614    }
615
616    /// Create a child context for a function return expression.
617    pub fn for_return(&self) -> Self {
618        match self.get_return_type() {
619            Some(ty) => ContextualTypeContext::with_expected(self.interner, ty),
620            None => ContextualTypeContext::new(self.interner),
621        }
622    }
623
624    /// Get the contextual yield type for a generator function.
625    ///
626    /// If the expected type is `Generator<Y, R, N>`, this returns Y.
627    /// This is used to contextually type yield expressions.
628    ///
629    /// Example:
630    /// ```typescript
631    /// function* gen(): Generator<number, void, unknown> {
632    ///     yield 1;  // 1 is contextually typed as number
633    /// }
634    /// ```
635    pub fn get_generator_yield_type(&self) -> Option<TypeId> {
636        let expected = self.expected?;
637
638        // Handle Union explicitly - collect yield types from all members
639        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
640            let members = self.interner.type_list(members);
641            let yield_types: Vec<TypeId> = members
642                .iter()
643                .filter_map(|&m| {
644                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
645                    ctx.get_generator_yield_type()
646                })
647                .collect();
648
649            return collect_single_or_union(self.interner, yield_types);
650        }
651
652        // Generator<Y, R, N> — yield type is arg 0
653        let mut extractor = ApplicationArgExtractor::new(self.interner, 0);
654        extractor.extract(expected)
655    }
656
657    /// Get the contextual return type for a generator function (`TReturn` from Generator<Y, `TReturn`, N>).
658    ///
659    /// This is used to contextually type return statements in generators.
660    pub fn get_generator_return_type(&self) -> Option<TypeId> {
661        let expected = self.expected?;
662
663        // Handle Union explicitly - collect return types from all members
664        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
665            let members = self.interner.type_list(members);
666            let return_types: Vec<TypeId> = members
667                .iter()
668                .filter_map(|&m| {
669                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
670                    ctx.get_generator_return_type()
671                })
672                .collect();
673
674            return collect_single_or_union(self.interner, return_types);
675        }
676
677        // Generator<Y, R, N> — return type is arg 1
678        let mut extractor = ApplicationArgExtractor::new(self.interner, 1);
679        extractor.extract(expected)
680    }
681
682    /// Get the contextual next type for a generator function (`TNext` from Generator<Y, R, `TNext`>).
683    ///
684    /// This is used to determine the type of values passed to .`next()` and
685    /// the type of the yield expression result.
686    pub fn get_generator_next_type(&self) -> Option<TypeId> {
687        let expected = self.expected?;
688
689        // Handle Union explicitly - collect next types from all members
690        if let Some(TypeData::Union(members)) = self.interner.lookup(expected) {
691            let members = self.interner.type_list(members);
692            let next_types: Vec<TypeId> = members
693                .iter()
694                .filter_map(|&m| {
695                    let ctx = ContextualTypeContext::with_expected(self.interner, m);
696                    ctx.get_generator_next_type()
697                })
698                .collect();
699
700            return collect_single_or_union(self.interner, next_types);
701        }
702
703        // Generator<Y, R, N> — next type is arg 2
704        let mut extractor = ApplicationArgExtractor::new(self.interner, 2);
705        extractor.extract(expected)
706    }
707
708    /// Create a child context for a yield expression in a generator.
709    pub fn for_yield(&self) -> Self {
710        match self.get_generator_yield_type() {
711            Some(ty) => ContextualTypeContext::with_expected(self.interner, ty),
712            None => ContextualTypeContext::new(self.interner),
713        }
714    }
715}
716
717/// Apply contextual type to infer a more specific type.
718///
719/// This implements bidirectional type inference:
720/// 1. If `expr_type` is any/unknown/error, use contextual type
721/// 2. If `expr_type` is a literal and contextual type is a union containing that literal's base type, preserve literal
722/// 3. If `expr_type` is assignable to contextual type and is more specific, use `expr_type`
723/// 4. Otherwise, prefer `expr_type` (don't widen to contextual type)
724pub fn apply_contextual_type(
725    interner: &dyn TypeDatabase,
726    expr_type: TypeId,
727    contextual_type: Option<TypeId>,
728) -> TypeId {
729    let ctx_type = match contextual_type {
730        Some(t) => t,
731        None => return expr_type,
732    };
733
734    // If expression type is any, unknown, or error, use contextual type
735    if expr_type.is_any_or_unknown() || expr_type.is_error() {
736        return ctx_type;
737    }
738
739    // If expression type is the same, just return it
740    if expr_type == ctx_type {
741        return expr_type;
742    }
743
744    // Check if expr_type is a literal type that should be preserved
745    // When contextual type is a union like string | number, we should preserve literal types
746    if let Some(expr_key) = interner.lookup(expr_type) {
747        // Literal types should be preserved when context is a union
748        if matches!(expr_key, TypeData::Literal(_))
749            && let Some(ctx_key) = interner.lookup(ctx_type)
750            && matches!(ctx_key, TypeData::Union(_))
751        {
752            // Preserve the literal type - it's more specific than the union
753            return expr_type;
754        }
755    }
756
757    // PERF: Reuse a single SubtypeChecker across all subtype checks in this function
758    let mut checker = crate::relations::subtype::SubtypeChecker::new(interner);
759
760    // Check if contextual type is a union
761    if let Some(TypeData::Union(members)) = interner.lookup(ctx_type) {
762        let members = interner.type_list(members);
763        // If expr_type is in the union, it's valid - use the more specific expr_type
764        for &member in members.iter() {
765            if member == expr_type {
766                return expr_type;
767            }
768        }
769        // If expr_type is assignable to any union member, use expr_type
770        for &member in members.iter() {
771            checker.reset();
772            if checker.is_subtype_of(expr_type, member) {
773                return expr_type;
774            }
775        }
776    }
777
778    // If expr_type is assignable to contextual type, use expr_type (it's more specific)
779    checker.reset();
780    if checker.is_subtype_of(expr_type, ctx_type) {
781        return expr_type;
782    }
783
784    // If contextual type is assignable to expr_type, use contextual type (it's more specific)
785    // BUT: Skip for function/callable types — the solver's bivariant SubtypeChecker can
786    // incorrectly say that a wider function type (e.g. (x: number|string) => void) is a
787    // subtype of a narrower one (e.g. (x: number) => void), which would widen the property
788    // type and suppress valid TS2322 errors under strict function types.
789    let is_function_type = matches!(
790        interner.lookup(expr_type),
791        Some(TypeData::Function(_) | TypeData::Object(_))
792    );
793    if !is_function_type {
794        checker.reset();
795        if checker.is_subtype_of(ctx_type, expr_type) {
796            return ctx_type;
797        }
798    }
799
800    // Default: prefer the expression type (don't widen to contextual type)
801    // This prevents incorrectly widening concrete types to generic type parameters
802    expr_type
803}
804
805#[cfg(test)]
806#[path = "../../tests/contextual_tests.rs"]
807mod tests;