mago_codex/ttype/
expander.rs

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