Skip to main content

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

1//! `array_column()` return type provider.
2
3use std::borrow::Cow;
4
5use mago_atom::concat_atom;
6use mago_codex::ttype::atomic::TAtomic;
7use mago_codex::ttype::atomic::array::TArray;
8use mago_codex::ttype::atomic::array::keyed::TKeyedArray;
9use mago_codex::ttype::atomic::array::list::TList;
10use mago_codex::ttype::atomic::object::TObject;
11use mago_codex::ttype::atomic::scalar::TScalar;
12use mago_codex::ttype::get_array_parameters;
13use mago_codex::ttype::union::TUnion;
14
15use crate::plugin::context::InvocationInfo;
16use crate::plugin::context::ProviderContext;
17use crate::plugin::provider::Provider;
18use crate::plugin::provider::ProviderMeta;
19use crate::plugin::provider::function::FunctionReturnTypeProvider;
20use crate::plugin::provider::function::FunctionTarget;
21
22static META: ProviderMeta = ProviderMeta::new(
23    "php::array::array_column",
24    "array_column",
25    "Returns list or array based on column_key and index_key arguments",
26);
27
28/// Provider for the `array_column()` function.
29///
30/// Returns typed arrays based on the `column_key` and `index_key` arguments.
31#[derive(Default)]
32pub struct ArrayColumnProvider;
33
34impl Provider for ArrayColumnProvider {
35    fn meta() -> &'static ProviderMeta {
36        &META
37    }
38}
39
40impl FunctionReturnTypeProvider for ArrayColumnProvider {
41    fn targets() -> FunctionTarget {
42        FunctionTarget::Exact("array_column")
43    }
44
45    fn get_return_type(
46        &self,
47        context: &ProviderContext<'_, '_, '_>,
48        invocation: &InvocationInfo<'_, '_, '_>,
49    ) -> Option<TUnion> {
50        let array_argument = invocation.get_argument(0, &["array"])?;
51        let array_type = context.get_expression_type(array_argument)?;
52
53        let array = array_type.get_single_array()?;
54        let codebase = context.codebase();
55
56        let array_parameters = get_array_parameters(array, codebase);
57        let obj = array_parameters.1.get_single_named_object()?;
58
59        let class_like = codebase.get_class_like(&obj.name)?;
60
61        let column_key_argument = invocation.get_argument(1, &["column_key"])?;
62        let column_key_type = context.get_expression_type(column_key_argument)?;
63
64        let column_type = if column_key_type.is_null() {
65            TUnion::from_atomic(TAtomic::Object(TObject::Named(obj.clone())))
66        } else {
67            let column_key_property_name = column_key_type.get_single_literal_string_value()?;
68            let column_key_property = class_like.properties.get(&concat_atom!("$", column_key_property_name))?;
69
70            column_key_property.type_metadata.as_ref()?.type_union.clone()
71        };
72
73        let index_key_argument = invocation.get_argument(2, &["index_key"]);
74        let index_key_type = index_key_argument.and_then(|argument| context.get_expression_type(argument));
75
76        let mut index_type = None;
77        if let Some(index_key_type) = index_key_type {
78            let index_key_property_name = index_key_type.get_single_literal_string_value();
79            let index_key_property = index_key_property_name
80                .and_then(|property_name| class_like.properties.get(&concat_atom!("$", property_name)));
81
82            if let Some(index_key_property) = index_key_property {
83                index_type = match index_key_property.type_metadata.as_ref()?.type_union.get_single() {
84                    TAtomic::Scalar(
85                        scalar @ (TScalar::ArrayKey
86                        | TScalar::Integer(_)
87                        | TScalar::String(_)
88                        | TScalar::ClassLikeString(_)),
89                    ) => Some(scalar),
90                    _ => None,
91                };
92            }
93        }
94
95        if let Some(index_type) = index_type {
96            let keyed_array = TKeyedArray::new_with_parameters(
97                Box::new(TUnion::from_atomic(TAtomic::Scalar(index_type.clone()))),
98                Box::new(column_type),
99            );
100
101            return Some(TUnion::from_atomic(TAtomic::Array(TArray::Keyed(keyed_array))));
102        }
103
104        let list = TList::new(Box::new(column_type));
105
106        Some(TUnion::from_single(Cow::Owned(TAtomic::Array(TArray::List(list)))))
107    }
108}