Skip to main content

mago_analyzer/plugin/libraries/stdlib/array/
array_map.rs

1//! `array_map()` return type provider.
2//!
3//! Preserves array shape (known items/elements) through `array_map` calls
4//! by replacing value types with the callback's return type. When the
5//! callback's return type contains conditional types (e.g.
6//! `($s is non-empty-string ? non-empty-lowercase-string : '')`), they are
7//! resolved against the input array's element type so the result is a
8//! concrete type rather than the unresolved conditional itself.
9
10use std::borrow::Cow;
11use std::collections::BTreeMap;
12use std::sync::Arc;
13
14use mago_codex::metadata::CodebaseMetadata;
15use mago_codex::ttype::add_union_type;
16use mago_codex::ttype::atomic::TAtomic;
17use mago_codex::ttype::atomic::array::TArray;
18use mago_codex::ttype::atomic::array::key::ArrayKey;
19use mago_codex::ttype::atomic::array::keyed::TKeyedArray;
20use mago_codex::ttype::atomic::array::list::TList;
21use mago_codex::ttype::combiner::CombinerOptions;
22use mago_codex::ttype::comparator::ComparisonResult;
23use mago_codex::ttype::comparator::union_comparator;
24use mago_codex::ttype::get_mixed;
25use mago_codex::ttype::union::TUnion;
26use mago_codex::ttype::wrap_atomic;
27use mago_word::Word;
28
29use crate::plugin::context::InvocationInfo;
30use crate::plugin::context::ProviderContext;
31use crate::plugin::provider::Provider;
32use crate::plugin::provider::ProviderMeta;
33use crate::plugin::provider::function::FunctionReturnTypeProvider;
34use crate::plugin::provider::function::FunctionTarget;
35
36static META: ProviderMeta =
37    ProviderMeta::new("php::array::array_map", "array_map", "Preserves array shape through array_map");
38
39#[derive(Default)]
40pub struct ArrayMapProvider;
41
42impl Provider for ArrayMapProvider {
43    fn meta() -> &'static ProviderMeta {
44        &META
45    }
46}
47
48impl FunctionReturnTypeProvider for ArrayMapProvider {
49    fn targets() -> FunctionTarget {
50        FunctionTarget::Exact(b"array_map")
51    }
52
53    fn get_return_type(
54        &self,
55        context: &ProviderContext<'_, '_, '_>,
56        invocation: &InvocationInfo<'_, '_, '_>,
57    ) -> Option<TUnion> {
58        // array_map(?callable $callback, array $array, array ...$arrays)
59        let argument_count = invocation.argument_count();
60        if argument_count < 2 {
61            return None;
62        }
63
64        let callback_arg = invocation.get_argument(0, &[b"callback"])?;
65        let callback_type = context.get_expression_type(callback_arg)?;
66        let callback_is_null = callback_type.is_null();
67
68        if callback_is_null && argument_count == 2 {
69            let array_arg = invocation.get_argument(1, &[b"array"])?;
70            return context.get_expression_type(array_arg).cloned();
71        }
72
73        if callback_is_null && argument_count > 2 {
74            return zip_input_arrays(context, invocation, argument_count);
75        }
76
77        if argument_count != 2 {
78            return None;
79        }
80
81        let array_arg = invocation.get_argument(1, &[b"array"])?;
82        let callback_metadata = context.get_callable_metadata(callback_arg);
83
84        // Fall back to the inferred return type when none is declared, e.g. for
85        // an untyped arrow function.
86        let declared_return_type =
87            callback_metadata.and_then(|metadata| metadata.return_type_metadata.as_ref()).map(|r| &r.type_union);
88        let inferred_return_type = if callback_type.is_single()
89            && let TAtomic::Callable(callable) = callback_type.get_single()
90        {
91            callable.get_signature().and_then(|signature| signature.get_return_type())
92        } else {
93            None
94        };
95        let raw_return_type = declared_return_type.or(inferred_return_type)?.clone();
96
97        let array_type = context.get_expression_type(array_arg)?;
98        let array = array_type.get_single_array()?;
99
100        let codebase = context.codebase();
101        // Only a declared return type can reference the callback's parameter.
102        let first_parameter_name = declared_return_type
103            .and(callback_metadata)
104            .and_then(|metadata| metadata.parameters.first())
105            .map(|parameter| parameter.name.0);
106
107        let resolve_for = |element: &TUnion| -> TUnion {
108            match first_parameter_name {
109                Some(parameter_name) => {
110                    resolve_conditionals_in_return(codebase, &raw_return_type, parameter_name, element)
111                }
112                None => raw_return_type.clone(),
113            }
114        };
115
116        match array {
117            TArray::Keyed(keyed) if keyed.has_known_items() => {
118                let known_items = keyed.get_known_items()?;
119                let new_items: BTreeMap<_, _> = known_items
120                    .iter()
121                    .map(|(key, (optional, value))| (*key, (*optional, resolve_for(value))))
122                    .collect();
123
124                let mut result = TKeyedArray::new().with_known_items(new_items).with_non_empty(keyed.is_non_empty());
125
126                if let Some((key_type, value_type)) = keyed.parameters.as_ref() {
127                    result = result.with_parameters(Arc::clone(key_type), Arc::new(resolve_for(value_type)));
128                }
129
130                Some(wrap_atomic(TAtomic::Array(TArray::Keyed(result))))
131            }
132            TArray::List(list) if list.known_elements.is_some() => {
133                let known_elements = list.known_elements.as_ref()?;
134                let new_elements: BTreeMap<_, _> = known_elements
135                    .iter()
136                    .map(|(idx, (optional, value))| (*idx, (*optional, resolve_for(value))))
137                    .collect();
138
139                let result = TList {
140                    element_type: if list.element_type.is_never() {
141                        Arc::clone(&list.element_type)
142                    } else {
143                        Arc::new(resolve_for(&list.element_type))
144                    },
145                    known_elements: Some(new_elements),
146                    known_count: list.known_count,
147                    non_empty: list.non_empty,
148                };
149
150                Some(wrap_atomic(TAtomic::Array(TArray::List(result))))
151            }
152            _ => None,
153        }
154    }
155}
156
157fn resolve_conditionals_in_return(
158    codebase: &CodebaseMetadata,
159    return_type: &TUnion,
160    parameter_name: Word,
161    argument_type: &TUnion,
162) -> TUnion {
163    let mut new_atomics: Vec<TAtomic> = Vec::with_capacity(return_type.types.len());
164    for atomic in return_type.types.as_ref() {
165        new_atomics.extend(resolve_atomic_with_argument(codebase, atomic, parameter_name, argument_type));
166    }
167
168    let mut result = return_type.clone();
169    result.types = Cow::Owned(new_atomics);
170    result
171}
172
173fn resolve_atomic_with_argument(
174    codebase: &CodebaseMetadata,
175    atomic: &TAtomic,
176    parameter_name: Word,
177    argument_type: &TUnion,
178) -> Vec<TAtomic> {
179    let TAtomic::Conditional(conditional) = atomic else {
180        return vec![atomic.clone()];
181    };
182
183    // Substitute the bound parameter inside each branch so a conditional
184    // whose then/otherwise references the parameter (e.g. `T is X ? T : null`)
185    // also reads as the argument type rather than a bare `TVariable`.
186    let then = substitute_parameter_in_union(&conditional.then, parameter_name, argument_type);
187    let otherwise = substitute_parameter_in_union(&conditional.otherwise, parameter_name, argument_type);
188
189    // Recurse to handle conditionals nested inside the branches.
190    let then = resolve_conditionals_in_return(codebase, &then, parameter_name, argument_type);
191    let otherwise = resolve_conditionals_in_return(codebase, &otherwise, parameter_name, argument_type);
192
193    let subject = substitute_parameter_in_union(&conditional.subject, parameter_name, argument_type);
194
195    if subject.is_never() {
196        return add_union_type(then, &otherwise, codebase, CombinerOptions::default()).types.into_owned();
197    }
198
199    let mut comparison_result = ComparisonResult::new();
200    let subject_is_contained = union_comparator::is_contained_by(
201        codebase,
202        &subject,
203        &conditional.target,
204        false,
205        false,
206        true,
207        &mut comparison_result,
208    );
209
210    let are_disjoint =
211        !union_comparator::can_expression_types_be_identical(codebase, &subject, &conditional.target, false, false);
212
213    if are_disjoint {
214        return if conditional.negated { then.types.into_owned() } else { otherwise.types.into_owned() };
215    }
216
217    if subject_is_contained {
218        return if conditional.negated { otherwise.types.into_owned() } else { then.types.into_owned() };
219    }
220
221    add_union_type(then, &otherwise, codebase, CombinerOptions::default()).types.into_owned()
222}
223
224fn substitute_parameter_in_union(union: &TUnion, parameter_name: Word, argument_type: &TUnion) -> TUnion {
225    let mut new_atomics: Vec<TAtomic> = Vec::with_capacity(union.types.len());
226    let mut changed = false;
227
228    for atomic in union.types.as_ref() {
229        if matches!(atomic, TAtomic::Variable(name) if *name == parameter_name) {
230            new_atomics.extend(argument_type.types.as_ref().iter().cloned());
231            changed = true;
232        } else {
233            new_atomics.push(atomic.clone());
234        }
235    }
236
237    if !changed {
238        return union.clone();
239    }
240
241    let mut result = union.clone();
242    result.types = Cow::Owned(new_atomics);
243    result
244}
245
246fn zip_input_arrays(
247    context: &ProviderContext<'_, '_, '_>,
248    invocation: &InvocationInfo<'_, '_, '_>,
249    argument_count: usize,
250) -> Option<TUnion> {
251    let array_count = argument_count - 1;
252    let mut tuple_items: BTreeMap<ArrayKey, (bool, TUnion)> = BTreeMap::new();
253    let mut all_inputs_non_empty = true;
254
255    for offset in 0..array_count {
256        let array_arg = invocation.get_argument(offset + 1, &[])?;
257        let array_type = context.get_expression_type(array_arg)?;
258        let array = array_type.get_single_array()?;
259
260        let value_type = array_value_type(array);
261        all_inputs_non_empty &= match array {
262            TArray::List(list) => list.non_empty,
263            TArray::Keyed(keyed) => keyed.is_non_empty(),
264        };
265
266        tuple_items.insert(ArrayKey::Integer(offset as i64), (false, value_type.as_nullable()));
267    }
268
269    Some(wrap_atomic(TAtomic::Array(TArray::List(TList {
270        element_type: Arc::new(wrap_atomic(TAtomic::Array(TArray::Keyed(
271            TKeyedArray::new().with_known_items(tuple_items).with_non_empty(true),
272        )))),
273        known_elements: None,
274        known_count: None,
275        non_empty: all_inputs_non_empty,
276    }))))
277}
278
279fn array_value_type(array: &TArray) -> TUnion {
280    match array {
281        TArray::List(list) => (*list.element_type).clone(),
282        TArray::Keyed(keyed) => keyed.get_value_type().cloned().unwrap_or_else(get_mixed),
283    }
284}