Skip to main content

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

1//! `array_merge()` return type provider.
2
3use std::borrow::Cow;
4use std::collections::BTreeMap;
5use std::sync::Arc;
6
7use mago_codex::ttype::atomic::TAtomic;
8use mago_codex::ttype::atomic::array::TArray;
9use mago_codex::ttype::atomic::array::key::ArrayKey;
10use mago_codex::ttype::atomic::array::keyed::TKeyedArray;
11use mago_codex::ttype::atomic::array::list::TList;
12use mago_codex::ttype::combine_union_types;
13use mago_codex::ttype::combiner::CombinerOptions;
14use mago_codex::ttype::get_array_parameters;
15use mago_codex::ttype::get_arraykey;
16use mago_codex::ttype::get_int;
17use mago_codex::ttype::get_iterable_parameters;
18use mago_codex::ttype::get_mixed;
19use mago_codex::ttype::get_never;
20use mago_codex::ttype::union::TUnion;
21
22use crate::plugin::context::InvocationInfo;
23use crate::plugin::context::ProviderContext;
24use crate::plugin::provider::Provider;
25use crate::plugin::provider::ProviderMeta;
26use crate::plugin::provider::function::FunctionReturnTypeProvider;
27use crate::plugin::provider::function::FunctionTarget;
28
29static META: ProviderMeta =
30    ProviderMeta::new("php::array::array_merge", "array_merge", "Returns merged array with combined types");
31
32static TARGETS: [&[u8]; 2] = [b"array_merge", b"psl\\dict\\merge"];
33
34/// Provider for the `array_merge()` and `Psl\Dict\merge()` functions.
35///
36/// Returns an array with types combined from all input arrays.
37#[derive(Default)]
38pub struct ArrayMergeProvider;
39
40impl Provider for ArrayMergeProvider {
41    fn meta() -> &'static ProviderMeta {
42        &META
43    }
44}
45
46impl FunctionReturnTypeProvider for ArrayMergeProvider {
47    fn targets() -> FunctionTarget {
48        FunctionTarget::ExactMultiple(&TARGETS)
49    }
50
51    fn get_return_type(
52        &self,
53        context: &ProviderContext<'_, '_, '_>,
54        invocation: &InvocationInfo<'_, '_, '_>,
55    ) -> Option<TUnion> {
56        let arguments = invocation.arguments();
57        if arguments.is_empty() {
58            return None;
59        }
60
61        let codebase = context.codebase();
62
63        let mut merged_items: BTreeMap<ArrayKey, (bool, TUnion)> = BTreeMap::new();
64        let mut merged_list_elements: BTreeMap<usize, (bool, TUnion)> = BTreeMap::new();
65        let mut next_list_index: usize = 0;
66        let mut has_parameters = false;
67        let mut merged_key_type: Option<TUnion> = None;
68        let mut merged_value_type: Option<TUnion> = None;
69        let mut any_argument_non_empty = false;
70        let mut all_arguments_are_lists = true;
71        let mut all_lists_are_closed = true;
72
73        for invocation_argument in arguments {
74            let argument_expr = invocation_argument.value()?;
75            let argument_type = context.get_expression_type(argument_expr)?;
76
77            if !argument_type.is_single() {
78                return None;
79            }
80
81            let argument_type = if invocation_argument.is_unpacked() {
82                let inner = get_iterable_parameters(argument_type.get_single(), codebase);
83                let (_, value_type) = inner?;
84                if !value_type.is_single() {
85                    return None;
86                }
87
88                if argument_type.get_single().is_non_empty_list() || argument_type.get_single().is_non_empty_array() {
89                    any_argument_non_empty = true;
90                }
91
92                Cow::Owned(value_type)
93            } else {
94                Cow::Borrowed(argument_type)
95            };
96
97            let iterable = argument_type.get_single();
98
99            if let TAtomic::Array(array) = iterable {
100                match array {
101                    TArray::Keyed(keyed) => {
102                        let is_empty_array = keyed.known_items.is_none() && keyed.parameters.is_none();
103
104                        if !is_empty_array {
105                            all_arguments_are_lists = false;
106                        }
107
108                        if keyed.non_empty {
109                            any_argument_non_empty = true;
110                        }
111
112                        if let Some(items) = keyed.known_items.as_ref() {
113                            for (key, value) in items {
114                                merged_items.insert(*key, value.clone());
115                            }
116                        }
117
118                        if let Some((key_type, value_type)) = &keyed.parameters {
119                            has_parameters = true;
120                            merged_key_type = Some(match merged_key_type {
121                                Some(existing) => {
122                                    combine_union_types(&existing, key_type, codebase, CombinerOptions::default())
123                                }
124                                None => (**key_type).clone(),
125                            });
126                            merged_value_type = Some(match merged_value_type {
127                                Some(existing) => {
128                                    combine_union_types(&existing, value_type, codebase, CombinerOptions::default())
129                                }
130                                None => (**value_type).clone(),
131                            });
132                        }
133                    }
134                    TArray::List(list) => {
135                        if list.non_empty {
136                            any_argument_non_empty = true;
137                        }
138
139                        let is_list_closed = list.element_type.is_never();
140                        if !is_list_closed {
141                            all_lists_are_closed = false;
142                        }
143
144                        if let Some(known_elements) = list.known_elements.as_ref() {
145                            for (idx, (optional, element_type)) in known_elements {
146                                let new_idx = next_list_index + idx;
147                                merged_list_elements.insert(new_idx, (*optional, element_type.clone()));
148                            }
149                            if let Some(max_idx) = known_elements.keys().max() {
150                                next_list_index += max_idx + 1;
151                            }
152                        } else if list.non_empty {
153                            next_list_index += 1; // At least one element
154                        } else {
155                            // list has no known elements and may be empty; leave the running index untouched
156                        }
157
158                        let (_, list_value_type) = get_array_parameters(&TArray::List(list.clone()), codebase);
159
160                        has_parameters = true;
161                        merged_value_type = Some(match merged_value_type {
162                            Some(existing) => {
163                                combine_union_types(&existing, &list_value_type, codebase, CombinerOptions::default())
164                            }
165                            None => list_value_type,
166                        });
167
168                        if !all_arguments_are_lists {
169                            let key_type = get_int();
170                            merged_key_type = Some(match merged_key_type {
171                                Some(existing) => {
172                                    combine_union_types(&existing, &key_type, codebase, CombinerOptions::default())
173                                }
174                                None => key_type,
175                            });
176                        }
177                    }
178                }
179            } else if let Some((iterable_key, iterable_value)) = get_iterable_parameters(iterable, codebase) {
180                all_arguments_are_lists = false;
181                has_parameters = true;
182                merged_key_type = Some(match merged_key_type {
183                    Some(existing) => {
184                        combine_union_types(&existing, &iterable_key, codebase, CombinerOptions::default())
185                    }
186                    None => iterable_key,
187                });
188                merged_value_type = Some(match merged_value_type {
189                    Some(existing) => {
190                        combine_union_types(&existing, &iterable_value, codebase, CombinerOptions::default())
191                    }
192                    None => iterable_value,
193                });
194            } else {
195                return None;
196            }
197        }
198
199        if all_arguments_are_lists {
200            let element_type =
201                if all_lists_are_closed { get_never() } else { merged_value_type.unwrap_or_else(get_mixed) };
202
203            let mut result_list = TList::new(Arc::new(element_type));
204            result_list.non_empty = any_argument_non_empty;
205
206            if !merged_list_elements.is_empty() {
207                result_list.known_elements = Some(merged_list_elements);
208            }
209
210            Some(TUnion::from_atomic(TAtomic::Array(TArray::List(result_list))))
211        } else {
212            let mut result_array = TKeyedArray::new();
213
214            let has_merged_items = !merged_items.is_empty();
215            if has_merged_items {
216                result_array.known_items = Some(merged_items);
217            }
218
219            result_array.non_empty = any_argument_non_empty || has_merged_items;
220
221            if has_parameters {
222                result_array.parameters = Some((
223                    Arc::new(merged_key_type.unwrap_or_else(get_arraykey)),
224                    Arc::new(merged_value_type.unwrap_or_else(get_mixed)),
225                ));
226            }
227
228            Some(TUnion::from_atomic(TAtomic::Array(TArray::Keyed(result_array))))
229        }
230    }
231}