linkrs_query/optimizer/factorization/

group_dependency_analyzer.rs

use std::collections::{HashMap, HashSet};

use linkrs_core::function::{is_list_lambda, list_lambda_arg_index};
use linkrs_core::types::expr::ExpressionId;

use crate::planning::plan::factorization::FactorizedSchema;

/// Analysis result for group dependencies.
#[derive(Debug, Clone, Default)]
pub struct GroupDependencyAnalysis {
    pub dependent_groups: HashSet<u32>,
    pub required_flat_groups: HashSet<u32>,
    pub dependent_exprs: HashSet<ExpressionId>,
}

/// Analyzer that walks an expression tree to collect which factorization
/// groups the expression depends on.
///
/// Mirrors `lbug::planner::GroupDependencyAnalyzer` in
/// `ref/ladybug/src/planner/operator/factorization/flatten_resolver.cpp`.
pub struct GroupDependencyAnalyzer<'a> {
    schema: &'a FactorizedSchema,
    collect_dependent_expr: bool,
    dependent_groups: HashSet<u32>,
    required_flat_groups: HashSet<u32>,
    dependent_exprs: HashSet<ExpressionId>,
    /// For function handling: maps expr id to inner expression for recursion.
    /// Borrowed so multi-column operators do not clone the whole store per
    /// expression.
    expr_store: Option<&'a HashMap<ExpressionId, linkrs_core::Expression>>,
    /// Descriptions of expressions that could not be resolved to any group.
    /// Unresolved expressions conservatively require every unflat group to
    /// be flattened (see `mark_unresolved`) instead of silently contributing
    /// no groups.
    unresolved: Vec<String>,
}

impl<'a> GroupDependencyAnalyzer<'a> {
    pub fn new(schema: &'a FactorizedSchema, collect_dependent_expr: bool) -> Self {
        Self {
            schema,
            collect_dependent_expr,
            dependent_groups: HashSet::new(),
            required_flat_groups: HashSet::new(),
            dependent_exprs: HashSet::new(),
            expr_store: None,
            unresolved: Vec::new(),
        }
    }

    pub fn with_expr_store(
        schema: &'a FactorizedSchema,
        collect_dependent_expr: bool,
        store: &'a HashMap<ExpressionId, linkrs_core::Expression>,
    ) -> Self {
        Self {
            schema,
            collect_dependent_expr,
            dependent_groups: HashSet::new(),
            required_flat_groups: HashSet::new(),
            dependent_exprs: HashSet::new(),
            expr_store: Some(store),
            unresolved: Vec::new(),
        }
    }

    pub fn dependent_groups(&self) -> &HashSet<u32> {
        &self.dependent_groups
    }

    pub fn required_flat_groups(&self) -> &HashSet<u32> {
        &self.required_flat_groups
    }

    pub fn dependent_exprs(&self) -> &HashSet<ExpressionId> {
        &self.dependent_exprs
    }

    /// Whether any expression could not be resolved to a group.
    pub fn has_unresolved(&self) -> bool {
        !self.unresolved.is_empty()
    }

    /// Descriptions of the unresolved expressions, for diagnostics.
    pub fn unresolved(&self) -> &[String] {
        &self.unresolved
    }

    /// Record an unresolvable expression and fall back to flattening every
    /// unflat group, so callers never silently miss a required flatten.
    /// The recorded description carries the schema shape (group count and
    /// scope size) to pinpoint missing alias registrations in the producing
    /// operator's `compute_factorized_schema`.
    fn mark_unresolved(&mut self, what: String) {
        let context = format!(
            "{what} (schema: {} groups, {} expressions in scope)",
            self.schema.num_groups(),
            self.schema.expressions_in_scope().len()
        );
        log::warn!("GroupDependencyAnalyzer: unresolved {context}");
        self.unresolved.push(context);
        for (pos, group) in self.schema.groups().iter().enumerate() {
            if !group.is_flat() {
                self.required_flat_groups.insert(pos as u32);
            }
        }
    }

    pub fn into_analysis(self) -> GroupDependencyAnalysis {
        GroupDependencyAnalysis {
            dependent_groups: self.dependent_groups,
            required_flat_groups: self.required_flat_groups,
            dependent_exprs: self.dependent_exprs,
        }
    }

    /// Analyze a single expression id: if it is in scope, record its group.
    /// Otherwise attempt to walk its inner expression tree if available.
    ///
    /// Resolution goes through the canonical `resolve_group_pos` order
    /// (id, then id-linked name, then bare name) so the three alias paths
    /// stay ordered in one place instead of drifting across call sites.
    pub fn visit(&mut self, expr_id: &ExpressionId) {
        let owned_name = self.schema.expression_name(expr_id).map(str::to_string);
        if let Some(pos) = self
            .schema
            .resolve_group_pos(Some(expr_id), owned_name.as_deref())
        {
            self.dependent_groups.insert(pos);
            if self.collect_dependent_expr {
                self.dependent_exprs.insert(expr_id.clone());
            }
            return;
        }

        // Not in scope: try to expand via stored expression.
        if let Some(expr) = self.expr_store.and_then(|store| store.get(expr_id)) {
            self.visit_expression(expr);
        } else {
            self.mark_unresolved(format!("expression {expr_id:?} not in scope"));
        }
    }

    pub fn visit_expression(&mut self, expr: &linkrs_core::Expression) {
        use linkrs_core::Expression;
        match expr {
            Expression::Variable(name) => {
                // Variables carry no expression id, so resolution uses the
                // canonical alias order with no id (bare-name fallback last).
                if let Some(pos) = self.schema.resolve_group_pos(None, Some(name)) {
                    self.dependent_groups.insert(pos);
                } else {
                    self.mark_unresolved(format!("variable `{name}` not bound in scope"));
                }
            }
            Expression::Property { object, .. } => {
                self.visit_expression(object);
            }
            Expression::Binary { left, right, .. } => {
                self.visit_expression(left);
                self.visit_expression(right);
            }
            Expression::Unary { operand, .. } => {
                self.visit_expression(operand);
            }
            Expression::Function { args, name } => {
                for arg in args {
                    self.visit_expression(arg.as_expr());
                }
                // List lambda functions require their lambda body to be flat.
                if is_list_lambda(name) {
                    // For list lambda, all dependent groups of the lambda
                    // expression must be flat. The body position comes from
                    // the function registry, not a hardcoded index.
                    if let Some(body) = list_lambda_arg_index(name).and_then(|i| args.get(i)) {
                        let mut lambda_analyzer =
                            GroupDependencyAnalyzer::new(self.schema, self.collect_dependent_expr);
                        lambda_analyzer.expr_store = self.expr_store;
                        lambda_analyzer.visit_expression(body.as_expr());
                        self.required_flat_groups
                            .extend(lambda_analyzer.dependent_groups.iter().copied());
                        self.required_flat_groups
                            .extend(lambda_analyzer.required_flat_groups.iter().copied());
                        self.unresolved.extend(lambda_analyzer.unresolved);
                    }
                }
            }
            Expression::Aggregate { args, filter, .. } => {
                for a in args {
                    self.visit_expression(a);
                }
                if let Some(f) = filter {
                    self.visit_expression(f);
                }
            }
            Expression::List(items) => {
                for item in items {
                    self.visit_expression(item);
                }
            }
            Expression::Map(pairs) => {
                for (_, v) in pairs {
                    self.visit_expression(v);
                }
            }
            Expression::Case {
                test_expr,
                conditions,
                default,
            } => {
                if let Some(t) = test_expr {
                    self.visit_expression(t);
                }
                for (when, then) in conditions {
                    self.visit_expression(when);
                    self.visit_expression(then);
                }
                if let Some(d) = default {
                    self.visit_expression(d);
                }
            }
            Expression::Subscript { collection, index } => {
                self.visit_expression(collection);
                self.visit_expression(index);
            }
            Expression::Range {
                collection,
                start,
                end,
            } => {
                self.visit_expression(collection);
                if let Some(s) = start {
                    self.visit_expression(s);
                }
                if let Some(e) = end {
                    self.visit_expression(e);
                }
            }
            Expression::Literal(_)
            | Expression::Parameter(_)
            | Expression::Label(_)
            | Expression::SessionVariable(_) => {}
            Expression::Path(items) | Expression::PathBuild(items) => {
                for item in items {
                    self.visit_expression(item);
                }
            }
            Expression::TypeCast {
                expression: inner, ..
            } => {
                self.visit_expression(inner);
            }
            Expression::Reduce {
                initial,
                source,
                mapping,
                ..
            } => {
                self.visit_expression(initial);
                self.visit_expression(source);
                self.visit_expression(mapping);
            }
            Expression::ListComprehension {
                variable: _,
                source,
                filter,
                map,
            } => {
                self.visit_expression(source);
                if let Some(f) = filter {
                    self.visit_expression(f);
                }
                if let Some(m) = map {
                    self.visit_expression(m);
                }
            }
            _ => {
                // For any other variants, walk children generically.
                for child in expr.children() {
                    self.visit_expression(child);
                }
            }
        }
    }

    /// Convenience: analyze a set of expression ids.
    pub fn analyze_ids(
        schema: &'a FactorizedSchema,
        expr_ids: &[ExpressionId],
        store: &'a HashMap<ExpressionId, linkrs_core::Expression>,
    ) -> GroupDependencyAnalysis {
        let mut analyzer = GroupDependencyAnalyzer::with_expr_store(schema, false, store);
        for id in expr_ids {
            analyzer.visit(id);
        }
        analyzer.into_analysis()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::planning::plan::factorization::FactorizedSchema;
    use linkrs_core::types::expr::ExpressionId;
    use linkrs_core::types::expr::FunctionArg;

    fn expr(id: u64) -> ExpressionId {
        ExpressionId::new(id)
    }

    #[test]
    fn single_expr_dependency() {
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
        schema.insert_to_group_and_scope(expr(2), g1).unwrap();

        let mut analyzer = GroupDependencyAnalyzer::new(&schema, false);
        analyzer.visit(&expr(2));
        assert!(analyzer.dependent_groups().contains(&g1));
        assert!(!analyzer.dependent_groups().contains(&g0));
    }

    #[test]
    fn expression_tree_walk() {
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        let id_a = expr(10);
        let id_b = expr(20);
        schema.insert_to_group_and_scope(id_a.clone(), g0).unwrap();
        schema.insert_to_group_and_scope(id_b.clone(), g1).unwrap();

        // Build store where combined expr 30 = a + b
        let store = HashMap::new();
        let _combined_expr = linkrs_core::Expression::Binary {
            left: Box::new(linkrs_core::Expression::Variable("a".to_string())),
            op: linkrs_core::types::operators::BinaryOperator::Add,
            right: Box::new(linkrs_core::Expression::Variable("b".to_string())),
        };
        // Simulate that id_a and id_b are the leaves but combined not in scope,
        // analyzer should walk its children when store contains mapping.
        // Here we use visit_expression directly.
        let mut analyzer = GroupDependencyAnalyzer::with_expr_store(&schema, false, &store);
        // Direct visit of Binary should not add groups because variable names not ids
        // So we test that visit by id works, and expression walk for property.
        analyzer.visit(&id_a);
        analyzer.visit(&id_b);
        assert_eq!(analyzer.dependent_groups().len(), 2);
    }

    #[test]
    fn variable_id_primary() {
        // Resolution by expression id is the primary path: it works without
        // any name registration.
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        schema.insert_to_group_and_scope(expr(10), g0).unwrap();
        schema.insert_to_group_and_scope(expr(20), g1).unwrap();

        let mut analyzer = GroupDependencyAnalyzer::new(&schema, false);
        analyzer.visit(&expr(20));
        assert!(analyzer.dependent_groups().contains(&g1));
        assert!(!analyzer.has_unresolved());
    }

    #[test]
    fn variable_name_fallback() {
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        schema
            .insert_to_group_and_scope_with_name(expr(10), Some("a".to_string()), g0)
            .unwrap();
        schema
            .insert_to_group_and_scope_with_name(expr(20), Some("b".to_string()), g1)
            .unwrap();

        let the_expr = linkrs_core::Expression::Binary {
            left: Box::new(linkrs_core::Expression::Variable("a".to_string())),
            op: linkrs_core::types::operators::BinaryOperator::Add,
            right: Box::new(linkrs_core::Expression::Variable("b".to_string())),
        };
        let mut store = HashMap::new();
        let fake_id = expr(999);
        store.insert(fake_id.clone(), the_expr);
        let mut analyzer = GroupDependencyAnalyzer::with_expr_store(&schema, false, &store);
        analyzer.visit(&fake_id);
        assert!(analyzer.dependent_groups().contains(&g0));
        assert!(analyzer.dependent_groups().contains(&g1));
        assert_eq!(analyzer.dependent_groups().len(), 2);
    }

    #[test]
    fn unscoped_expression_requires_flatten_all() {
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
        schema.insert_to_group_and_scope(expr(2), g1).unwrap();

        let empty_store = HashMap::new();
        let mut analyzer = GroupDependencyAnalyzer::with_expr_store(&schema, false, &empty_store);
        analyzer.visit(&expr(999));
        assert!(analyzer.has_unresolved());
        assert_eq!(analyzer.unresolved().len(), 1);
        assert!(analyzer.required_flat_groups().contains(&g1));
        assert!(!analyzer.required_flat_groups().contains(&g0));
    }

    #[test]
    fn unbound_variable_requires_flatten_all() {
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        schema.insert_to_group_and_scope(expr(1), g0).unwrap();
        schema.insert_to_group_and_scope(expr(2), g1).unwrap();

        let mut analyzer = GroupDependencyAnalyzer::new(&schema, false);
        analyzer.visit_expression(&linkrs_core::Expression::Variable("ghost".to_string()));
        assert!(analyzer.has_unresolved());
        assert!(analyzer.required_flat_groups().contains(&g1));
    }

    #[test]
    fn list_filter_lambda_body_requires_flat() {
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        schema
            .insert_to_group_and_scope_with_name(expr(10), Some("a".to_string()), g0)
            .unwrap();
        schema
            .insert_to_group_and_scope_with_name(expr(20), Some("items".to_string()), g1)
            .unwrap();

        // list_filter(a, items.age > 30): the lambda body over the unflat
        // group must be flattened.
        let body = linkrs_core::Expression::Binary {
            left: Box::new(linkrs_core::Expression::Property {
                object: Box::new(linkrs_core::Expression::Variable("items".to_string())),
                property: "age".to_string(),
            }),
            op: linkrs_core::types::operators::BinaryOperator::GreaterThan,
            right: Box::new(linkrs_core::Expression::Literal(
                linkrs_core::value::Value::Int(30),
            )),
        };
        let the_expr = linkrs_core::Expression::Function {
            name: "list_filter".to_string(),
            args: vec![
                FunctionArg::Positional(linkrs_core::Expression::Variable("a".to_string())),
                FunctionArg::Positional(body),
            ],
        };
        let mut store = HashMap::new();
        let fake_id = expr(999);
        store.insert(fake_id.clone(), the_expr);
        let mut analyzer = GroupDependencyAnalyzer::with_expr_store(&schema, false, &store);
        analyzer.visit(&fake_id);
        assert!(analyzer.required_flat_groups().contains(&g1));
        assert!(!analyzer.has_unresolved());
    }

    #[test]
    fn path_expression_visits_items() {
        let mut schema = FactorizedSchema::new();
        let g0 = schema.create_flat_group(false);
        let g1 = schema.create_group();
        schema
            .insert_to_group_and_scope_with_name(expr(10), Some("a".to_string()), g0)
            .unwrap();
        schema
            .insert_to_group_and_scope_with_name(expr(20), Some("b".to_string()), g1)
            .unwrap();

        let the_expr =
            linkrs_core::Expression::Path(vec![linkrs_core::Expression::Variable("b".to_string())]);
        let mut analyzer = GroupDependencyAnalyzer::new(&schema, false);
        analyzer.visit_expression(&the_expr);
        assert!(analyzer.dependent_groups().contains(&g1));
        assert!(!analyzer.has_unresolved());
    }

    #[test]
    fn list_any_all_single_lambda_body_requires_flat() {
        for name in ["list_any", "list_all", "list_single"] {
            let mut schema = FactorizedSchema::new();
            let g0 = schema.create_flat_group(false);
            let g1 = schema.create_group();
            schema
                .insert_to_group_and_scope_with_name(expr(10), Some("a".to_string()), g0)
                .unwrap();
            schema
                .insert_to_group_and_scope_with_name(expr(20), Some("items".to_string()), g1)
                .unwrap();

            let body = linkrs_core::Expression::Binary {
                left: Box::new(linkrs_core::Expression::Property {
                    object: Box::new(linkrs_core::Expression::Variable("items".to_string())),
                    property: "age".to_string(),
                }),
                op: linkrs_core::types::operators::BinaryOperator::GreaterThan,
                right: Box::new(linkrs_core::Expression::Literal(
                    linkrs_core::value::Value::Int(30),
                )),
            };
            let the_expr = linkrs_core::Expression::Function {
                name: name.to_string(),
                args: vec![
                    FunctionArg::Positional(linkrs_core::Expression::Variable("a".to_string())),
                    FunctionArg::Positional(body),
                ],
            };
            let mut store = HashMap::new();
            let fake_id = expr(999);
            store.insert(fake_id.clone(), the_expr);
            let mut analyzer = GroupDependencyAnalyzer::with_expr_store(&schema, false, &store);
            analyzer.visit(&fake_id);
            assert!(
                analyzer.required_flat_groups().contains(&g1),
                "{name} lambda body must require flat"
            );
            assert!(!analyzer.has_unresolved());
        }
    }
}