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