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        let raw_return_type = &callback_metadata.return_type_metadata.as_ref()?.type_union;
84
85        let array_type = context.get_expression_type(array_arg)?;
86        let array = array_type.get_single_array()?;
87
88        let codebase = context.codebase();
89        let first_parameter_name = callback_metadata.parameters.first().map(|parameter| parameter.name.0);
90
91        let resolve_for = |element: &TUnion| -> TUnion {
92            match first_parameter_name {
93                Some(parameter_name) => {
94                    resolve_conditionals_in_return(codebase, raw_return_type, parameter_name, element)
95                }
96                None => raw_return_type.clone(),
97            }
98        };
99
100        match array {
101            TArray::Keyed(keyed) if keyed.has_known_items() => {
102                let known_items = keyed.get_known_items()?;
103                let new_items: BTreeMap<_, _> = known_items
104                    .iter()
105                    .map(|(key, (optional, value))| (*key, (*optional, resolve_for(value))))
106                    .collect();
107
108                let mut result = TKeyedArray::new().with_known_items(new_items).with_non_empty(keyed.is_non_empty());
109
110                if let Some((key_type, value_type)) = keyed.parameters.as_ref() {
111                    result = result.with_parameters(Arc::clone(key_type), Arc::new(resolve_for(value_type)));
112                }
113
114                Some(wrap_atomic(TAtomic::Array(TArray::Keyed(result))))
115            }
116            TArray::List(list) if list.known_elements.is_some() => {
117                let known_elements = list.known_elements.as_ref()?;
118                let new_elements: BTreeMap<_, _> = known_elements
119                    .iter()
120                    .map(|(idx, (optional, value))| (*idx, (*optional, resolve_for(value))))
121                    .collect();
122
123                let result = TList {
124                    element_type: if list.element_type.is_never() {
125                        Arc::clone(&list.element_type)
126                    } else {
127                        Arc::new(resolve_for(&list.element_type))
128                    },
129                    known_elements: Some(new_elements),
130                    known_count: list.known_count,
131                    non_empty: list.non_empty,
132                };
133
134                Some(wrap_atomic(TAtomic::Array(TArray::List(result))))
135            }
136            _ => None,
137        }
138    }
139}
140
141fn resolve_conditionals_in_return(
142    codebase: &CodebaseMetadata,
143    return_type: &TUnion,
144    parameter_name: Word,
145    argument_type: &TUnion,
146) -> TUnion {
147    let mut new_atomics: Vec<TAtomic> = Vec::with_capacity(return_type.types.len());
148    for atomic in return_type.types.as_ref() {
149        new_atomics.extend(resolve_atomic_with_argument(codebase, atomic, parameter_name, argument_type));
150    }
151
152    let mut result = return_type.clone();
153    result.types = Cow::Owned(new_atomics);
154    result
155}
156
157fn resolve_atomic_with_argument(
158    codebase: &CodebaseMetadata,
159    atomic: &TAtomic,
160    parameter_name: Word,
161    argument_type: &TUnion,
162) -> Vec<TAtomic> {
163    let TAtomic::Conditional(conditional) = atomic else {
164        return vec![atomic.clone()];
165    };
166
167    // Substitute the bound parameter inside each branch so a conditional
168    // whose then/otherwise references the parameter (e.g. `T is X ? T : null`)
169    // also reads as the argument type rather than a bare `TVariable`.
170    let then = substitute_parameter_in_union(&conditional.then, parameter_name, argument_type);
171    let otherwise = substitute_parameter_in_union(&conditional.otherwise, parameter_name, argument_type);
172
173    // Recurse to handle conditionals nested inside the branches.
174    let then = resolve_conditionals_in_return(codebase, &then, parameter_name, argument_type);
175    let otherwise = resolve_conditionals_in_return(codebase, &otherwise, parameter_name, argument_type);
176
177    let subject = substitute_parameter_in_union(&conditional.subject, parameter_name, argument_type);
178
179    if subject.is_never() {
180        return add_union_type(then, &otherwise, codebase, CombinerOptions::default()).types.into_owned();
181    }
182
183    let mut comparison_result = ComparisonResult::new();
184    let subject_is_contained = union_comparator::is_contained_by(
185        codebase,
186        &subject,
187        &conditional.target,
188        false,
189        false,
190        true,
191        &mut comparison_result,
192    );
193
194    let are_disjoint =
195        !union_comparator::can_expression_types_be_identical(codebase, &subject, &conditional.target, false, false);
196
197    if are_disjoint {
198        return if conditional.negated { then.types.into_owned() } else { otherwise.types.into_owned() };
199    }
200
201    if subject_is_contained {
202        return if conditional.negated { otherwise.types.into_owned() } else { then.types.into_owned() };
203    }
204
205    add_union_type(then, &otherwise, codebase, CombinerOptions::default()).types.into_owned()
206}
207
208fn substitute_parameter_in_union(union: &TUnion, parameter_name: Word, argument_type: &TUnion) -> TUnion {
209    let mut new_atomics: Vec<TAtomic> = Vec::with_capacity(union.types.len());
210    let mut changed = false;
211
212    for atomic in union.types.as_ref() {
213        if matches!(atomic, TAtomic::Variable(name) if *name == parameter_name) {
214            new_atomics.extend(argument_type.types.as_ref().iter().cloned());
215            changed = true;
216        } else {
217            new_atomics.push(atomic.clone());
218        }
219    }
220
221    if !changed {
222        return union.clone();
223    }
224
225    let mut result = union.clone();
226    result.types = Cow::Owned(new_atomics);
227    result
228}
229
230fn zip_input_arrays(
231    context: &ProviderContext<'_, '_, '_>,
232    invocation: &InvocationInfo<'_, '_, '_>,
233    argument_count: usize,
234) -> Option<TUnion> {
235    let array_count = argument_count - 1;
236    let mut tuple_items: BTreeMap<ArrayKey, (bool, TUnion)> = BTreeMap::new();
237    let mut all_inputs_non_empty = true;
238
239    for offset in 0..array_count {
240        let array_arg = invocation.get_argument(offset + 1, &[])?;
241        let array_type = context.get_expression_type(array_arg)?;
242        let array = array_type.get_single_array()?;
243
244        let value_type = array_value_type(array);
245        all_inputs_non_empty &= match array {
246            TArray::List(list) => list.non_empty,
247            TArray::Keyed(keyed) => keyed.is_non_empty(),
248        };
249
250        tuple_items.insert(ArrayKey::Integer(offset as i64), (false, value_type.as_nullable()));
251    }
252
253    Some(wrap_atomic(TAtomic::Array(TArray::List(TList {
254        element_type: Arc::new(wrap_atomic(TAtomic::Array(TArray::Keyed(
255            TKeyedArray::new().with_known_items(tuple_items).with_non_empty(true),
256        )))),
257        known_elements: None,
258        known_count: None,
259        non_empty: all_inputs_non_empty,
260    }))))
261}
262
263fn array_value_type(array: &TArray) -> TUnion {
264    match array {
265        TArray::List(list) => (*list.element_type).clone(),
266        TArray::Keyed(keyed) => keyed.get_value_type().cloned().unwrap_or_else(get_mixed),
267    }
268}