hamelin_lib 0.21.5

//! Map functions

use std::any::TypeId;
use std::sync::Arc;

use anyhow::bail;

use crate::func::def::{FunctionDef, ParameterBinding, ParameterBindingProvider, Parameters};
use crate::func::utils::must_be_same_type;
use crate::tree::ast::expression::{Expression, ExpressionKind};
use crate::tree::typed_ast::context::ExpressionTranslationContext;
use crate::tree::typed_ast::expression::TypedExpression;
use crate::types::array::Array;
use crate::types::map::Map;
use crate::types::matcher::{
    AnyMatcher, ArrayMatcher, BaseMatcher, ExactMatcher, MapKeyMatcher, MapMatcher, TupleMatcher,
};
use crate::types::tuple::Tuple;
use crate::types::{Type, STRING, UNKNOWN, VARIANT};

/// 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(MapKeyMatcher::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(
                MapKeyMatcher::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(
            MapKeyMatcher::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>()
    }
}

/// transform_values(map, lambda) - Apply a lambda to every value of a map.
///
/// The lambda receives each value and returns a new value. Keys are preserved.
/// The result is a map with the original keys and the lambda's return values.
///
/// Example: `transform_values(map("a":1, "b":2), v -> v * 10)` returns
/// `map("a":10, "b":20)`.
#[derive(Default, Clone)]
pub struct TransformValues;

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

    fn parameters(&self) -> Parameters {
        // AnyMatcher used for arity checking only - custom_bind validates types
        Parameters::new()
            .with("map", AnyMatcher)
            .with("lambda", AnyMatcher)
    }

    fn custom_bind(
        &self,
        ast_binding: &ParameterBinding<Arc<Expression>>,
        ctx: &mut ExpressionTranslationContext,
    ) -> anyhow::Result<Option<ParameterBinding<Arc<TypedExpression>>>> {
        let map_ast = ast_binding.get_by_name("map")?.clone();
        let typed_map = Arc::new(TypedExpression::from_ast_with_context(map_ast, ctx));

        let value_type = match typed_map.resolved_type.as_ref() {
            Type::Map(m) => m.value_type.clone(),
            other => bail!("Expected map type, got {}", other),
        };

        let lambda_ast = ast_binding.get_by_name("lambda")?.clone();

        let ExpressionKind::Lambda(lambda) = &lambda_ast.kind else {
            bail!("Expected lambda expression");
        };

        let lambda_with_hints = lambda.with_param_types(&[value_type])?;

        let lambda_expr = Expression {
            kind: lambda_with_hints.into(),
            span: lambda_ast.span.clone(),
        };
        let typed_lambda = Arc::new(TypedExpression::from_ast_with_context(
            Arc::new(lambda_expr),
            ctx,
        ));

        Ok(Some(ParameterBinding::from_named([
            ("map".to_string(), typed_map),
            ("lambda".to_string(), typed_lambda),
        ])))
    }

    fn return_type(&self, bindings: &dyn ParameterBindingProvider) -> anyhow::Result<Type> {
        let map_type = bindings.get_by_name("map")?.typ().clone();
        let key_type = match &map_type {
            Type::Map(m) => m.key_type.as_ref().clone(),
            Type::Unknown => return Ok(Type::Unknown),
            other => bail!("Expected map type for 'map', got {}", other),
        };

        let lambda_arg = bindings.get_by_name("lambda")?;
        let return_type = match lambda_arg.typ() {
            Type::Function(f) => f.return_type.as_ref().clone(),
            Type::Unknown => Type::Unknown,
            other => bail!("Expected function type for lambda, got {}", other),
        };

        Ok(Map::new(key_type, return_type).into())
    }

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

/// map_variant_get_values(map(K, variant), path) - Extract a variant field from every value of a map.
///
/// Internal function used by transform_values lowering. Applies `variant_get` element-wise
/// to the value column of a map, avoiding the explode/reaggregate pipeline.
#[derive(Default, Clone)]
pub struct MapVariantGetValues;

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

    fn parameters(&self) -> Parameters {
        Parameters::new()
            .with(
                "map",
                MapMatcher::of(BaseMatcher, ExactMatcher::of(VARIANT)),
            )
            .with("path", ExactMatcher::of(STRING))
    }

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

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

/// map_variant_to_json_values(map(K, variant)) - Convert every variant value of a map to its JSON string.
///
/// Internal function used by transform_values lowering. Applies `to_json_string` element-wise
/// to the value column of a map, avoiding the explode/reaggregate pipeline.
#[derive(Default, Clone)]
pub struct MapVariantToJsonValues;

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

    fn parameters(&self) -> Parameters {
        Parameters::new().with(
            "map",
            MapMatcher::of(BaseMatcher, ExactMatcher::of(VARIANT)),
        )
    }

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

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