Skip to main content

mago_analyzer/plugin/libraries/psl/regex/
capture_groups.rs

1//! `Psl\Regex\capture_groups()` 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::key::ArrayKey;
9use mago_codex::ttype::atomic::array::keyed::TKeyedArray;
10use mago_codex::ttype::atomic::object::TObject;
11use mago_codex::ttype::atomic::object::named::TNamedObject;
12use mago_codex::ttype::get_arraykey;
13use mago_codex::ttype::get_string;
14use mago_codex::ttype::union::TUnion;
15use mago_word::word;
16
17use crate::plugin::context::InvocationInfo;
18use crate::plugin::context::ProviderContext;
19use crate::plugin::provider::Provider;
20use crate::plugin::provider::ProviderMeta;
21use crate::plugin::provider::function::FunctionReturnTypeProvider;
22use crate::plugin::provider::function::FunctionTarget;
23
24static META: ProviderMeta = ProviderMeta::new(
25    "psl::regex::capture_groups",
26    "Psl\\Regex\\capture_groups",
27    "Returns TypeInterface with capture group array shape",
28);
29
30/// Provider for the `Psl\Regex\capture_groups()` function.
31///
32/// Returns a `TypeInterface` with an array type that has keys for each capture group.
33#[derive(Default)]
34pub struct CaptureGroupsProvider;
35
36impl Provider for CaptureGroupsProvider {
37    fn meta() -> &'static ProviderMeta {
38        &META
39    }
40}
41
42impl FunctionReturnTypeProvider for CaptureGroupsProvider {
43    fn targets() -> FunctionTarget {
44        FunctionTarget::Exact(b"psl\\regex\\capture_groups")
45    }
46
47    fn get_return_type(
48        &self,
49        context: &ProviderContext<'_, '_, '_>,
50        invocation: &InvocationInfo<'_, '_, '_>,
51    ) -> Option<TUnion> {
52        let Some(groups) = invocation.get_argument(0, &[b"groups"]) else {
53            return Some(capture_groups_fallback_type());
54        };
55
56        let Some(groups_type) = context.get_expression_type(groups) else {
57            return Some(capture_groups_fallback_type());
58        };
59
60        let Some(array_atomic) = groups_type.get_single_array() else {
61            return Some(capture_groups_fallback_type());
62        };
63
64        let mut known_items = BTreeMap::from([(ArrayKey::Integer(0), (false, get_string()))]);
65
66        let has_extra = match array_atomic {
67            TArray::Keyed(keyed_array) => {
68                let Some(groups_known_items) = keyed_array.known_items.as_ref() else {
69                    return Some(capture_groups_fallback_type());
70                };
71
72                let mut has_unknown = false;
73                for (optional, group_known_item) in groups_known_items.values() {
74                    let Some(key) = group_known_item.get_single_array_key() else {
75                        has_unknown = true;
76                        continue;
77                    };
78
79                    known_items.insert(key, (*optional, get_string()));
80                }
81
82                has_unknown || keyed_array.parameters.is_some()
83            }
84            TArray::List(list) => {
85                let Some(groups_known_elements) = list.known_elements.as_ref() else {
86                    return Some(capture_groups_fallback_type());
87                };
88
89                let mut has_unknown = false;
90                for (optional, groups_known_element) in groups_known_elements.values() {
91                    let Some(key) = groups_known_element.get_single_array_key() else {
92                        has_unknown = true;
93                        continue;
94                    };
95
96                    known_items.insert(key, (*optional, get_string()));
97                }
98
99                has_unknown || !list.element_type.is_never()
100            }
101        };
102
103        Some(TUnion::from_atomic(TAtomic::Object(TObject::Named(TNamedObject::new_with_type_parameters(
104            word("Psl\\Type\\TypeInterface"),
105            Some(vec![TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray {
106                parameters: if has_extra { Some((Arc::new(get_arraykey()), Arc::new(get_string()))) } else { None },
107                non_empty: true,
108                known_items: Some(known_items),
109            })))]),
110        )))))
111    }
112}
113
114fn capture_groups_fallback_type() -> TUnion {
115    TUnion::from_atomic(TAtomic::Object(TObject::Named(TNamedObject::new_with_type_parameters(
116        word("Psl\\Type\\TypeInterface"),
117        Some(vec![TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray::new_with_parameters(
118            Arc::new(get_arraykey()),
119            Arc::new(get_string()),
120        ))))]),
121    ))))
122}