mago-analyzer 1.10.0

//! `Psl\Regex\capture_groups()` return type provider.

use std::collections::BTreeMap;
use std::sync::Arc;

use mago_atom::atom;
use mago_codex::ttype::atomic::TAtomic;
use mago_codex::ttype::atomic::array::TArray;
use mago_codex::ttype::atomic::array::key::ArrayKey;
use mago_codex::ttype::atomic::array::keyed::TKeyedArray;
use mago_codex::ttype::atomic::object::TObject;
use mago_codex::ttype::atomic::object::named::TNamedObject;
use mago_codex::ttype::get_arraykey;
use mago_codex::ttype::get_string;
use mago_codex::ttype::union::TUnion;

use crate::plugin::context::InvocationInfo;
use crate::plugin::context::ProviderContext;
use crate::plugin::provider::Provider;
use crate::plugin::provider::ProviderMeta;
use crate::plugin::provider::function::FunctionReturnTypeProvider;
use crate::plugin::provider::function::FunctionTarget;

static META: ProviderMeta = ProviderMeta::new(
    "psl::regex::capture_groups",
    "Psl\\Regex\\capture_groups",
    "Returns TypeInterface with capture group array shape",
);

/// Provider for the `Psl\Regex\capture_groups()` function.
///
/// Returns a `TypeInterface` with an array type that has keys for each capture group.
#[derive(Default)]
pub struct CaptureGroupsProvider;

impl Provider for CaptureGroupsProvider {
    fn meta() -> &'static ProviderMeta {
        &META
    }
}

impl FunctionReturnTypeProvider for CaptureGroupsProvider {
    fn targets() -> FunctionTarget {
        FunctionTarget::Exact("psl\\regex\\capture_groups")
    }

    fn get_return_type(
        &self,
        context: &ProviderContext<'_, '_, '_>,
        invocation: &InvocationInfo<'_, '_, '_>,
    ) -> Option<TUnion> {
        let Some(groups) = invocation.get_argument(0, &["groups"]) else {
            return Some(capture_groups_fallback_type());
        };

        let Some(groups_type) = context.get_expression_type(groups) else {
            return Some(capture_groups_fallback_type());
        };

        let Some(array_atomic) = groups_type.get_single_array() else {
            return Some(capture_groups_fallback_type());
        };

        let mut known_items = BTreeMap::from([(ArrayKey::Integer(0), (false, get_string()))]);

        let has_extra = match array_atomic {
            TArray::Keyed(keyed_array) => {
                let Some(groups_known_items) = keyed_array.known_items.as_ref() else {
                    return Some(capture_groups_fallback_type());
                };

                let mut has_unknown = false;
                for (optional, group_known_item) in groups_known_items.values() {
                    let Some(key) = group_known_item.get_single_array_key() else {
                        has_unknown = true;
                        continue;
                    };

                    known_items.insert(key, (*optional, get_string()));
                }

                has_unknown || keyed_array.parameters.is_some()
            }
            TArray::List(list) => {
                let Some(groups_known_elements) = list.known_elements.as_ref() else {
                    return Some(capture_groups_fallback_type());
                };

                let mut has_unknown = false;
                for (optional, groups_known_element) in groups_known_elements.values() {
                    let Some(key) = groups_known_element.get_single_array_key() else {
                        has_unknown = true;
                        continue;
                    };

                    known_items.insert(key, (*optional, get_string()));
                }

                has_unknown || !list.element_type.is_never()
            }
        };

        Some(TUnion::from_atomic(TAtomic::Object(TObject::Named(TNamedObject::new_with_type_parameters(
            atom("Psl\\Type\\TypeInterface"),
            Some(vec![TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray {
                parameters: if has_extra { Some((Arc::new(get_arraykey()), Arc::new(get_string()))) } else { None },
                non_empty: true,
                known_items: Some(known_items),
            })))]),
        )))))
    }
}

fn capture_groups_fallback_type() -> TUnion {
    TUnion::from_atomic(TAtomic::Object(TObject::Named(TNamedObject::new_with_type_parameters(
        atom("Psl\\Type\\TypeInterface"),
        Some(vec![TUnion::from_atomic(TAtomic::Array(TArray::Keyed(TKeyedArray::new_with_parameters(
            Arc::new(get_arraykey()),
            Arc::new(get_string()),
        ))))]),
    ))))
}