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