content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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.
   */
  package org.apache.asterix.optimizer.rules;
  
  import java.util.ArrayList;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  
  import org.apache.asterix.om.functions.BuiltinFunctions;
  import org.apache.commons.lang3.mutable.Mutable;
  import org.apache.commons.lang3.mutable.MutableObject;
  import org.apache.hyracks.algebricks.common.exceptions.AlgebricksException;
  import org.apache.hyracks.algebricks.common.utils.ListSet;
  import org.apache.hyracks.algebricks.core.algebra.base.ILogicalExpression;
  import org.apache.hyracks.algebricks.core.algebra.base.ILogicalOperator;
  import org.apache.hyracks.algebricks.core.algebra.base.ILogicalPlan;
  import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext;
  import org.apache.hyracks.algebricks.core.algebra.base.LogicalExpressionTag;
  import org.apache.hyracks.algebricks.core.algebra.base.LogicalOperatorTag;
  import org.apache.hyracks.algebricks.core.algebra.base.LogicalVariable;
  import org.apache.hyracks.algebricks.core.algebra.expressions.AbstractFunctionCallExpression;
  import org.apache.hyracks.algebricks.core.algebra.expressions.StatefulFunctionCallExpression;
  import org.apache.hyracks.algebricks.core.algebra.expressions.VariableReferenceExpression;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractLogicalOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractOperatorWithNestedPlans;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AggregateOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AssignOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.RunningAggregateOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.UnnestOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.visitors.CardinalityInferenceVisitor;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.visitors.VariableUtilities;
  import org.apache.hyracks.algebricks.core.algebra.properties.UnpartitionedPropertyComputer;
  import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;
  
  /**
   * The rule cancels out redundant pairs of operators unnest-listify aggregate
   * <p>
   * Case 1 (direct):
   * <p>
   * Before plan:
   * <pre>
   * unnest $x [[ at $p ]] <- function-call:scan-collection($y)
   * aggregate $y <- function-call:listify($z)
   * </pre>
   * After plan:
   * <pre>
   * [[ runningaggregate $p <- tid]]
   * assign $x <- $z
   * </pre>
   *
   * Case 2 (reverse):
   * <p>
   * Before plan:
   * <pre>
   * aggregate $x <- function-call:listify($y)
   * unnest $y <- function-call:scan-collection($z)
   * </pre>
   * After plan:
   * <pre>
   * assign $x <- $z
   * </pre>
   *
   * Notes regarding Case 2(reverse):
   * <ol>
   * <li> Case 2 rewriting is only considered if unnest's input operator has cardinality "exactly 1".
   *      If unnest's input operator produces 0 tuples then the aggregate operator would produce 1 tuple with $x = [],
   *      while the assign rewriting would produce 0 tuples. Therefore it's not equivalent.
   *      If unnest's input operator produces N tuples (where N > 1) then the aggregate operator would produce 1 tuple
   *      with a concatenated list of all unnested values from all input tuples,
   *      while the assign rewriting would produce N tuples. Therefore it's not equivalent.
   *
   * <li> It's configurable whether Case 2 rewriting is attempted or not if the 'aggregate' operator
   *      is the root of a subplan's nested plan. The reason for allowing disabling Case 2 is because
   *      aggregate-unnest elimination may prevent further subplan inlining rewritings.
   * </ol>
   */
  public class RemoveRedundantListifyRule implements IAlgebraicRewriteRule {
  
      private final boolean allowReverseCaseAtSubplanRoot;
  
      public RemoveRedundantListifyRule(boolean allowReverseCaseAtSubplanRoot) {
          this.allowReverseCaseAtSubplanRoot = allowReverseCaseAtSubplanRoot;
      }
  
      @Override
      public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context) {
          return false;
      }
  
      @Override
      public boolean rewritePre(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
              throws AlgebricksException {
          // apply it only at the top of the plan
          ILogicalOperator op = opRef.getValue();
          if (context.checkIfInDontApplySet(this, op)) {
              return false;
          }
          Set<LogicalVariable> varSet = new HashSet<>();
          return applyRuleDown(opRef, false, varSet, context);
      }
  
      private boolean applyRuleDown(Mutable<ILogicalOperator> opRef, boolean isSubplanRoot, Set<LogicalVariable> varSet,
              IOptimizationContext context) throws AlgebricksException {
          boolean changed = applies(opRef, varSet, context);
          boolean skipReverseCase = isSubplanRoot && !allowReverseCaseAtSubplanRoot;
          if (!skipReverseCase) {
              changed |= appliesForReverseCase(opRef, varSet, context);
          }
          AbstractLogicalOperator op = (AbstractLogicalOperator) opRef.getValue();
          VariableUtilities.getUsedVariables(op, varSet);
          if (op.hasNestedPlans()) {
              boolean isSubplanOp = op.getOperatorTag() == LogicalOperatorTag.SUBPLAN;
              // Variables used by the parent operators should be live at op.
              Set<LogicalVariable> localLiveVars = new ListSet<>();
              VariableUtilities.getLiveVariables(op, localLiveVars);
              varSet.retainAll(localLiveVars);
              AbstractOperatorWithNestedPlans aonp = (AbstractOperatorWithNestedPlans) op;
              for (ILogicalPlan p : aonp.getNestedPlans()) {
                  for (Mutable<ILogicalOperator> r : p.getRoots()) {
                      if (applyRuleDown(r, isSubplanOp, varSet, context)) {
                          changed = true;
                      }
                      context.addToDontApplySet(this, r.getValue());
                  }
              }
          }
          for (Mutable<ILogicalOperator> i : op.getInputs()) {
              if (applyRuleDown(i, false, varSet, context)) {
                  changed = true;
              }
              context.addToDontApplySet(this, i.getValue());
          }
          return changed;
      }
  
      private boolean applies(Mutable<ILogicalOperator> opRef, Set<LogicalVariable> varUsedAbove,
              IOptimizationContext context) throws AlgebricksException {
          AbstractLogicalOperator op1 = (AbstractLogicalOperator) opRef.getValue();
          if (op1.getOperatorTag() != LogicalOperatorTag.UNNEST) {
              return false;
          }
          UnnestOperator unnest1 = (UnnestOperator) op1;
          ILogicalExpression expr = unnest1.getExpressionRef().getValue();
  
          if (expr.getExpressionTag() != LogicalExpressionTag.FUNCTION_CALL) {
              return false;
          }
          if (((AbstractFunctionCallExpression) expr).getFunctionIdentifier() != BuiltinFunctions.SCAN_COLLECTION) {
              return false;
          }
          AbstractFunctionCallExpression functionCall = (AbstractFunctionCallExpression) expr;
          ILogicalExpression functionCallArgExpr = functionCall.getArguments().get(0).getValue();
          if (functionCallArgExpr.getExpressionTag() != LogicalExpressionTag.VARIABLE) {
              return false;
          }
          LogicalVariable unnestedVar = ((VariableReferenceExpression) functionCallArgExpr).getVariableReference();
  
          if (varUsedAbove.contains(unnestedVar)) {
              return false;
          }
          Mutable<ILogicalOperator> aggregateParentRef = opRef;
          AbstractLogicalOperator r = op1;
          boolean metAggregate = false;
          while (r.getInputs().size() == 1) {
              aggregateParentRef = r.getInputs().get(0);
              r = (AbstractLogicalOperator) aggregateParentRef.getValue();
              if (r.getOperatorTag() == LogicalOperatorTag.ASSIGN) {
                  AssignOperator assign = (AssignOperator) r;
                  List<LogicalVariable> variables = assign.getVariables();
                  // The assign operator doesn't produce any variable that is used by the unnest.
                  if (variables.contains(unnestedVar)) {
                      return false;
                  }
              } else {
                  if (r.getOperatorTag() == LogicalOperatorTag.AGGREGATE) {
                      metAggregate = true;
                  }
                  break;
              }
          }
  
          if (!metAggregate) {
              return false;
          }
          AggregateOperator agg = (AggregateOperator) r;
          if (agg.getVariables().size() != 1) {
              return false;
          }
          LogicalVariable aggVar = agg.getVariables().get(0);
          ILogicalExpression aggFun = agg.getExpressions().get(0).getValue();
          if (!aggVar.equals(unnestedVar) || aggFun.getExpressionTag() != LogicalExpressionTag.FUNCTION_CALL) {
              return false;
          }
          AbstractFunctionCallExpression f = (AbstractFunctionCallExpression) aggFun;
          if (!BuiltinFunctions.LISTIFY.equals(f.getFunctionIdentifier())) {
              return false;
          }
          if (f.getArguments().size() != 1) {
              return false;
          }
          ILogicalExpression arg0 = f.getArguments().get(0).getValue();
          if (arg0.getExpressionTag() != LogicalExpressionTag.VARIABLE) {
              return false;
          }
          LogicalVariable paramVar = ((VariableReferenceExpression) arg0).getVariableReference();
  
          List<LogicalVariable> assgnVars = new ArrayList<>(1);
          assgnVars.add(unnest1.getVariable());
          List<Mutable<ILogicalExpression>> assgnExprs = new ArrayList<>(1);
          VariableReferenceExpression paramVarRef = new VariableReferenceExpression(paramVar);
          paramVarRef.setSourceLocation(arg0.getSourceLocation());
          assgnExprs.add(new MutableObject<>(paramVarRef));
          AssignOperator assign = new AssignOperator(assgnVars, assgnExprs);
          assign.setSourceLocation(agg.getSourceLocation());
          assign.getInputs().add(agg.getInputs().get(0));
          context.computeAndSetTypeEnvironmentForOperator(assign);
          LogicalVariable posVar = unnest1.getPositionalVariable();
  
          if (posVar == null) {
              // Removes the aggregate operator.
              aggregateParentRef.setValue(assign);
          } else {
              List<LogicalVariable> raggVars = new ArrayList<>(1);
              raggVars.add(posVar);
              List<Mutable<ILogicalExpression>> rAggExprs = new ArrayList<>(1);
              StatefulFunctionCallExpression tidFun =
                      new StatefulFunctionCallExpression(BuiltinFunctions.getBuiltinFunctionInfo(BuiltinFunctions.TID),
                              UnpartitionedPropertyComputer.INSTANCE);
              tidFun.setSourceLocation(agg.getSourceLocation());
              rAggExprs.add(new MutableObject<>(tidFun));
              RunningAggregateOperator rAgg = new RunningAggregateOperator(raggVars, rAggExprs);
              rAgg.setSourceLocation(agg.getSourceLocation());
              rAgg.getInputs().add(new MutableObject<>(assign));
              aggregateParentRef.setValue(rAgg);
              context.computeAndSetTypeEnvironmentForOperator(rAgg);
          }
          // Removes the unnest operator.
          opRef.setValue(unnest1.getInputs().get(0).getValue());
          return true;
      }
  
      private boolean appliesForReverseCase(Mutable<ILogicalOperator> opRef, Set<LogicalVariable> varUsedAbove,
              IOptimizationContext context) throws AlgebricksException {
          AbstractLogicalOperator op1 = (AbstractLogicalOperator) opRef.getValue();
          if (op1.getOperatorTag() != LogicalOperatorTag.AGGREGATE) {
              return false;
          }
          AggregateOperator agg = (AggregateOperator) op1;
          if (agg.getVariables().size() != 1) {
              return false;
          }
          LogicalVariable aggVar = agg.getVariables().get(0);
          ILogicalExpression aggFun = agg.getExpressions().get(0).getValue();
          AbstractFunctionCallExpression f = (AbstractFunctionCallExpression) aggFun;
          if (!BuiltinFunctions.LISTIFY.equals(f.getFunctionIdentifier())) {
              return false;
          }
          if (f.getArguments().size() != 1) {
              return false;
          }
          ILogicalExpression arg0 = f.getArguments().get(0).getValue();
          if (arg0.getExpressionTag() != LogicalExpressionTag.VARIABLE) {
              return false;
          }
          LogicalVariable aggInputVar = ((VariableReferenceExpression) arg0).getVariableReference();
          if (varUsedAbove.contains(aggInputVar)) {
              return false;
          }
  
          if (agg.getInputs().size() == 0) {
              return false;
          }
  
          AbstractLogicalOperator op2 = (AbstractLogicalOperator) agg.getInputs().get(0).getValue();
          if (op2.getOperatorTag() != LogicalOperatorTag.UNNEST) {
              return false;
          }
          UnnestOperator unnest = (UnnestOperator) op2;
          if (unnest.getPositionalVariable() != null) {
              return false;
          }
          if (!unnest.getVariable().equals(aggInputVar)) {
              return false;
          }
          ILogicalExpression unnestArg = unnest.getExpressionRef().getValue();
          if (unnestArg.getExpressionTag() != LogicalExpressionTag.FUNCTION_CALL) {
              return false;
          }
          AbstractFunctionCallExpression scanFunc = (AbstractFunctionCallExpression) unnestArg;
          if (scanFunc.getFunctionIdentifier() != BuiltinFunctions.SCAN_COLLECTION) {
              return false;
          }
          if (scanFunc.getArguments().size() != 1) {
              return false;
          }
          ILogicalOperator unnestInputOp = unnest.getInputs().get(0).getValue();
          if (!CardinalityInferenceVisitor.isCardinalityExactOne(unnestInputOp)) {
              return false;
          }
  
          List<LogicalVariable> assgnVars = new ArrayList<>(1);
          assgnVars.add(aggVar);
          AssignOperator assign = new AssignOperator(assgnVars, scanFunc.getArguments());
          assign.setSourceLocation(agg.getSourceLocation());
          assign.getInputs().add(new MutableObject<>(unnestInputOp));
          context.computeAndSetTypeEnvironmentForOperator(assign);
          opRef.setValue(assign);
          return true;
      }
  }
  