mago_analyzer/plugin/libraries/psl/async_/
all.rs

1//! `Psl\Async\all()` return type provider.
2
3use std::collections::BTreeMap;
4use std::sync::Arc;
5
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::get_array_parameters;
12use mago_codex::ttype::get_never;
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    "psl::async::all",
24    "Psl\\Async\\all",
25    "Unwraps Awaitable types from array values, preserving array shape",
26);
27
28/// Provider for the `Psl\Async\all()` function.
29///
30/// Transforms `array<K, Awaitable<V>>` → `array<K, V>`, preserving
31/// sealed array shapes, list structure, and non-empty status.
32#[derive(Default)]
33pub struct AllProvider;
34
35impl Provider for AllProvider {
36    fn meta() -> &'static ProviderMeta {
37        &META
38    }
39}
40
41impl FunctionReturnTypeProvider for AllProvider {
42    fn targets() -> FunctionTarget {
43        FunctionTarget::Exact(b"psl\\async\\all")
44    }
45
46    fn get_return_type(
47        &self,
48        context: &ProviderContext<'_, '_, '_>,
49        invocation: &InvocationInfo<'_, '_, '_>,
50    ) -> Option<TUnion> {
51        let awaitables_arg = invocation.get_argument(0, &[b"awaitables"])?;
52        let awaitables_type = context.get_expression_type(awaitables_arg)?;
53
54        let array = awaitables_type.get_single_array()?;
55
56        unwrap_awaitable_array(array, context)
57    }
58}
59
60/// Unwraps `Awaitable<V>` to `V` for each element in an array type,
61/// preserving array shape.
62pub(super) fn unwrap_awaitable_array(array: &TArray, context: &ProviderContext<'_, '_, '_>) -> Option<TUnion> {
63    match array {
64        TArray::List(list) => Some(TUnion::from_atomic(TAtomic::Array(TArray::List(TList {
65            element_type: Arc::new(if list.element_type.is_never() {
66                get_never()
67            } else {
68                unwrap_awaitable_type(&list.element_type)?
69            }),
70            known_count: list.known_count,
71            non_empty: list.non_empty,
72            known_elements: if let Some(known_elements) = &list.known_elements {
73                let mut new_elements = BTreeMap::new();
74                for (index, (possibly_undefined, element_type)) in known_elements {
75                    let inner = unwrap_awaitable_type(element_type)?;
76                    new_elements.insert(*index, (*possibly_undefined, inner));
77                }
78
79                Some(new_elements)
80            } else {
81                None
82            },
83        })))),
84        TArray::Keyed(keyed) => {
85            if let Some(known_items) = &keyed.known_items {
86                let mut new_items = BTreeMap::new();
87                for (key, (possibly_undefined, item_type)) in known_items {
88                    let inner = unwrap_awaitable_type(item_type)?;
89                    new_items.insert(*key, (*possibly_undefined, inner));
90                }
91
92                let parameters = if let Some((key_type, value_type)) = &keyed.parameters {
93                    Some((Arc::clone(key_type), Arc::new(unwrap_awaitable_type(value_type)?)))
94                } else {
95                    None
96                };
97
98                return Some(TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray {
99                    parameters,
100                    non_empty: keyed.non_empty,
101                    known_items: Some(new_items),
102                }))));
103            }
104
105            let (key_type, value_type) = get_array_parameters(array, context.codebase());
106            let inner = unwrap_awaitable_type(&value_type)?;
107
108            Some(TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray {
109                parameters: Some((Arc::new(key_type), Arc::new(inner))),
110                non_empty: keyed.non_empty,
111                known_items: None,
112            }))))
113        }
114    }
115}
116
117/// Extracts `V` from `Awaitable<V>`.
118///
119/// Looks for a named object whose name ends with `awaitable` (case-insensitive)
120/// and extracts the first type parameter.
121fn unwrap_awaitable_type(union: &TUnion) -> Option<TUnion> {
122    let mut result_types = Vec::new();
123    let mut found_awaitable = false;
124
125    for atomic in union.types.as_ref() {
126        match atomic {
127            TAtomic::Object(TObject::Named(named))
128                if {
129                    let name_bytes = named.name.as_bytes();
130                    let stripped = name_bytes.strip_prefix(b"\\").unwrap_or(name_bytes);
131                    stripped.eq_ignore_ascii_case(b"Psl\\Async\\Awaitable")
132                } =>
133            {
134                found_awaitable = true;
135
136                if let Some(type_params) = named.type_parameters.as_ref()
137                    && let Some(inner) = type_params.first()
138                {
139                    result_types.extend(inner.types.iter().cloned());
140                } else {
141                    return None;
142                }
143            }
144            _ => {
145                result_types.push(atomic.clone());
146            }
147        }
148    }
149
150    if !found_awaitable || result_types.is_empty() {
151        return None;
152    }
153
154    Some(TUnion::from_vec(result_types))
155}
mago_analyzer/plugin/libraries/psl/async_/all.rs

mago_analyzer/plugin/libraries/psl/async_/
all.rs
