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