// Copyright 2022 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::prelude::{PlanRef, *};
use crate::expr::{ExprImpl, ExprRewriter, ExprVisitor};
use crate::optimizer::plan_expr_visitor::InputRefCounter;
use crate::optimizer::plan_node::*;
use crate::utils::Substitute;

/// Merge contiguous [`LogicalProject`] nodes.
pub struct ProjectMergeRule {}
impl ProjectMergeRule {
    pub fn merge_project_exprs(
        outer_project_exprs: &[ExprImpl],
        inner_project_exprs: &[ExprImpl],
        keep_common_sub_func: bool,
    ) -> Option<Vec<ExprImpl>> {
        let mut input_ref_counter = InputRefCounter::default();
        if keep_common_sub_func {
            for expr in outer_project_exprs {
                input_ref_counter.visit_expr(expr);
            }
            // bail out if it is a project generated by `CommonSubExprExtractRule`.
            for (index, count) in &input_ref_counter.counter {
                if *count > 1 && matches!(inner_project_exprs[*index], ExprImpl::FunctionCall(_)) {
                    return None;
                }
            }
        }

        let mut subst = Substitute {
            mapping: inner_project_exprs.to_vec(),
        };
        let exprs = outer_project_exprs
            .iter()
            .cloned()
            .map(|expr| subst.rewrite_expr(expr))
            .collect();
        Some(exprs)
    }
}

impl ProjectMergeRule {
    pub fn create() -> BoxedRule {
        Box::new(ProjectMergeRule {})
    }
}

impl Rule<Logical> for ProjectMergeRule {
    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
        let outer_project: &LogicalProject = plan.as_logical_project()?;
        let input = outer_project.input();
        let inner_project: &LogicalProject = input.as_logical_project()?;

        let exprs = Self::merge_project_exprs(outer_project.exprs(), inner_project.exprs(), true)?;
        Some(LogicalProject::new(inner_project.input(), exprs).into())
    }
}
