Skip to main content

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

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