mago_codex/ttype/
expander.rs

1use mago_interner::StringIdentifier;
2use mago_interner::ThreadedInterner;
3
4use crate::get_class_like;
5use crate::get_closure;
6use crate::get_declaring_method;
7use crate::get_function;
8use crate::identifier::function_like::FunctionLikeIdentifier;
9use crate::metadata::CodebaseMetadata;
10use crate::metadata::function_like::FunctionLikeMetadata;
11use crate::ttype::atomic::TAtomic;
12use crate::ttype::atomic::array::TArray;
13use crate::ttype::atomic::callable::TCallable;
14use crate::ttype::atomic::callable::TCallableSignature;
15use crate::ttype::atomic::callable::parameter::TCallableParameter;
16use crate::ttype::atomic::derived::TDerived;
17use crate::ttype::atomic::derived::key_of::TKeyOf;
18use crate::ttype::atomic::derived::value_of::TValueOf;
19use crate::ttype::atomic::mixed::TMixed;
20use crate::ttype::atomic::object::TObject;
21use crate::ttype::atomic::object::named::TNamedObject;
22use crate::ttype::atomic::reference::TReference;
23use crate::ttype::atomic::reference::TReferenceMemberSelector;
24use crate::ttype::atomic::scalar::TScalar;
25use crate::ttype::atomic::scalar::class_like_string::TClassLikeString;
26use crate::ttype::combiner;
27use crate::ttype::get_mixed;
28use crate::ttype::union::TUnion;
29
30#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Default)]
31pub enum StaticClassType {
32    #[default]
33    None,
34    Name(StringIdentifier),
35    Object(TObject),
36}
37
38#[derive(Debug)]
39pub struct TypeExpansionOptions<'a> {
40    pub self_class: Option<&'a StringIdentifier>,
41    pub static_class_type: StaticClassType,
42    pub parent_class: Option<&'a StringIdentifier>,
43    pub evaluate_class_constants: bool,
44    pub evaluate_conditional_types: bool,
45    pub function_is_final: bool,
46    pub expand_generic: bool,
47    pub expand_templates: bool,
48}
49
50impl Default for TypeExpansionOptions<'_> {
51    fn default() -> Self {
52        Self {
53            self_class: None,
54            static_class_type: StaticClassType::default(),
55            parent_class: None,
56            evaluate_class_constants: true,
57            evaluate_conditional_types: false,
58            function_is_final: false,
59            expand_generic: false,
60            expand_templates: true,
61        }
62    }
63}
64
65pub fn expand_union(
66    codebase: &CodebaseMetadata,
67    interner: &ThreadedInterner,
68    return_type: &mut TUnion,
69    options: &TypeExpansionOptions,
70) {
71    let previous_types = std::mem::take(&mut return_type.types);
72    return_type.types = combiner::combine(previous_types, codebase, interner, false);
73
74    let mut new_return_type_parts = vec![];
75    let mut skipped_keys = vec![];
76
77    for (i, return_type_part) in return_type.types.iter_mut().enumerate() {
78        let mut skip_key = false;
79        expand_atomic(return_type_part, codebase, interner, options, &mut skip_key, &mut new_return_type_parts);
80
81        if skip_key {
82            skipped_keys.push(i);
83        }
84    }
85
86    if !skipped_keys.is_empty() {
87        let mut i = 0;
88
89        return_type.types.retain(|_| {
90            let to_retain = !skipped_keys.contains(&i);
91            i += 1;
92            to_retain
93        });
94
95        new_return_type_parts.append(&mut return_type.types);
96
97        if new_return_type_parts.is_empty() {
98            new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
99        }
100
101        if new_return_type_parts.len() > 1 {
102            return_type.types = combiner::combine(new_return_type_parts, codebase, interner, false)
103        } else {
104            return_type.types = new_return_type_parts;
105        }
106    }
107}
108
109fn expand_atomic(
110    return_type_part: &mut TAtomic,
111    codebase: &CodebaseMetadata,
112    interner: &ThreadedInterner,
113    options: &TypeExpansionOptions,
114    skip_key: &mut bool,
115    new_return_type_parts: &mut Vec<TAtomic>,
116) {
117    match return_type_part {
118        TAtomic::Array(array_type) => match array_type {
119            TArray::Keyed(keyed_data) => {
120                if let Some((key_parameter, value_parameter)) = &mut keyed_data.parameters {
121                    expand_union(codebase, interner, key_parameter, options);
122                    expand_union(codebase, interner, value_parameter, options);
123                }
124
125                if let Some(known_items) = &mut keyed_data.known_items {
126                    for (_, item_type) in known_items.values_mut() {
127                        expand_union(codebase, interner, item_type, options);
128                    }
129                }
130            }
131            TArray::List(list_data) => {
132                expand_union(codebase, interner, &mut list_data.element_type, options);
133
134                if let Some(known_elements) = &mut list_data.known_elements {
135                    for (_, element_type) in known_elements.values_mut() {
136                        expand_union(codebase, interner, element_type, options);
137                    }
138                }
139            }
140        },
141        TAtomic::Object(TObject::Named(named_object)) => {
142            expand_named_object(named_object, codebase, interner, options);
143        }
144        TAtomic::Callable(TCallable::Signature(signature)) => {
145            if let Some(return_type) = signature.get_return_type_mut() {
146                expand_union(codebase, interner, return_type, options);
147            }
148
149            for param in signature.get_parameters_mut() {
150                if let Some(param_type) = param.get_type_signature_mut() {
151                    expand_union(codebase, interner, param_type, options);
152                }
153            }
154        }
155        TAtomic::GenericParameter(parameter) => {
156            expand_union(codebase, interner, parameter.constraint.as_mut(), options);
157        }
158        TAtomic::Scalar(TScalar::ClassLikeString(TClassLikeString::OfType { constraint, .. })) => {
159            let mut atomic_return_type_parts = vec![];
160            expand_atomic(constraint, codebase, interner, options, &mut false, &mut atomic_return_type_parts);
161
162            if !atomic_return_type_parts.is_empty() {
163                *constraint = Box::new(atomic_return_type_parts.remove(0));
164            }
165        }
166        TAtomic::Reference(TReference::Member { class_like_name, member_selector }) => {
167            *skip_key = true;
168
169            match member_selector {
170                TReferenceMemberSelector::Wildcard => {
171                    let Some(class_like) = get_class_like(codebase, interner, class_like_name) else {
172                        new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
173
174                        return;
175                    };
176
177                    for constant in class_like.constants.values() {
178                        if let Some(inferred_type) = constant.inferred_type.as_ref() {
179                            let mut inferred_type = inferred_type.clone();
180
181                            let mut skip_inferred_type = false;
182                            expand_atomic(
183                                &mut inferred_type,
184                                codebase,
185                                interner,
186                                options,
187                                &mut skip_inferred_type,
188                                new_return_type_parts,
189                            );
190
191                            if !skip_inferred_type {
192                                new_return_type_parts.push(inferred_type);
193                            }
194                        } else if let Some(type_metadata) = constant.type_metadata.as_ref() {
195                            let mut constant_type = type_metadata.type_union.clone();
196
197                            expand_union(codebase, interner, &mut constant_type, options);
198
199                            new_return_type_parts.extend(constant_type.types);
200                        } else {
201                            new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
202                        }
203                    }
204
205                    for enum_case_name in class_like.enum_cases.keys() {
206                        new_return_type_parts
207                            .push(TAtomic::Object(TObject::new_enum_case(class_like.original_name, *enum_case_name)));
208                    }
209                }
210                TReferenceMemberSelector::StartsWith(prefix) => {
211                    let Some(class_like) = get_class_like(codebase, interner, class_like_name) else {
212                        new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
213
214                        return;
215                    };
216
217                    let prefix_str = interner.lookup(prefix);
218
219                    for (constant_name, constant) in class_like.constants.iter() {
220                        let constant_name_str = interner.lookup(constant_name);
221
222                        if !constant_name_str.starts_with(prefix_str) {
223                            continue;
224                        }
225
226                        if let Some(inferred_type) = constant.inferred_type.as_ref() {
227                            let mut inferred_type = inferred_type.clone();
228
229                            let mut skip_inferred_type = false;
230                            expand_atomic(
231                                &mut inferred_type,
232                                codebase,
233                                interner,
234                                options,
235                                &mut skip_inferred_type,
236                                new_return_type_parts,
237                            );
238
239                            if !skip_inferred_type {
240                                new_return_type_parts.push(inferred_type);
241                            }
242                        } else if let Some(type_metadata) = constant.type_metadata.as_ref() {
243                            let mut constant_type = type_metadata.type_union.clone();
244
245                            expand_union(codebase, interner, &mut constant_type, options);
246
247                            new_return_type_parts.extend(constant_type.types);
248                        } else {
249                            new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
250                        }
251                    }
252
253                    for enum_case_name in class_like.enum_cases.keys() {
254                        let enum_case_name_str = interner.lookup(enum_case_name);
255
256                        if !enum_case_name_str.starts_with(prefix_str) {
257                            continue;
258                        }
259
260                        new_return_type_parts
261                            .push(TAtomic::Object(TObject::new_enum_case(class_like.original_name, *enum_case_name)));
262                    }
263                }
264                TReferenceMemberSelector::EndsWith(suffix) => {
265                    let Some(class_like) = get_class_like(codebase, interner, class_like_name) else {
266                        new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
267
268                        return;
269                    };
270
271                    let suffix_str = interner.lookup(suffix);
272
273                    for (constant_name, constant) in class_like.constants.iter() {
274                        let constant_name_str = interner.lookup(constant_name);
275
276                        if !constant_name_str.ends_with(suffix_str) {
277                            continue;
278                        }
279
280                        if let Some(inferred_type) = constant.inferred_type.as_ref() {
281                            let mut inferred_type = inferred_type.clone();
282
283                            let mut skip_inferred_type = false;
284                            expand_atomic(
285                                &mut inferred_type,
286                                codebase,
287                                interner,
288                                options,
289                                &mut skip_inferred_type,
290                                new_return_type_parts,
291                            );
292
293                            if !skip_inferred_type {
294                                new_return_type_parts.push(inferred_type);
295                            }
296                        } else if let Some(type_metadata) = constant.type_metadata.as_ref() {
297                            let mut constant_type = type_metadata.type_union.clone();
298
299                            expand_union(codebase, interner, &mut constant_type, options);
300
301                            new_return_type_parts.extend(constant_type.types);
302                        } else {
303                            new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
304                        }
305                    }
306
307                    for enum_case_name in class_like.enum_cases.keys() {
308                        let enum_case_name_str = interner.lookup(enum_case_name);
309
310                        if !enum_case_name_str.ends_with(suffix_str) {
311                            continue;
312                        }
313
314                        new_return_type_parts
315                            .push(TAtomic::Object(TObject::new_enum_case(class_like.original_name, *enum_case_name)));
316                    }
317                }
318                TReferenceMemberSelector::Identifier(member_name) => {
319                    let Some(class_like) = get_class_like(codebase, interner, class_like_name) else {
320                        new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
321
322                        return;
323                    };
324
325                    if class_like.enum_cases.contains_key(member_name) {
326                        new_return_type_parts
327                            .push(TAtomic::Object(TObject::new_enum_case(class_like.original_name, *member_name)));
328                    } else if let Some(constant) = class_like.constants.get(member_name) {
329                        if let Some(inferred_type) = constant.inferred_type.as_ref() {
330                            let mut inferred_type = inferred_type.clone();
331
332                            let mut skip_inferred_type = false;
333                            expand_atomic(
334                                &mut inferred_type,
335                                codebase,
336                                interner,
337                                options,
338                                &mut skip_inferred_type,
339                                new_return_type_parts,
340                            );
341
342                            if !skip_inferred_type {
343                                new_return_type_parts.push(inferred_type);
344                            }
345                        } else if let Some(type_metadata) = constant.type_metadata.as_ref() {
346                            let mut constant_type = type_metadata.type_union.clone();
347
348                            expand_union(codebase, interner, &mut constant_type, options);
349
350                            new_return_type_parts.extend(constant_type.types);
351                        } else {
352                            new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
353                        }
354                    } else {
355                        new_return_type_parts.push(TAtomic::Mixed(TMixed::new()));
356                    }
357                }
358            }
359        }
360        TAtomic::Callable(TCallable::Alias(id)) => {
361            if let Some(value) = get_atomic_of_function_like_identifier(id, codebase, interner) {
362                *skip_key = true;
363                new_return_type_parts.push(value);
364            }
365        }
366        TAtomic::Conditional(conditional) => {
367            *skip_key = true;
368
369            let mut then = conditional.then.clone();
370            let mut otherwise = conditional.otherwise.clone();
371
372            expand_union(codebase, interner, &mut then, options);
373            expand_union(codebase, interner, &mut otherwise, options);
374
375            new_return_type_parts.extend(then.types);
376            new_return_type_parts.extend(otherwise.types);
377        }
378        TAtomic::Derived(derived) => match derived {
379            TDerived::KeyOf(key_of) => {
380                *skip_key = true;
381                new_return_type_parts.extend(expand_key_of(key_of, codebase, interner, options));
382            }
383            TDerived::ValueOf(value_of) => {
384                *skip_key = true;
385                new_return_type_parts.extend(expand_value_of(value_of, codebase, interner, options));
386            }
387            TDerived::PropertiesOf(_) => todo!("expand_properties_of"),
388        },
389        _ => {}
390    }
391}
392
393fn expand_named_object(
394    named_object: &mut TNamedObject,
395    codebase: &CodebaseMetadata,
396    interner: &ThreadedInterner,
397    options: &TypeExpansionOptions,
398) {
399    let name_str_lc = interner.lookup(&named_object.name).to_lowercase();
400
401    if named_object.is_this() || name_str_lc == "static" || name_str_lc == "$this" {
402        match &options.static_class_type {
403            StaticClassType::Object(TObject::Named(static_object)) => {
404                if let Some(static_object_intersections) = &static_object.intersection_types {
405                    let intersections = named_object.intersection_types.get_or_insert_with(Vec::new);
406                    intersections.extend(static_object_intersections.iter().cloned());
407                }
408
409                if named_object.type_parameters.is_none() {
410                    named_object.type_parameters = static_object.type_parameters.clone();
411                }
412
413                named_object.name = static_object.name;
414                named_object.is_this = true;
415            }
416            StaticClassType::Name(static_class_name) => {
417                named_object.name = *static_class_name;
418                named_object.is_this = options.function_is_final;
419            }
420            _ => {}
421        }
422    } else if name_str_lc == "self" {
423        if let Some(self_class_name) = options.self_class {
424            named_object.name = *self_class_name;
425        }
426    } else if name_str_lc == "parent"
427        && let Some(self_class_name) = options.self_class
428        && let Some(class_metadata) = get_class_like(codebase, interner, self_class_name)
429        && let Some(parent_name) = class_metadata.direct_parent_class
430    {
431        named_object.name = parent_name;
432    }
433
434    if named_object.type_parameters.is_none()
435        && let Some(class_like_metadata) = get_class_like(codebase, interner, &named_object.name)
436        && !class_like_metadata.template_types.is_empty()
437    {
438        let default_params: Vec<TUnion> = class_like_metadata
439            .template_types
440            .iter()
441            .map(|(_, template_map)| template_map.iter().map(|(_, t)| t).next().cloned().unwrap_or_else(get_mixed))
442            .collect();
443
444        if !default_params.is_empty() {
445            named_object.type_parameters = Some(default_params);
446        }
447    }
448}
449
450pub fn get_signature_of_function_like_identifier(
451    function_like_identifier: &FunctionLikeIdentifier,
452    codebase: &CodebaseMetadata,
453    interner: &ThreadedInterner,
454) -> Option<TCallableSignature> {
455    Some(match function_like_identifier {
456        FunctionLikeIdentifier::Function(name) => {
457            let function_like_metadata = get_function(codebase, interner, name)?;
458
459            get_signature_of_function_like_metadata(
460                function_like_identifier,
461                function_like_metadata,
462                codebase,
463                interner,
464                &TypeExpansionOptions::default(),
465            )
466        }
467        FunctionLikeIdentifier::Closure(file_id, position) => {
468            let function_like_metadata = get_closure(codebase, interner, file_id, position)?;
469
470            get_signature_of_function_like_metadata(
471                function_like_identifier,
472                function_like_metadata,
473                codebase,
474                interner,
475                &TypeExpansionOptions::default(),
476            )
477        }
478        FunctionLikeIdentifier::Method(classlike_name, method_name) => {
479            let function_like_metadata = get_declaring_method(codebase, interner, classlike_name, method_name)?;
480
481            get_signature_of_function_like_metadata(
482                function_like_identifier,
483                function_like_metadata,
484                codebase,
485                interner,
486                &TypeExpansionOptions {
487                    self_class: Some(classlike_name),
488                    static_class_type: StaticClassType::Name(*classlike_name),
489                    ..Default::default()
490                },
491            )
492        }
493    })
494}
495
496pub fn get_atomic_of_function_like_identifier(
497    function_like_identifier: &FunctionLikeIdentifier,
498    codebase: &CodebaseMetadata,
499    interner: &ThreadedInterner,
500) -> Option<TAtomic> {
501    let signature = get_signature_of_function_like_identifier(function_like_identifier, codebase, interner)?;
502
503    Some(TAtomic::Callable(TCallable::Signature(signature)))
504}
505
506pub fn get_signature_of_function_like_metadata(
507    function_like_identifier: &FunctionLikeIdentifier,
508    function_like_metadata: &FunctionLikeMetadata,
509    codebase: &CodebaseMetadata,
510    interner: &ThreadedInterner,
511    options: &TypeExpansionOptions,
512) -> TCallableSignature {
513    let parameters: Vec<_> = function_like_metadata
514        .parameters
515        .iter()
516        .map(|parameter_metadata| {
517            let type_signature = if let Some(t) = parameter_metadata.get_type_metadata() {
518                let mut t = t.type_union.clone();
519                expand_union(codebase, interner, &mut t, options);
520                Some(Box::new(t))
521            } else {
522                None
523            };
524
525            TCallableParameter::new(
526                type_signature,
527                parameter_metadata.flags.is_by_reference(),
528                parameter_metadata.flags.is_variadic(),
529                parameter_metadata.flags.has_default(),
530            )
531        })
532        .collect();
533
534    let return_type = if let Some(type_metadata) = function_like_metadata.return_type_metadata.as_ref() {
535        let mut return_type = type_metadata.type_union.clone();
536        expand_union(codebase, interner, &mut return_type, options);
537        Some(Box::new(return_type))
538    } else {
539        None
540    };
541
542    let mut signature = TCallableSignature::new(function_like_metadata.flags.is_pure(), true)
543        .with_parameters(parameters)
544        .with_return_type(return_type)
545        .with_source(Some(*function_like_identifier));
546
547    if let FunctionLikeIdentifier::Closure(file_id, closure_position) = function_like_identifier {
548        signature = signature.with_closure_location(Some((*file_id, *closure_position)));
549    }
550
551    signature
552}
553
554fn expand_key_of(
555    return_type_key_of: &TKeyOf,
556    codebase: &CodebaseMetadata,
557    interner: &ThreadedInterner,
558    options: &TypeExpansionOptions,
559) -> Vec<TAtomic> {
560    let mut type_atomics = vec![];
561
562    let mut target_type = return_type_key_of.get_target_type().clone();
563    let mut new_atomics = vec![];
564    let mut remove_target_atomic = false;
565    expand_atomic(&mut target_type, codebase, interner, options, &mut remove_target_atomic, &mut new_atomics);
566    type_atomics.extend(new_atomics);
567    if !remove_target_atomic {
568        type_atomics.push(target_type);
569    }
570
571    let Some(new_return_types) = TKeyOf::get_key_of_targets(type_atomics, codebase, interner, false) else {
572        return vec![TAtomic::Derived(TDerived::KeyOf(return_type_key_of.clone()))];
573    };
574
575    new_return_types.types
576}
577
578fn expand_value_of(
579    return_type_value_of: &TValueOf,
580    codebase: &CodebaseMetadata,
581    interner: &ThreadedInterner,
582    options: &TypeExpansionOptions,
583) -> Vec<TAtomic> {
584    let mut type_atomics = vec![];
585
586    let mut target_type = return_type_value_of.get_target_type().clone();
587    let mut new_atomics = vec![];
588    let mut remove_target_atomic = false;
589    expand_atomic(&mut target_type, codebase, interner, options, &mut remove_target_atomic, &mut new_atomics);
590    type_atomics.extend(new_atomics);
591    if !remove_target_atomic {
592        type_atomics.push(target_type);
593    }
594
595    let Some(new_return_types) = TValueOf::get_value_of_targets(type_atomics, codebase, interner, false) else {
596        return vec![TAtomic::Derived(TDerived::ValueOf(return_type_value_of.clone()))];
597    };
598
599    new_return_types.types
600}