mago_analyzer/plugin/libraries/psl/regex/
capture_groups.rs1use 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#[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}