Skip to main content

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

1//! `array_column()` return type provider.
2
3use std::borrow::Cow;
4use std::sync::Arc;
5
6use mago_codex::metadata::class_like::ClassLikeMetadata;
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::atomic::object::TObject;
13use mago_codex::ttype::atomic::scalar::TScalar;
14use mago_codex::ttype::get_array_parameters;
15use mago_codex::ttype::union::TUnion;
16use mago_word::concat_word;
17use mago_word::word;
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 = ProviderMeta::new(
27    "php::array::array_column",
28    "array_column",
29    "Returns list or array based on column_key and index_key arguments",
30);
31
32/// Provider for the `array_column()` function.
33///
34/// Returns typed arrays based on the `column_key` and `index_key` arguments.
35#[derive(Default)]
36pub struct ArrayColumnProvider;
37
38impl Provider for ArrayColumnProvider {
39    fn meta() -> &'static ProviderMeta {
40        &META
41    }
42}
43
44impl FunctionReturnTypeProvider for ArrayColumnProvider {
45    fn targets() -> FunctionTarget {
46        FunctionTarget::Exact(b"array_column")
47    }
48
49    fn get_return_type(
50        &self,
51        context: &ProviderContext<'_, '_, '_>,
52        invocation: &InvocationInfo<'_, '_, '_>,
53    ) -> Option<TUnion> {
54        let array_argument = invocation.get_argument(0, &[b"array"])?;
55        let array_type = context.get_expression_type(array_argument)?;
56
57        let array = array_type.get_single_array()?;
58        let codebase = context.codebase();
59        let element_type = get_array_parameters(array, codebase).1;
60
61        let column_key_argument = invocation.get_argument(1, &[b"column_key"])?;
62        let column_key_type = context.get_expression_type(column_key_argument)?;
63
64        let index_key_argument = invocation.get_argument(2, &[b"index_key"]);
65        let index_key_type = index_key_argument.and_then(|arg| context.get_expression_type(arg));
66
67        if let Some(result) = try_resolve_from_named_object(&element_type, column_key_type, index_key_type, codebase) {
68            return Some(result);
69        }
70
71        if let Some(result) = try_resolve_from_keyed_array(&element_type, column_key_type, index_key_type) {
72            return Some(result);
73        }
74
75        None
76    }
77}
78
79/// Resolve column and index types from an object element type by looking up
80/// class properties.
81fn try_resolve_from_named_object(
82    element_type: &TUnion,
83    column_key_type: &TUnion,
84    index_key_type: Option<&TUnion>,
85    codebase: &mago_codex::metadata::CodebaseMetadata,
86) -> Option<TUnion> {
87    let obj = element_type.get_single_named_object()?;
88    let class_like = codebase.get_class_like(obj.name.as_bytes())?;
89
90    let column_type = if column_key_type.is_null() {
91        TUnion::from_atomic(TAtomic::Object(TObject::Named(obj.clone())))
92    } else {
93        let prop_name = column_key_type.get_single_literal_string_value()?;
94        let prop = class_like.properties.get(&concat_word!(b"$", prop_name))?;
95        prop.type_metadata.as_ref()?.type_union.clone()
96    };
97
98    let index_type = resolve_index_type_from_property(index_key_type, class_like);
99
100    Some(build_result(column_type, index_type))
101}
102
103/// Resolve column and index types from a keyed-array element type by looking
104/// up known items.
105fn try_resolve_from_keyed_array(
106    element_type: &TUnion,
107    column_key_type: &TUnion,
108    index_key_type: Option<&TUnion>,
109) -> Option<TUnion> {
110    if !element_type.is_single() {
111        return None;
112    }
113
114    let TAtomic::Array(TArray::Keyed(keyed)) = element_type.get_single() else {
115        return None;
116    };
117
118    let known_items = keyed.get_known_items()?;
119
120    let column_type = if column_key_type.is_null() {
121        element_type.clone()
122    } else {
123        let key_str = column_key_type.get_single_literal_string_value()?;
124        let (_, value_type) = known_items.get(&ArrayKey::String(word(key_str)))?;
125        value_type.clone()
126    };
127
128    let index_type = if let Some(index_key_type) = index_key_type {
129        if index_key_type.is_null() {
130            None
131        } else {
132            let key_str = index_key_type.get_single_literal_string_value()?;
133            let (_, value_type) = known_items.get(&ArrayKey::String(word(key_str)))?;
134            extract_scalar_for_key(value_type)
135        }
136    } else {
137        None
138    };
139
140    Some(build_result(column_type, index_type))
141}
142
143/// Try to extract the key scalar type from a value type (for use as array index).
144fn extract_scalar_for_key(value_type: &TUnion) -> Option<&TScalar> {
145    if !value_type.is_single() {
146        return None;
147    }
148
149    match value_type.get_single() {
150        TAtomic::Scalar(
151            scalar @ (TScalar::ArrayKey | TScalar::Integer(_) | TScalar::String(_) | TScalar::ClassLikeString(_)),
152        ) => Some(scalar),
153        _ => None,
154    }
155}
156
157fn resolve_index_type_from_property<'meta>(
158    index_key_type: Option<&TUnion>,
159    class_like: &'meta ClassLikeMetadata,
160) -> Option<&'meta TScalar> {
161    let index_key_type = index_key_type?;
162    if index_key_type.is_null() {
163        return None;
164    }
165
166    let prop_name = index_key_type.get_single_literal_string_value()?;
167    let prop = class_like.properties.get(&concat_word!("$", prop_name))?;
168    let prop_type = &prop.type_metadata.as_ref()?.type_union;
169
170    extract_scalar_for_key(prop_type)
171}
172
173fn build_result(column_type: TUnion, index_type: Option<&TScalar>) -> TUnion {
174    if let Some(index_scalar) = index_type {
175        let keyed_array = TKeyedArray::new_with_parameters(
176            Arc::new(TUnion::from_atomic(TAtomic::Scalar(index_scalar.clone()))),
177            Arc::new(column_type),
178        );
179
180        TUnion::from_atomic(TAtomic::Array(TArray::Keyed(keyed_array)))
181    } else {
182        let list = TList::new(Arc::new(column_type));
183
184        TUnion::from_single(Cow::Owned(TAtomic::Array(TArray::List(list))))
185    }
186}