hamelin_lib 0.7.12

//! Map functions

use std::any::TypeId;

use anyhow::bail;

use crate::func::def::{FunctionDef, ParameterBindingProvider, Parameters};
use crate::func::utils::must_be_same_type;
use crate::types::array::Array;
use crate::types::map::Map;
use crate::types::matcher::{AnyMatcher, ArrayMatcher, BaseMatcher, MapMatcher, TupleMatcher};
use crate::types::tuple::Tuple;
use crate::types::{Type, UNKNOWN};

/// map(keys_array, values_array) - Create map from parallel arrays.
#[derive(Default, Clone)]
pub struct MapFromArrays;

impl FunctionDef for MapFromArrays {
    fn name(&self) -> &'static str {
        "map"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new()
            .with("keys", ArrayMatcher::of(BaseMatcher::default()))
            .with("values", ArrayMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        let keys = match bindings.get_by_index(0)?.typ() {
            Type::Array(array) => array,
            _ => bail!("checked by type"),
        };
        let values = match bindings.get_by_index(1)?.typ() {
            Type::Array(array) => array,
            _ => bail!("checked by type"),
        };
        Ok(Map::new(
            keys.element_type.as_ref().clone(),
            values.element_type.as_ref().clone(),
        )
        .into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}

/// map(array_of_pairs) - Create map from array of key-value pairs.
#[derive(Default, Clone)]
pub struct MapFromPairs;

impl FunctionDef for MapFromPairs {
    fn name(&self) -> &'static str {
        "map"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with(
            "elements",
            ArrayMatcher::of(TupleMatcher::of(
                BaseMatcher::default(),
                AnyMatcher::default(),
            )),
        )
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        if let Type::Array(array) = bindings.get_by_index(0)?.typ() {
            if let Type::Tuple(Tuple { elements }) = array.element_type.as_ref() {
                return Ok(Map::new(
                    elements
                        .get(0)
                        .ok_or_else(|| anyhow::anyhow!("checked by type"))?
                        .as_ref()
                        .clone(),
                    elements
                        .get(1)
                        .ok_or_else(|| anyhow::anyhow!("checked by type"))?
                        .as_ref()
                        .clone(),
                )
                .into());
            }
        };

        bail!("checked by type");
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}

/// map() - Create empty map.
#[derive(Default, Clone)]
pub struct MapEmpty;

impl FunctionDef for MapEmpty {
    fn name(&self) -> &'static str {
        "map"
    }

    fn return_type(&self, _bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        Ok(Map::new(UNKNOWN.into(), UNKNOWN.into()).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}

/// map(pair1, pair2, ...) - Create map from key-value pair arguments.
#[derive(Default, Clone)]
pub struct MapFromKeyValue;

impl FunctionDef for MapFromKeyValue {
    fn name(&self) -> &'static str {
        "map"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with_var_args(TupleMatcher::of(
            BaseMatcher::default(),
            AnyMatcher::default(),
        ))
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        let tuples: Vec<_> = bindings
            .iter()
            .map(|t| match t.typ() {
                Type::Tuple(tup) => Ok(tup.clone()),
                other => bail!("Expected tuple type, got: {other}"),
            })
            .collect::<anyhow::Result<_>>()?;
        let key_type = must_be_same_type(tuples.iter().map(|tup| tup.elements[0].as_ref()))?;
        let value_type = must_be_same_type(tuples.iter().map(|tup| tup.elements[1].as_ref()))?;

        Ok(Map::new(key_type.clone(), value_type.clone()).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}

/// map_keys(map) - Get array of map keys.
#[derive(Default, Clone)]
pub struct MapKeys;

impl FunctionDef for MapKeys {
    fn name(&self) -> &'static str {
        "map_keys"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("map", MapMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        let key_type = match bindings.get_by_index(0)?.typ() {
            Type::Map(m) => m.key_type.as_ref(),
            _ => bail!("expected map type"),
        };

        Ok(Array::new(key_type.clone()).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}

/// map_values(map) - Get array of map values.
#[derive(Default, Clone)]
pub struct MapValues;

impl FunctionDef for MapValues {
    fn name(&self) -> &'static str {
        "map_values"
    }

    fn parameters(&self) -> Parameters {
        Parameters::new().with("map", MapMatcher::default())
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        let value_type = match bindings.get_by_index(0)?.typ() {
            Type::Map(m) => m.value_type.as_ref(),
            _ => bail!("expected map type"),
        };

        Ok(Array::new(value_type.clone()).into())
    }

    fn type_id(&self) -> TypeId {
        TypeId::of::<Self>()
    }
}