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

1//! `Psl\Async\concurrently()` 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::callable::TCallable;
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::concurrently",
24    "Psl\\Async\\concurrently",
25    "Extracts return types from closure array values, preserving array shape",
26);
27
28/// Provider for the `Psl\Async\concurrently()` function.
29///
30/// Transforms `array<K, (Closure(): V)>` → `array<K, V>`, preserving
31/// sealed array shapes, list structure, and non-empty status.
32#[derive(Default)]
33pub struct ConcurrentlyProvider;
34
35impl Provider for ConcurrentlyProvider {
36    fn meta() -> &'static ProviderMeta {
37        &META
38    }
39}
40
41impl FunctionReturnTypeProvider for ConcurrentlyProvider {
42    fn targets() -> FunctionTarget {
43        FunctionTarget::Exact(b"psl\\async\\concurrently")
44    }
45
46    fn get_return_type(
47        &self,
48        context: &ProviderContext<'_, '_, '_>,
49        invocation: &InvocationInfo<'_, '_, '_>,
50    ) -> Option<TUnion> {
51        let tasks_arg = invocation.get_argument(0, &[b"tasks"])?;
52        let tasks_type = context.get_expression_type(tasks_arg)?;
53
54        let array = tasks_type.get_single_array()?;
55
56        unwrap_closure_return_array(array, context)
57    }
58}
59
60/// Extracts closure return types from each element in an array type,
61/// preserving array shape.
62fn unwrap_closure_return_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                extract_closure_return_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 = extract_closure_return_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 = extract_closure_return_type(item_type)?;
89                    new_items.insert(*key, (*possibly_undefined, inner));
90                }
91
92                return Some(TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray {
93                    parameters: keyed.parameters.as_ref().map(|(k, v)| {
94                        let unwrapped_v = extract_closure_return_type(v).unwrap_or_else(|| (**v).clone());
95                        (Arc::clone(k), Arc::new(unwrapped_v))
96                    }),
97                    non_empty: keyed.non_empty,
98                    known_items: Some(new_items),
99                }))));
100            }
101
102            let (key_type, value_type) = get_array_parameters(array, context.codebase());
103            let inner = extract_closure_return_type(&value_type)?;
104
105            Some(TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray {
106                parameters: Some((Arc::new(key_type), Arc::new(inner))),
107                non_empty: keyed.non_empty,
108                known_items: None,
109            }))))
110        }
111    }
112}
113
114/// Extracts the return type from a `Closure(): V` type.
115fn extract_closure_return_type(union: &TUnion) -> Option<TUnion> {
116    let mut result_types = Vec::new();
117
118    for atomic in union.types.as_ref() {
119        match atomic {
120            TAtomic::Callable(TCallable::Signature(sig)) => {
121                let return_type = sig.get_return_type()?;
122                result_types.extend(return_type.types.iter().cloned());
123            }
124            _ => {
125                return None;
126            }
127        }
128    }
129
130    if result_types.is_empty() {
131        return None;
132    }
133
134    Some(TUnion::from_vec(result_types))
135}
mago_analyzer/plugin/libraries/psl/async_/concurrently.rs

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