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