// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.

//! Transformation based on pushing demand information about columns toward sources.

use itertools::Itertools;
use mz_ore::assert_none;
use std::collections::{BTreeMap, BTreeSet};

use mz_expr::{
    AggregateExpr, AggregateFunc, Id, JoinInputMapper, MirRelationExpr, MirScalarExpr,
    RECURSION_LIMIT,
};
use mz_ore::stack::{CheckedRecursion, RecursionGuard};
use mz_repr::{Datum, Row};

use crate::TransformCtx;

/// Drive demand from the root through operators.
///
/// This transform alerts operators to their columns that influence the
/// ultimate output of the expression, and gives them permission to swap
/// other columns for dummy values. As part of this, operators should not
/// actually use any of these dummy values, lest they run-time error.
///
/// This transformation primarily informs the `Join` operator, which can
/// simplify its intermediate state when it knows that certain columns are
/// not observed in its output. Internal arrangements need not maintain
/// columns that are no longer required in the join pipeline, which are
/// those columns not required by the output nor any further equalities.
///
/// Nowadays, this transform is mostly obsoleted by `ProjectionPushdown`.
/// However, I know of one thing that it still does that `ProjectionPushdown`
/// doesn't do (there might be more such things):
/// if you have something like
/// ```code
///     Project (#0, #1)
///       Join on=(#0 = #1)
/// ```
/// then this is turned into
/// ```code
///     Project (#0, #0)
///       Join on=(#0 = #1)
/// ```
/// This can be beneficial for projecting out some columns earlier inside a complex join (by the LIR
/// planning), and then recovering them after the join (if needed) by duplicating existing columns.
///
/// After the last run of `Demand`, there should always be a `ProjectionPushdown`, so that dummies
/// are eliminated from plans.
#[derive(Debug)]
pub struct Demand {
    recursion_guard: RecursionGuard,
}

impl Default for Demand {
    fn default() -> Demand {
        Demand {
            recursion_guard: RecursionGuard::with_limit(RECURSION_LIMIT),
        }
    }
}

impl CheckedRecursion for Demand {
    fn recursion_guard(&self) -> &RecursionGuard {
        &self.recursion_guard
    }
}

impl crate::Transform for Demand {
    fn name(&self) -> &'static str {
        "Demand"
    }

    #[mz_ore::instrument(
        target = "optimizer",
        level = "debug",
        fields(path.segment = "demand")
    )]
    fn actually_perform_transform(
        &self,
        relation: &mut MirRelationExpr,
        _: &mut TransformCtx,
    ) -> Result<(), crate::TransformError> {
        let result = self.action(
            relation,
            (0..relation.arity()).collect(),
            &mut BTreeMap::new(),
        );
        mz_repr::explain::trace_plan(&*relation);
        result
    }
}

impl Demand {
    /// Columns to be produced.
    fn action(
        &self,
        relation: &mut MirRelationExpr,
        mut columns: BTreeSet<usize>,
        gets: &mut BTreeMap<Id, BTreeSet<usize>>,
    ) -> Result<(), crate::TransformError> {
        self.checked_recur(|_| {
            // A valid relation type is only needed for Maps, but we can't borrow
            // the relation in the corresponding branch of the match statement, since
            // it is already borrowed mutably.
            let relation_type = if matches!(relation, MirRelationExpr::Map { .. }) {
                Some(relation.typ())
            } else {
                None
            };
            match relation {
                MirRelationExpr::Constant { .. } => {
                    // Nothing clever to do with constants, that I can think of.
                    Ok(())
                }
                MirRelationExpr::Get { id, .. } => {
                    gets.entry(*id)
                        .or_insert_with(BTreeSet::new)
                        .extend(columns);
                    Ok(())
                }
                MirRelationExpr::Let { id, value, body } => {
                    // Let harvests any requirements of get from its body,
                    // and pushes the union of the requirements at its value.
                    let id = Id::Local(*id);
                    let prior = gets.insert(id, BTreeSet::new());
                    assert_none!(prior); // no shadowing
                    self.action(body, columns, gets)?;
                    let needs = gets.remove(&id).expect("existing gets entry");
                    if let Some(prior) = prior {
                        gets.insert(id, prior);
                    }

                    self.action(value, needs, gets)
                }
                MirRelationExpr::LetRec {
                    ids,
                    values,
                    limits: _,
                    body,
                } => {
                    let ids_used_across_iterations = MirRelationExpr::recursive_ids(ids, values)
                        .iter()
                        .map(|id| Id::Local(*id))
                        .collect::<BTreeSet<_>>();
                    let ids = ids.iter().map(|id| Id::Local(*id)).collect_vec();
                    for id in ids.iter() {
                        let prior = gets.insert(id.clone(), BTreeSet::new());
                        assert_none!(prior); // no shadowing
                    }
                    self.action(body, columns, gets)?;
                    for (id, value) in ids.iter().rev().zip_eq(values.iter_mut().rev()) {
                        let needs = if !ids_used_across_iterations.contains(id) {
                            gets.remove(id).expect("existing gets entry")
                        } else {
                            // Remove, but ignore the collected needs
                            gets.remove(id).expect("existing gets entry");
                            // Instead of using `gets`, we'll say we need all columns for a
                            // recursive id
                            (0..value.arity()).collect::<BTreeSet<_>>()
                        };
                        self.action(value, needs, gets)?;
                    }
                    Ok(())
                }
                MirRelationExpr::Project { input, outputs } => self.action(
                    input,
                    columns.into_iter().map(|c| outputs[c]).collect(),
                    gets,
                ),
                MirRelationExpr::Map { input, scalars } => {
                    let relation_type = relation_type.as_ref().unwrap();
                    let arity = input.arity();
                    // contains columns whose supports have yet to be explored
                    let mut new_columns = columns.clone();
                    new_columns.retain(|c| *c >= arity);
                    while !new_columns.is_empty() {
                        // explore supports
                        new_columns = new_columns
                            .iter()
                            .flat_map(|c| scalars[*c - arity].support())
                            .filter(|c| !columns.contains(c))
                            .collect();
                        // add those columns to the seen list
                        columns.extend(new_columns.clone());
                        new_columns.retain(|c| *c >= arity);
                    }

                    // Replace un-read expressions with literals to prevent evaluation.
                    for (index, scalar) in scalars.iter_mut().enumerate() {
                        if !columns.contains(&(arity + index)) {
                            // Leave literals as they are, to benefit explain.
                            if !scalar.is_literal() {
                                let typ = relation_type.column_types[arity + index].clone();
                                *scalar = MirScalarExpr::Literal(
                                    Ok(Row::pack_slice(&[Datum::Dummy])),
                                    typ,
                                );
                            }
                        }
                    }

                    columns.retain(|c| *c < arity);
                    self.action(input, columns, gets)
                }
                MirRelationExpr::FlatMap {
                    input,
                    func: _,
                    exprs,
                } => {
                    // A FlatMap which returns zero rows acts like a filter
                    // so we always need to execute it
                    for expr in exprs {
                        expr.support_into(&mut columns);
                    }
                    columns.retain(|c| *c < input.arity());
                    self.action(input, columns, gets)
                }
                MirRelationExpr::Filter { input, predicates } => {
                    for predicate in predicates {
                        predicate.support_into(&mut columns)
                    }
                    self.action(input, columns, gets)
                }
                MirRelationExpr::Join {
                    inputs,
                    equivalences,
                    implementation: _,
                } => {
                    let input_mapper = JoinInputMapper::new(inputs);

                    // Each produced column that is equivalent to a prior column should be remapped
                    // so that upstream uses depend only on the first column, simplifying the demand
                    // analysis. In principle we could choose any representative, if it turns out
                    // that some other column would have been more helpful, but we don't have a great
                    // reason to do that at the moment.
                    let mut permutation: Vec<usize> = (0..input_mapper.total_columns()).collect();
                    for equivalence in equivalences.iter() {
                        let mut first_column = None;
                        for expr in equivalence.iter() {
                            if let MirScalarExpr::Column(c) = expr {
                                if let Some(prior) = &first_column {
                                    permutation[*c] = *prior;
                                } else {
                                    first_column = Some(*c);
                                }
                            }
                        }
                    }

                    let should_permute = columns.iter().any(|c| permutation[*c] != *c);

                    // Each equivalence class imposes internal demand for columns.
                    for equivalence in equivalences.iter() {
                        for expr in equivalence.iter() {
                            expr.support_into(&mut columns);
                        }
                    }

                    // Populate child demands from external and internal demands.
                    let new_columns = input_mapper.split_column_set_by_input(columns.iter());

                    // Recursively indicate the requirements.
                    for (input, columns) in inputs.iter_mut().zip(new_columns) {
                        self.action(input, columns, gets)?;
                    }

                    // Install a permutation if any demanded column is not the
                    // canonical column.
                    if should_permute {
                        *relation = relation.take_dangerous().project(permutation);
                    }

                    Ok(())
                }
                MirRelationExpr::Reduce {
                    input,
                    group_key,
                    aggregates,
                    monotonic: _,
                    expected_group_size: _,
                } => {
                    let mut new_columns = BTreeSet::new();
                    // Group keys determine aggregation granularity and are
                    // each crucial in determining aggregates and even the
                    // multiplicities of other keys.
                    for k in group_key.iter() {
                        k.support_into(&mut new_columns)
                    }
                    for column in columns.iter() {
                        // No obvious requirements on aggregate columns.
                        // A "non-empty" requirement, I guess?
                        if *column >= group_key.len() {
                            aggregates[*column - group_key.len()]
                                .expr
                                .support_into(&mut new_columns);
                        }
                    }

                    // Replace un-demanded aggregations with dummies.
                    let input_type = input.typ();
                    for index in (0..aggregates.len()).rev() {
                        if !columns.contains(&(group_key.len() + index)) {
                            let typ = aggregates[index].typ(&input_type.column_types);
                            aggregates[index] = AggregateExpr {
                                func: AggregateFunc::Dummy,
                                expr: MirScalarExpr::literal_ok(Datum::Dummy, typ.scalar_type),
                                distinct: false,
                            };
                        }
                    }

                    self.action(input, new_columns, gets)
                }
                MirRelationExpr::TopK {
                    input,
                    group_key,
                    order_key,
                    limit,
                    ..
                } => {
                    // Group and order keys and limit must be retained, as they
                    // define which rows are retained.
                    columns.extend(group_key.iter().cloned());
                    columns.extend(order_key.iter().map(|o| o.column));
                    if let Some(limit) = limit {
                        // Strictly speaking not needed because the
                        // `limit` support should be a subset of the
                        // `group_key` support, but we don't want to
                        // take this for granted here.
                        limit.support_into(&mut columns)
                    }
                    self.action(input, columns, gets)
                }
                MirRelationExpr::Negate { input } => self.action(input, columns, gets),
                MirRelationExpr::Threshold { input } => {
                    // Threshold requires all columns, as collapsing any distinct values
                    // has the potential to change how it thresholds counts. This could
                    // be improved with reasoning about distinctness or non-negativity.
                    let arity = input.arity();
                    self.action(input, (0..arity).collect(), gets)
                }
                MirRelationExpr::Union { base, inputs } => {
                    self.action(base, columns.clone(), gets)?;
                    for input in inputs {
                        self.action(input, columns.clone(), gets)?;
                    }
                    Ok(())
                }
                MirRelationExpr::ArrangeBy { input, keys } => {
                    for key_set in keys {
                        for key in key_set {
                            key.support_into(&mut columns);
                        }
                    }
                    self.action(input, columns, gets)
                }
            }
        })
    }
}