tsz_solver/evaluation/evaluate_rules/
keyof.rs

1//! keyof operator evaluation.
2//!
3//! Handles TypeScript's keyof operator: `keyof T`
4
5use crate::TypeDatabase;
6use crate::relations::subtype::TypeResolver;
7use crate::types::{IntrinsicKind, LiteralValue, TupleElement, TypeData, TypeId, TypeListId};
8use rustc_hash::FxHashSet;
9use tsz_common::interner::Atom;
10
11use super::super::evaluate::{
12    ARRAY_METHODS_RETURN_ANY, ARRAY_METHODS_RETURN_BOOLEAN, ARRAY_METHODS_RETURN_NUMBER,
13    ARRAY_METHODS_RETURN_STRING, ARRAY_METHODS_RETURN_VOID, TypeEvaluator,
14};
15
16/// Tracks the types of keys found during keyof evaluation.
17pub(crate) struct KeyofKeySet {
18    pub has_string: bool,
19    pub has_number: bool,
20    pub has_symbol: bool,
21    pub string_literals: FxHashSet<Atom>,
22}
23
24impl KeyofKeySet {
25    pub fn new() -> Self {
26        Self {
27            has_string: false,
28            has_number: false,
29            has_symbol: false,
30            string_literals: FxHashSet::default(),
31        }
32    }
33
34    pub fn insert_type(&mut self, interner: &dyn TypeDatabase, type_id: TypeId) -> bool {
35        let Some(key) = interner.lookup(type_id) else {
36            return false;
37        };
38
39        match key {
40            TypeData::Union(members) => {
41                let members = interner.type_list(members);
42                members
43                    .iter()
44                    .all(|&member| self.insert_type(interner, member))
45            }
46            TypeData::Intrinsic(kind) => match kind {
47                IntrinsicKind::String => {
48                    self.has_string = true;
49                    true
50                }
51                IntrinsicKind::Number => {
52                    self.has_number = true;
53                    true
54                }
55                IntrinsicKind::Symbol => {
56                    self.has_symbol = true;
57                    true
58                }
59                IntrinsicKind::Never => true,
60                _ => false,
61            },
62            TypeData::Literal(LiteralValue::String(atom)) => {
63                self.string_literals.insert(atom);
64                true
65            }
66            _ => false,
67        }
68    }
69}
70
71impl<'a, R: TypeResolver> TypeEvaluator<'a, R> {
72    /// Helper to recursively evaluate keyof while respecting depth limits.
73    /// Creates a `KeyOf` type and evaluates it through the main `evaluate()` method.
74    fn recurse_keyof(&mut self, operand: TypeId) -> TypeId {
75        let keyof = self.interner().keyof(operand);
76        self.evaluate(keyof)
77    }
78
79    /// Evaluate keyof T - extract the keys of an object type
80    pub fn evaluate_keyof(&mut self, operand: TypeId) -> TypeId {
81        // CRITICAL: Handle TemplateLiteral BEFORE Union to avoid incorrect intersection.
82        // Template literals that expand to unions should return apparent keys of string,
83        // not the intersection of individual literal keys.
84        if let Some(TypeData::TemplateLiteral(_)) = self.interner().lookup(operand) {
85            return self.apparent_primitive_keyof(IntrinsicKind::String);
86        }
87
88        // CRITICAL: Handle Union types BEFORE general evaluation to avoid union simplification.
89        // keyof (A | B) = keyof A & keyof B (distributive contravariance)
90        // If we call evaluate(operand) first, unions get simplified and we lose members.
91        // See test_keyof_union_string_index_and_literal_narrows
92        if let Some(TypeData::Union(members)) = self.interner().lookup(operand) {
93            let member_list = self.interner().type_list(members);
94
95            // Recursively compute keyof for each member (this resolves Lazy/Ref/etc.)
96            let mut key_types: Vec<TypeId> = Vec::with_capacity(member_list.len());
97            for &member in member_list.iter() {
98                key_types.push(self.recurse_keyof(member));
99            }
100
101            // keyof (A | B) = keyof A & keyof B - compute intersection of all key sets
102            // Prefer explicit key-set intersection to avoid opaque literal intersections
103            if let Some(intersection) = self.intersect_keyof_sets(&key_types) {
104                intersection
105            } else {
106                // Fallback: use general intersection
107                self.interner().intersection(key_types)
108            }
109        } else {
110            // For non-union types, evaluate normally
111            let evaluated_operand = self.evaluate(operand);
112
113            let key = match self.interner().lookup(evaluated_operand) {
114                Some(k) => k,
115                None => return TypeId::NEVER,
116            };
117
118            match key {
119                TypeData::Mapped(mapped_id) => {
120                    let mapped = self.interner().mapped_type(mapped_id);
121                    if let Some(name_type) = mapped.name_type {
122                        self.evaluate(name_type)
123                    } else {
124                        self.evaluate(mapped.constraint)
125                    }
126                }
127                TypeData::ReadonlyType(inner) => self.recurse_keyof(inner),
128                TypeData::TypeQuery(sym) => {
129                    // Resolve typeof query before computing keyof
130                    let resolved = self.resolver().resolve_symbol_ref(sym, self.interner());
131                    if let Some(resolved) = resolved {
132                        self.recurse_keyof(resolved)
133                    } else {
134                        TypeId::ERROR
135                    }
136                }
137                TypeData::TypeParameter(param) | TypeData::Infer(param) => {
138                    if let Some(constraint) = param.constraint {
139                        if constraint == evaluated_operand {
140                            self.interner().keyof(operand)
141                        } else {
142                            self.recurse_keyof(constraint)
143                        }
144                    } else {
145                        self.interner().keyof(operand)
146                    }
147                }
148                TypeData::Object(shape_id) => {
149                    let shape = self.interner().object_shape(shape_id);
150                    if shape.properties.is_empty() {
151                        return TypeId::NEVER;
152                    }
153                    let key_types: Vec<TypeId> = shape
154                        .properties
155                        .iter()
156                        .map(|p| self.interner().literal_string_atom(p.name))
157                        .collect();
158                    self.interner().union(key_types)
159                }
160                TypeData::ObjectWithIndex(shape_id) => {
161                    let shape = self.interner().object_shape(shape_id);
162                    let mut key_types: Vec<TypeId> = shape
163                        .properties
164                        .iter()
165                        .map(|p| self.interner().literal_string_atom(p.name))
166                        .collect();
167
168                    if shape.string_index.is_some() {
169                        key_types.push(TypeId::STRING);
170                        key_types.push(TypeId::NUMBER);
171                    } else if shape.number_index.is_some() {
172                        key_types.push(TypeId::NUMBER);
173                    }
174
175                    if key_types.is_empty() {
176                        TypeId::NEVER
177                    } else {
178                        self.interner().union(key_types)
179                    }
180                }
181                TypeData::Array(_) => self.interner().union(self.array_keyof_keys()),
182                TypeData::Tuple(elements) => {
183                    let elements = self.interner().tuple_list(elements);
184                    let mut key_types: Vec<TypeId> = Vec::new();
185                    self.append_tuple_indices(&elements, 0, &mut key_types);
186                    let mut array_keys = self.array_keyof_keys();
187                    key_types.append(&mut array_keys);
188                    if key_types.is_empty() {
189                        return TypeId::NEVER;
190                    }
191                    self.interner().union(key_types)
192                }
193                TypeData::Intrinsic(kind) => match kind {
194                    IntrinsicKind::Any => {
195                        // keyof any = string | number | symbol
196                        self.interner()
197                            .union3(TypeId::STRING, TypeId::NUMBER, TypeId::SYMBOL)
198                    }
199                    IntrinsicKind::Unknown => {
200                        // keyof unknown = never
201                        TypeId::NEVER
202                    }
203                    IntrinsicKind::Never
204                    | IntrinsicKind::Void
205                    | IntrinsicKind::Null
206                    | IntrinsicKind::Undefined
207                    | IntrinsicKind::Object
208                    | IntrinsicKind::Function => TypeId::NEVER,
209                    IntrinsicKind::String
210                    | IntrinsicKind::Number
211                    | IntrinsicKind::Boolean
212                    | IntrinsicKind::Bigint
213                    | IntrinsicKind::Symbol => self.apparent_primitive_keyof(kind),
214                },
215                TypeData::Literal(literal) => {
216                    if let Some(kind) = self.apparent_literal_kind(&literal) {
217                        self.apparent_primitive_keyof(kind)
218                    } else {
219                        self.interner().keyof(operand)
220                    }
221                }
222                TypeData::TemplateLiteral(_) => {
223                    self.apparent_primitive_keyof(IntrinsicKind::String)
224                }
225                // NOTE: Union is handled at the top of this function to avoid union simplification
226                TypeData::Intersection(members) => {
227                    // keyof (A & B) = keyof A | keyof B (covariance)
228                    self.keyof_intersection(members, operand)
229                }
230                // CRITICAL: Handle Lazy (type aliases) by attempting resolution via resolver
231                TypeData::Lazy(def_id) => {
232                    match self.resolver().resolve_lazy(def_id, self.interner()) {
233                        Some(resolved) => {
234                            // Recursively compute keyof of the resolved type
235                            self.recurse_keyof(resolved)
236                        }
237                        None => {
238                            // Keep as deferred KeyOf if resolution fails
239                            self.interner().keyof(operand)
240                        }
241                    }
242                }
243                // CRITICAL: Handle Application (generic types) by evaluating them first
244                TypeData::Application(_app_id) => {
245                    // Evaluate the application to get the instantiated type
246                    let evaluated = self.evaluate(evaluated_operand);
247                    // Then compute keyof of the evaluated result
248                    self.recurse_keyof(evaluated)
249                }
250                // For other types (type parameters, etc.), keep as KeyOf (deferred)
251                _ => self.interner().keyof(operand),
252            }
253        }
254    }
255
256    /// Compute keyof for an intersection type: keyof (A & B) = keyof A | keyof B
257    pub(crate) fn keyof_intersection(&mut self, members: TypeListId, _operand: TypeId) -> TypeId {
258        let members = self.interner().type_list(members);
259        // Use recurse_keyof to respect depth limits
260        // Use loop instead of closure to allow mutable self access
261        let mut key_sets: Vec<TypeId> = Vec::with_capacity(members.len());
262        for &m in members.iter() {
263            key_sets.push(self.recurse_keyof(m));
264        }
265        self.interner().union(key_sets)
266    }
267
268    /// Get the keyof keys for an array type (includes all array methods and number index).
269    pub(crate) fn array_keyof_keys(&self) -> Vec<TypeId> {
270        let mut keys = Vec::new();
271        keys.push(TypeId::NUMBER);
272        keys.push(self.interner().literal_string("length"));
273        for &name in ARRAY_METHODS_RETURN_ANY {
274            keys.push(self.interner().literal_string(name));
275        }
276        for &name in ARRAY_METHODS_RETURN_BOOLEAN {
277            keys.push(self.interner().literal_string(name));
278        }
279        for &name in ARRAY_METHODS_RETURN_NUMBER {
280            keys.push(self.interner().literal_string(name));
281        }
282        for &name in ARRAY_METHODS_RETURN_VOID {
283            keys.push(self.interner().literal_string(name));
284        }
285        for &name in ARRAY_METHODS_RETURN_STRING {
286            keys.push(self.interner().literal_string(name));
287        }
288        keys
289    }
290
291    /// Append tuple indices as string literal keys to the output vector.
292    /// Returns the next index to use, or None if a rest element prevents fixed indexing.
293    pub(crate) fn append_tuple_indices(
294        &self,
295        elements: &[TupleElement],
296        base: usize,
297        out: &mut Vec<TypeId>,
298    ) -> Option<usize> {
299        let mut index = base;
300
301        for element in elements {
302            if element.rest {
303                match self.interner().lookup(element.type_id) {
304                    Some(TypeData::Tuple(rest_elements)) => {
305                        let rest_elements = self.interner().tuple_list(rest_elements);
306                        match self.append_tuple_indices(&rest_elements, index, out) {
307                            Some(next) => {
308                                index = next;
309                                continue;
310                            }
311                            None => return None,
312                        }
313                    }
314                    _ => return None,
315                }
316            }
317            out.push(self.interner().literal_string(&index.to_string()));
318            index += 1;
319        }
320
321        Some(index)
322    }
323
324    /// Compute the intersection of multiple keyof key sets.
325    /// Returns None if the intersection cannot be computed (e.g., non-literal keys).
326    pub(crate) fn intersect_keyof_sets(&self, key_sets: &[TypeId]) -> Option<TypeId> {
327        let mut parsed_sets = Vec::with_capacity(key_sets.len());
328        for &key_set in key_sets {
329            let mut parsed = KeyofKeySet::new();
330            if !parsed.insert_type(self.interner(), key_set) {
331                return None;
332            }
333            parsed_sets.push(parsed);
334        }
335
336        let mut all_string = true;
337        let mut string_possible = true;
338        let mut common_literals: Option<FxHashSet<Atom>> = None;
339        let mut all_number = true;
340        let mut all_symbol = true;
341
342        for set in &parsed_sets {
343            if set.has_string {
344                // string index signatures don't restrict literal key overlap
345            } else {
346                all_string = false;
347                if set.string_literals.is_empty() {
348                    string_possible = false;
349                } else {
350                    common_literals = Some(match common_literals {
351                        Some(mut existing) => {
352                            existing.retain(|atom| set.string_literals.contains(atom));
353                            existing
354                        }
355                        None => set.string_literals.clone(),
356                    });
357                }
358            }
359
360            if !set.has_number {
361                all_number = false;
362            }
363            if !set.has_symbol {
364                all_symbol = false;
365            }
366        }
367
368        let mut result_keys = Vec::new();
369        if string_possible {
370            if all_string {
371                result_keys.push(TypeId::STRING);
372            } else if let Some(common) = common_literals {
373                for atom in common {
374                    result_keys.push(self.interner().literal_string_atom(atom));
375                }
376            }
377        }
378        if all_number {
379            result_keys.push(TypeId::NUMBER);
380        }
381        if all_symbol {
382            result_keys.push(TypeId::SYMBOL);
383        }
384
385        if result_keys.is_empty() {
386            Some(TypeId::NEVER)
387        } else if result_keys.len() == 1 {
388            Some(result_keys[0])
389        } else {
390            Some(self.interner().union(result_keys))
391        }
392    }
393}
tsz_solver/evaluation/evaluate_rules/keyof.rs

tsz_solver/evaluation/evaluate_rules/
keyof.rs
