mago_analyzer/plugin/libraries/psl/dict/
select_keys.rs

1//! `Psl\Dict\select_keys()` 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::combine_union_types;
10use mago_codex::ttype::combiner::CombinerOptions;
11use mago_codex::ttype::get_iterable_parameters;
12use mago_codex::ttype::union::TUnion;
13
14use crate::plugin::context::InvocationInfo;
15use crate::plugin::context::ProviderContext;
16use crate::plugin::provider::Provider;
17use crate::plugin::provider::ProviderMeta;
18use crate::plugin::provider::function::FunctionReturnTypeProvider;
19use crate::plugin::provider::function::FunctionTarget;
20
21static META: ProviderMeta = ProviderMeta::new(
22    "psl::dict::select_keys",
23    "Psl\\Dict\\select_keys",
24    "Returns array with only selected keys from input",
25);
26
27/// Provider for the `Psl\Dict\select_keys()` function.
28///
29/// Narrows the return type based on the literal keys in the `$keys` argument:
30/// - If the input is a keyed array with known items AND keys are literals,
31///   returns only the matching known items.
32/// - If the input is a generic iterable/array and keys are literals,
33///   returns a shaped array with optional entries for each key.
34/// - Falls back to `None` (generic return type) otherwise.
35#[derive(Default)]
36pub struct SelectKeysProvider;
37
38impl Provider for SelectKeysProvider {
39    fn meta() -> &'static ProviderMeta {
40        &META
41    }
42}
43
44impl FunctionReturnTypeProvider for SelectKeysProvider {
45    fn targets() -> FunctionTarget {
46        FunctionTarget::Exact(b"psl\\dict\\select_keys")
47    }
48
49    fn get_return_type(
50        &self,
51        context: &ProviderContext<'_, '_, '_>,
52        invocation: &InvocationInfo<'_, '_, '_>,
53    ) -> Option<TUnion> {
54        let iterable_expr = invocation.get_argument(0, &[b"iterable"])?;
55        let keys_expr = invocation.get_argument(1, &[b"keys"])?;
56
57        let iterable_type = context.get_expression_type(iterable_expr)?;
58        let keys_type = context.get_expression_type(keys_expr)?;
59
60        if !keys_type.is_single() {
61            return None;
62        }
63
64        let selected_keys = extract_literal_keys(keys_type.get_single())?;
65        if selected_keys.is_empty() {
66            return None;
67        }
68
69        let codebase = context.codebase();
70
71        let mut known_items: BTreeMap<ArrayKey, (bool, TUnion)> = BTreeMap::new();
72        let mut generic_value_type: Option<TUnion> = None;
73
74        for atomic in iterable_type.types.as_ref() {
75            if let TAtomic::Array(TArray::Keyed(keyed)) = atomic {
76                if let Some(items) = &keyed.known_items {
77                    for key in &selected_keys {
78                        if let Some((optional, value_type)) = items.get(key) {
79                            known_items
80                                .entry(*key)
81                                .and_modify(|(is_optional, existing)| {
82                                    *is_optional = *is_optional || *optional;
83                                    *existing =
84                                        combine_union_types(existing, value_type, codebase, CombinerOptions::default());
85                                })
86                                .or_insert_with(|| (*optional, value_type.clone()));
87                        }
88                    }
89                }
90
91                if let Some((_, value_param)) = &keyed.parameters {
92                    generic_value_type = Some(match generic_value_type {
93                        Some(existing) => {
94                            combine_union_types(&existing, value_param, codebase, CombinerOptions::default())
95                        }
96                        None => (**value_param).clone(),
97                    });
98                }
99
100                continue;
101            }
102
103            if let Some((_, value_type)) = get_iterable_parameters(atomic, codebase) {
104                generic_value_type = Some(match generic_value_type {
105                    Some(existing) => combine_union_types(&existing, &value_type, codebase, CombinerOptions::default()),
106                    None => value_type,
107                });
108            }
109        }
110
111        let mut result_items: BTreeMap<ArrayKey, (bool, TUnion)> = BTreeMap::new();
112        for key in &selected_keys {
113            if let Some(item) = known_items.get(key) {
114                result_items.insert(*key, item.clone());
115            } else if let Some(value_type) = generic_value_type.as_ref() {
116                result_items.insert(*key, (true, value_type.clone()));
117            } else {
118                // key isn't known and there's no fallback value type; drop it from the result
119            }
120        }
121
122        if result_items.is_empty() {
123            return None;
124        }
125
126        let mut result = TKeyedArray::new();
127        result.known_items = Some(result_items);
128        result.non_empty = known_items.values().any(|(optional, _)| !optional);
129
130        Some(TUnion::from_atomic(TAtomic::Array(TArray::Keyed(result))))
131    }
132}
133
134/// Extracts literal key values from a list or keyed array type (the `$keys` argument).
135fn extract_literal_keys(atomic: &TAtomic) -> Option<Vec<ArrayKey>> {
136    match atomic {
137        TAtomic::Array(TArray::List(list)) => {
138            let known_elements = list.known_elements.as_ref()?;
139            let mut keys = Vec::new();
140
141            for (_, element_type) in known_elements.values() {
142                keys.push(union_to_array_key(element_type)?);
143            }
144
145            if keys.is_empty() { None } else { Some(keys) }
146        }
147        TAtomic::Array(TArray::Keyed(keyed)) => {
148            let known_items = keyed.known_items.as_ref()?;
149            let mut keys = Vec::new();
150
151            for (_, value_type) in known_items.values() {
152                keys.push(union_to_array_key(value_type)?);
153            }
154
155            if keys.is_empty() { None } else { Some(keys) }
156        }
157        _ => None,
158    }
159}
160
161/// Converts a union type to an ArrayKey if it's a single literal string or int.
162fn union_to_array_key(union: &TUnion) -> Option<ArrayKey> {
163    if !union.is_single() {
164        return None;
165    }
166
167    let atomic = union.get_single();
168    if let Some(value) = atomic.get_literal_string_value() {
169        Some(ArrayKey::String(mago_word::word(value)))
170    } else {
171        atomic.get_literal_int_value().map(ArrayKey::Integer)
172    }
173}
mago_analyzer/plugin/libraries/psl/dict/select_keys.rs

mago_analyzer/plugin/libraries/psl/dict/
select_keys.rs
