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}