/*
 * 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.beam.sdk.extensions.sql.impl.rel;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptCluster;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptPlanner;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelTraitSet;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.core.Calc;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.logical.LogicalCalc;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataType;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexCall;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexDynamicParam;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexFieldAccess;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexInputRef;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexLiteral;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexLocalRef;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexProgram;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexShuttle;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexUtil;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexVisitorImpl;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.tools.RelBuilder;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.Litmus;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.Util;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.graph.DefaultDirectedGraph;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.graph.DefaultEdge;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.graph.DirectedGraph;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.graph.TopologicalOrderIterator;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.primitives.Ints;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;

/**
 * CalcRelSplitter operates on a {@link Calc} with multiple {@link RexCall} sub-expressions that
 * cannot all be implemented by a single concrete {@link RelNode}.
 *
 * <p>This is a copy of {@link
 * org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rules.CalcRelSplitter} modified to
 * work with Beam. TODO(CALCITE-4538) consider contributing these changes back upstream.
 *
 * <p>For example, the Java and Fennel calculator do not implement an identical set of operators.
 * The Calc can be used to split a single Calc with mixed Java- and Fennel-only operators into a
 * tree of Calc object that can each be individually implemented by either Java or Fennel.and splits
 * it into several Calc instances.
 *
 * <p>Currently the splitter is only capable of handling two "rel types". That is, it can deal with
 * Java vs. Fennel Calcs, but not Java vs. Fennel vs. some other type of Calc.
 *
 * <p>See {@link ProjectToWindowRule} for an example of how this class is used.
 */
@SuppressWarnings({"all", "OperatorPrecedence"})
public class CalcRelSplitter {
  // ~ Static fields/initializers ---------------------------------------------

  private static final Logger RULE_LOGGER = RelOptPlanner.LOGGER;

  // ~ Instance fields --------------------------------------------------------

  protected final RexProgram program;
  private final RelDataTypeFactory typeFactory;

  private final RelType[] relTypes;
  private final RelOptCluster cluster;
  private final RelTraitSet traits;
  private final RelNode child;
  protected final RelBuilder relBuilder;

  // ~ Constructors -----------------------------------------------------------

  /**
   * Constructs a CalcRelSplitter.
   *
   * @param calc Calc to split
   * @param relTypes Array of rel types, e.g. {Java, Fennel}. Must be distinct.
   */
  public CalcRelSplitter(Calc calc, RelBuilder relBuilder, RelType[] relTypes) {
    this.relBuilder = relBuilder;
    for (int i = 0; i < relTypes.length; i++) {
      assert relTypes[i] != null;
      for (int j = 0; j < i; j++) {
        assert relTypes[i] != relTypes[j] : "Rel types must be distinct";
      }
    }
    this.program = calc.getProgram();
    this.cluster = calc.getCluster();
    this.traits = calc.getTraitSet();
    this.typeFactory = calc.getCluster().getTypeFactory();
    this.child = calc.getInput();
    this.relTypes = relTypes;
  }

  // ~ Methods ----------------------------------------------------------------

  public RelNode execute() {
    // Check that program is valid. In particular, this means that every
    // expression is trivial (either an atom, or a function applied to
    // references to atoms) and every expression depends only on
    // expressions to the left.
    assert program.isValid(Litmus.THROW, null);
    final List<RexNode> exprList = program.getExprList();
    final RexNode[] exprs = exprList.toArray(new RexNode[0]);
    assert !RexUtil.containComplexExprs(exprList);

    // Figure out what level each expression belongs to.
    int[] exprLevels = new int[exprs.length];

    // The type of a level is given by
    // relTypes[levelTypeOrdinals[level]].
    int[] levelTypeOrdinals = new int[exprs.length];

    int levelCount = chooseLevels(exprs, -1, exprLevels, levelTypeOrdinals);

    // For each expression, figure out which is the highest level where it
    // is used.
    int[] exprMaxUsingLevelOrdinals =
        new HighestUsageFinder(exprs, exprLevels).getMaxUsingLevelOrdinals();

    // If expressions are used as outputs, mark them as higher than that.
    final List<RexLocalRef> projectRefList = program.getProjectList();
    final RexLocalRef conditionRef = program.getCondition();
    for (RexLocalRef projectRef : projectRefList) {
      exprMaxUsingLevelOrdinals[projectRef.getIndex()] = levelCount;
    }
    if (conditionRef != null) {
      exprMaxUsingLevelOrdinals[conditionRef.getIndex()] = levelCount;
    }

    // Print out what we've got.
    if (RULE_LOGGER.isTraceEnabled()) {
      traceLevelExpressions(exprs, exprLevels, levelTypeOrdinals, levelCount);
    }

    // Now build the calcs.
    RelNode rel = child;
    final int inputFieldCount = program.getInputRowType().getFieldCount();
    int[] inputExprOrdinals = identityArray(inputFieldCount);
    boolean doneCondition = false;
    for (int level = 0; level < levelCount; level++) {
      final int[] projectExprOrdinals;
      final RelDataType outputRowType;
      if (level == (levelCount - 1)) {
        outputRowType = program.getOutputRowType();
        projectExprOrdinals = new int[projectRefList.size()];
        for (int i = 0; i < projectExprOrdinals.length; i++) {
          projectExprOrdinals[i] = projectRefList.get(i).getIndex();
        }
      } else {
        outputRowType = null;

        // Project the expressions which are computed at this level or
        // before, and will be used at later levels.
        List<Integer> projectExprOrdinalList = new ArrayList<>();
        for (int i = 0; i < exprs.length; i++) {
          RexNode expr = exprs[i];
          if (expr instanceof RexLiteral) {
            // Don't project literals. They are always created in
            // the level where they are used.
            exprLevels[i] = -1;
            continue;
          }
          if ((exprLevels[i] <= level) && (exprMaxUsingLevelOrdinals[i] > level)) {
            projectExprOrdinalList.add(i);
          }
        }
        projectExprOrdinals = Ints.toArray(projectExprOrdinalList);
      }

      final RelType relType = relTypes[levelTypeOrdinals[level]];

      // Can we do the condition this level?
      int conditionExprOrdinal = -1;
      if ((conditionRef != null) && !doneCondition) {
        conditionExprOrdinal = conditionRef.getIndex();
        if ((exprLevels[conditionExprOrdinal] > level) || !relType.supportsCondition()) {
          // stand down -- we're not ready to do the condition yet
          conditionExprOrdinal = -1;
        } else {
          doneCondition = true;
        }
      }

      RexProgram program1 =
          createProgramForLevel(
              level,
              levelCount,
              rel.getRowType(),
              exprs,
              exprLevels,
              inputExprOrdinals,
              projectExprOrdinals,
              conditionExprOrdinal,
              outputRowType);
      rel = relType.makeRel(cluster, traits, relBuilder, rel, program1);

      rel = handle(rel);

      // The outputs of this level will be the inputs to the next level.
      inputExprOrdinals = projectExprOrdinals;
    }

    Preconditions.checkArgument(doneCondition || (conditionRef == null), "unhandled condition");
    return rel;
  }

  /**
   * Opportunity to further refine the relational expression created for a given level. The default
   * implementation returns the relational expression unchanged.
   */
  protected RelNode handle(RelNode rel) {
    return rel;
  }

  /**
   * Figures out which expressions to calculate at which level.
   *
   * @param exprs Array of expressions
   * @param conditionOrdinal Ordinal of the condition expression, or -1 if no condition
   * @param exprLevels Level ordinal for each expression (output)
   * @param levelTypeOrdinals The type of each level (output)
   * @return Number of levels required
   */
  private int chooseLevels(
      final RexNode[] exprs, int conditionOrdinal, int[] exprLevels, int[] levelTypeOrdinals) {
    final int inputFieldCount = program.getInputRowType().getFieldCount();

    int levelCount = 0;
    final MaxInputFinder maxInputFinder = new MaxInputFinder(exprLevels);
    boolean[] relTypesPossibleForTopLevel = new boolean[relTypes.length];
    Arrays.fill(relTypesPossibleForTopLevel, true);

    // Compute the order in which to visit expressions.
    final List<Set<Integer>> cohorts = getCohorts();
    final List<Integer> permutation = computeTopologicalOrdering(exprs, cohorts);

    for (int i : permutation) {
      RexNode expr = exprs[i];
      final boolean condition = i == conditionOrdinal;

      if (i < inputFieldCount) {
        assert expr instanceof RexInputRef;
        exprLevels[i] = -1;
        continue;
      }

      // Deduce the minimum level of the expression. An expression must
      // be at a level greater than or equal to all of its inputs.
      int level = maxInputFinder.maxInputFor(expr);

      // If the expression is in a cohort, it can occur no lower than the
      // levels of other expressions in the same cohort.
      Set<Integer> cohort = findCohort(cohorts, i);
      if (cohort != null) {
        for (Integer exprOrdinal : cohort) {
          if (exprOrdinal == i) {
            // Already did this member of the cohort. It's a waste
            // of effort to repeat.
            continue;
          }
          final RexNode cohortExpr = exprs[exprOrdinal];
          int cohortLevel = maxInputFinder.maxInputFor(cohortExpr);
          if (cohortLevel > level) {
            level = cohortLevel;
          }
        }
      }

      // Try to implement this expression at this level.
      // If that is not possible, try to implement it at higher levels.
      levelLoop:
      for (; ; ++level) {
        if (level >= levelCount) {
          // This is a new level. We can use any type we like.
          for (int relTypeOrdinal = 0; relTypeOrdinal < relTypes.length; relTypeOrdinal++) {
            if (!relTypesPossibleForTopLevel[relTypeOrdinal]) {
              continue;
            }
            if (relTypes[relTypeOrdinal].canImplement(expr, condition)) {
              // Success. We have found a type where we can
              // implement this expression.
              exprLevels[i] = level;
              levelTypeOrdinals[level] = relTypeOrdinal;
              assert (level == 0) || (levelTypeOrdinals[level - 1] != levelTypeOrdinals[level])
                  : "successive levels of same type";

              // Figure out which of the other reltypes are
              // still possible for this level.
              // Previous reltypes are not possible.
              for (int j = 0; j < relTypeOrdinal; ++j) {
                relTypesPossibleForTopLevel[j] = false;
              }

              // Successive reltypes may be possible.
              for (int j = relTypeOrdinal + 1; j < relTypes.length; ++j) {
                if (relTypesPossibleForTopLevel[j]) {
                  relTypesPossibleForTopLevel[j] = relTypes[j].canImplement(expr, condition);
                }
              }

              // Move to next level.
              levelTypeOrdinals[levelCount] = firstSet(relTypesPossibleForTopLevel);
              ++levelCount;
              Arrays.fill(relTypesPossibleForTopLevel, true);
              break levelLoop;
            }
          }

          // None of the reltypes still active for this level could
          // implement expr. But maybe we could succeed with a new
          // level, with all options open?
          if (count(relTypesPossibleForTopLevel) >= relTypes.length) {
            // Cannot implement for any type.
            throw new AssertionError("cannot implement " + expr);
          }
          levelTypeOrdinals[levelCount] = firstSet(relTypesPossibleForTopLevel);
          ++levelCount;
          Arrays.fill(relTypesPossibleForTopLevel, true);
        } else {
          final int levelTypeOrdinal = levelTypeOrdinals[level];
          if (!relTypes[levelTypeOrdinal].canImplement(expr, condition)) {
            // Cannot implement this expression in this type;
            // continue to next level.
            continue;
          }
          exprLevels[i] = level;
          break;
        }
      }
    }
    if (levelCount == 0) {
      // At least one level is always required.
      levelCount = 1;
    }
    return levelCount;
  }

  /**
   * Computes the order in which to visit expressions, so that we decide the level of an expression
   * only after the levels of lower expressions have been decided.
   *
   * <p>First, we need to ensure that an expression is visited after all of its inputs.
   *
   * <p>Further, if the expression is a member of a cohort, we need to visit it after the inputs of
   * all other expressions in that cohort. With this condition, expressions in the same cohort will
   * very likely end up in the same level.
   *
   * <p>Note that if there are no cohorts, the expressions from the {@link RexProgram} are already
   * in a suitable order. We perform the topological sort just to ensure that the code path is
   * well-trodden.
   *
   * @param exprs Expressions
   * @param cohorts List of cohorts, each of which is a set of expr ordinals
   * @return Expression ordinals in topological order
   */
  private static List<Integer> computeTopologicalOrdering(
      RexNode[] exprs, List<Set<Integer>> cohorts) {
    final DirectedGraph<Integer, DefaultEdge> graph = DefaultDirectedGraph.create();
    for (int i = 0; i < exprs.length; i++) {
      graph.addVertex(i);
    }
    for (int i = 0; i < exprs.length; i++) {
      final RexNode expr = exprs[i];
      final Set<Integer> cohort = findCohort(cohorts, i);
      final Set<Integer> targets;
      if (cohort == null) {
        targets = Collections.singleton(i);
      } else {
        targets = cohort;
      }
      expr.accept(
          new RexVisitorImpl<Void>(true) {
            @Override
            public Void visitLocalRef(RexLocalRef localRef) {
              for (Integer target : targets) {
                graph.addEdge(localRef.getIndex(), target);
              }
              return null;
            }
          });
    }
    TopologicalOrderIterator<Integer, DefaultEdge> iter = new TopologicalOrderIterator<>(graph);
    final List<Integer> permutation = new ArrayList<>();
    while (iter.hasNext()) {
      permutation.add(iter.next());
    }
    return permutation;
  }

  /**
   * Finds the cohort that contains the given integer, or returns null.
   *
   * @param cohorts List of cohorts, each a set of integers
   * @param ordinal Integer to search for
   * @return Cohort that contains the integer, or null if not found
   */
  private static @Nullable Set<Integer> findCohort(List<Set<Integer>> cohorts, int ordinal) {
    for (Set<Integer> cohort : cohorts) {
      if (cohort.contains(ordinal)) {
        return cohort;
      }
    }
    return null;
  }

  private static int[] identityArray(int length) {
    final int[] ints = new int[length];
    for (int i = 0; i < ints.length; i++) {
      ints[i] = i;
    }
    return ints;
  }

  /**
   * Creates a program containing the expressions for a given level.
   *
   * <p>The expression list of the program will consist of all entries in the expression list <code>
   * allExprs[i]</code> for which the corresponding level ordinal <code>exprLevels[i]</code> is
   * equal to <code>level</code>. Expressions are mapped according to <code>inputExprOrdinals</code>
   * .
   *
   * @param level Level ordinal
   * @param levelCount Number of levels
   * @param inputRowType Input row type
   * @param allExprs Array of all expressions
   * @param exprLevels Array of the level ordinal of each expression
   * @param inputExprOrdinals Ordinals in the expression list of input expressions. Input expression
   *     <code>i</code> will be found at position <code>inputExprOrdinals[i]</code>.
   * @param projectExprOrdinals Ordinals of the expressions to be output this level.
   * @param conditionExprOrdinal Ordinal of the expression to form the condition for this level, or
   *     -1 if there is no condition.
   * @param outputRowType Output row type
   * @return Relational expression
   */
  private RexProgram createProgramForLevel(
      int level,
      int levelCount,
      RelDataType inputRowType,
      RexNode[] allExprs,
      int[] exprLevels,
      int[] inputExprOrdinals,
      final int[] projectExprOrdinals,
      int conditionExprOrdinal,
      @Nullable RelDataType outputRowType) {
    // Build a list of expressions to form the calc.
    List<RexNode> exprs = new ArrayList<>();

    // exprInverseOrdinals describes where an expression in allExprs comes
    // from -- from an input, from a calculated expression, or -1 if not
    // available at this level.
    int[] exprInverseOrdinals = new int[allExprs.length];
    Arrays.fill(exprInverseOrdinals, -1);
    int j = 0;

    // First populate the inputs. They were computed at some previous level
    // and are used here.
    for (int i = 0; i < inputExprOrdinals.length; i++) {
      final int inputExprOrdinal = inputExprOrdinals[i];
      exprs.add(new RexInputRef(i, allExprs[inputExprOrdinal].getType()));
      exprInverseOrdinals[inputExprOrdinal] = j;
      ++j;
    }

    // Next populate the computed expressions.
    final RexShuttle shuttle =
        new InputToCommonExprConverter(
            exprInverseOrdinals, exprLevels, level, inputExprOrdinals, allExprs);
    for (int i = 0; i < allExprs.length; i++) {
      if (exprLevels[i] == level
          || exprLevels[i] == -1
              && level == (levelCount - 1)
              && allExprs[i] instanceof RexLiteral) {
        RexNode expr = allExprs[i];
        final RexNode translatedExpr = expr.accept(shuttle);
        exprs.add(translatedExpr);
        assert exprInverseOrdinals[i] == -1;
        exprInverseOrdinals[i] = j;
        ++j;
      }
    }

    // Form the projection and condition list. Project and condition
    // ordinals are offsets into allExprs, so we need to map them into
    // exprs.
    final List<RexLocalRef> projectRefs = new ArrayList<>(projectExprOrdinals.length);
    final List<String> fieldNames = new ArrayList<>(projectExprOrdinals.length);
    for (int i = 0; i < projectExprOrdinals.length; i++) {
      final int projectExprOrdinal = projectExprOrdinals[i];
      final int index = exprInverseOrdinals[projectExprOrdinal];
      assert index >= 0;
      RexNode expr = allExprs[projectExprOrdinal];
      projectRefs.add(new RexLocalRef(index, expr.getType()));

      // Inherit meaningful field name if possible.
      fieldNames.add(deriveFieldName(expr, i));
    }
    RexLocalRef conditionRef;
    if (conditionExprOrdinal >= 0) {
      final int index = exprInverseOrdinals[conditionExprOrdinal];
      conditionRef = new RexLocalRef(index, allExprs[conditionExprOrdinal].getType());
    } else {
      conditionRef = null;
    }
    if (outputRowType == null) {
      outputRowType = RexUtil.createStructType(typeFactory, projectRefs, fieldNames, null);
    }
    final RexProgram program =
        new RexProgram(inputRowType, exprs, projectRefs, conditionRef, outputRowType);
    // Program is NOT normalized here (e.g. can contain literals in
    // call operands), since literals should be inlined.
    return program;
  }

  private String deriveFieldName(RexNode expr, int ordinal) {
    if (expr instanceof RexInputRef) {
      int inputIndex = ((RexInputRef) expr).getIndex();
      String fieldName = child.getRowType().getFieldList().get(inputIndex).getName();
      // Don't inherit field names like '$3' from child: that's
      // confusing.
      if (!fieldName.startsWith("$") || fieldName.startsWith("$EXPR")) {
        return fieldName;
      }
    }
    return "$" + ordinal;
  }

  /**
   * Traces the given array of level expression lists at the finer level.
   *
   * @param exprs Array expressions
   * @param exprLevels For each expression, the ordinal of its level
   * @param levelTypeOrdinals For each level, the ordinal of its type in the {@link #relTypes} array
   * @param levelCount The number of levels
   */
  private void traceLevelExpressions(
      RexNode[] exprs, int[] exprLevels, int[] levelTypeOrdinals, int levelCount) {
    StringWriter traceMsg = new StringWriter();
    PrintWriter traceWriter = new PrintWriter(traceMsg);
    traceWriter.println("FarragoAutoCalcRule result expressions for: ");
    traceWriter.println(program.toString());

    for (int level = 0; level < levelCount; level++) {
      traceWriter.println("Rel Level " + level + ", type " + relTypes[levelTypeOrdinals[level]]);

      for (int i = 0; i < exprs.length; i++) {
        RexNode expr = exprs[i];
        assert (exprLevels[i] >= -1) && (exprLevels[i] < levelCount)
            : "expression's level is out of range";
        if (exprLevels[i] == level) {
          traceWriter.println("\t" + i + ": " + expr);
        }
      }
      traceWriter.println();
    }
    String msg = traceMsg.toString();
    RULE_LOGGER.trace(msg);
  }

  /** Returns the number of bits set in an array. */
  private static int count(boolean[] booleans) {
    int count = 0;
    for (boolean b : booleans) {
      if (b) {
        ++count;
      }
    }
    return count;
  }

  /** Returns the index of the first set bit in an array. */
  private static int firstSet(boolean[] booleans) {
    for (int i = 0; i < booleans.length; i++) {
      if (booleans[i]) {
        return i;
      }
    }
    return -1;
  }

  /**
   * Searches for a value in a map, and returns the position where it was found, or -1.
   *
   * @param value Value to search for
   * @param map Map to search in
   * @return Ordinal of value in map, or -1 if not found
   */
  private static int indexOf(int value, int[] map) {
    for (int i = 0; i < map.length; i++) {
      if (value == map[i]) {
        return i;
      }
    }
    return -1;
  }

  /**
   * Returns whether a relational expression can be implemented solely in a given {@link RelType}.
   *
   * @param rel Calculation relational expression
   * @param relTypeName Name of a {@link RelType}
   * @return Whether relational expression can be implemented
   */
  protected boolean canImplement(LogicalCalc rel, String relTypeName) {
    for (RelType relType : relTypes) {
      if (relType.name.equals(relTypeName)) {
        return relType.canImplement(rel.getProgram());
      }
    }
    throw new AssertionError("unknown type " + relTypeName);
  }

  /**
   * Returns a list of sets of expressions that should be on the same level.
   *
   * <p>For example, if this method returns { {3, 5}, {4, 7} }, it means that expressions 3 and 5,
   * should be on the same level, and expressions 4 and 7 should be on the same level. The two
   * cohorts do not need to be on the same level.
   *
   * <p>The list is best effort. If it is not possible to arrange that the expressions in a cohort
   * are on the same level, the {@link #execute()} method will still succeed.
   *
   * <p>The default implementation of this method returns the empty list; expressions will be put on
   * the most suitable level. This is generally the lowest possible level, except for literals,
   * which are placed at the level where they are used.
   *
   * @return List of cohorts, that is sets of expressions, that the splitting algorithm should
   *     attempt to place on the same level
   */
  protected List<Set<Integer>> getCohorts() {
    return Collections.emptyList();
  }

  // ~ Inner Classes ----------------------------------------------------------

  /** Type of relational expression. Determines which kinds of expressions it can handle. */
  public abstract static class RelType {
    private final String name;

    protected RelType(String name) {
      this.name = name;
    }

    @Override
    public String toString() {
      return name;
    }

    protected abstract boolean canImplement(RexFieldAccess field);

    protected abstract boolean canImplement(RexDynamicParam param);

    protected abstract boolean canImplement(RexLiteral literal);

    protected abstract boolean canImplement(RexCall call);

    protected boolean supportsCondition() {
      return true;
    }

    protected RelNode makeRel(
        RelOptCluster cluster,
        RelTraitSet traitSet,
        RelBuilder relBuilder,
        RelNode input,
        RexProgram program) {
      return LogicalCalc.create(input, program);
    }

    /**
     * Returns whether this <code>RelType</code> can implement a given expression.
     *
     * @param expr Expression
     * @param condition Whether expression is a condition
     * @return Whether this <code>RelType</code> can implement a given expression.
     */
    public boolean canImplement(RexNode expr, boolean condition) {
      if (condition && !supportsCondition()) {
        return false;
      }
      try {
        expr.accept(new ImplementTester(this));
        return true;
      } catch (CannotImplement e) {
        Util.swallow(e, null);
        return false;
      }
    }

    /**
     * Returns whether this tester's <code>RelType</code> can implement a given program.
     *
     * @param program Program
     * @return Whether this tester's <code>RelType</code> can implement a given program.
     */
    public boolean canImplement(RexProgram program) {
      if ((program.getCondition() != null) && !canImplement(program.getCondition(), true)) {
        return false;
      }
      for (RexNode expr : program.getExprList()) {
        if (!canImplement(expr, false)) {
          return false;
        }
      }
      return true;
    }
  }

  /**
   * Visitor which returns whether an expression can be implemented in a given type of relational
   * expression.
   */
  private static class ImplementTester extends RexVisitorImpl<Void> {
    private final RelType relType;

    ImplementTester(RelType relType) {
      super(false);
      this.relType = relType;
    }

    @Override
    public Void visitCall(RexCall call) {
      if (!relType.canImplement(call)) {
        throw CannotImplement.INSTANCE;
      }
      return null;
    }

    @Override
    public Void visitDynamicParam(RexDynamicParam dynamicParam) {
      if (!relType.canImplement(dynamicParam)) {
        throw CannotImplement.INSTANCE;
      }
      return null;
    }

    @Override
    public Void visitFieldAccess(RexFieldAccess fieldAccess) {
      if (!relType.canImplement(fieldAccess)) {
        throw CannotImplement.INSTANCE;
      }
      return null;
    }

    @Override
    public Void visitLiteral(RexLiteral literal) {
      if (!relType.canImplement(literal)) {
        throw CannotImplement.INSTANCE;
      }
      return null;
    }
  }

  /** Control exception for {@link ImplementTester}. */
  private static class CannotImplement extends RuntimeException {
    @SuppressWarnings("ThrowableInstanceNeverThrown")
    static final CannotImplement INSTANCE = new CannotImplement();
  }

  /**
   * Shuttle which converts every reference to an input field in an expression to a reference to a
   * common sub-expression.
   */
  private static class InputToCommonExprConverter extends RexShuttle {
    private final int[] exprInverseOrdinals;
    private final int[] exprLevels;
    private final int level;
    private final int[] inputExprOrdinals;
    private final RexNode[] allExprs;

    InputToCommonExprConverter(
        int[] exprInverseOrdinals,
        int[] exprLevels,
        int level,
        int[] inputExprOrdinals,
        RexNode[] allExprs) {
      this.exprInverseOrdinals = exprInverseOrdinals;
      this.exprLevels = exprLevels;
      this.level = level;
      this.inputExprOrdinals = inputExprOrdinals;
      this.allExprs = allExprs;
    }

    @Override
    public RexNode visitInputRef(RexInputRef input) {
      final int index = exprInverseOrdinals[input.getIndex()];
      assert index >= 0;
      return new RexLocalRef(index, input.getType());
    }

    @Override
    public RexNode visitLocalRef(RexLocalRef local) {
      // A reference to a local variable becomes a reference to an input
      // if the local was computed at a previous level.
      final int localIndex = local.getIndex();
      final int exprLevel = exprLevels[localIndex];
      if (exprLevel < level) {
        if (allExprs[localIndex] instanceof RexLiteral) {
          // Expression is to be inlined. Use the original expression.
          return allExprs[localIndex];
        }
        int inputIndex = indexOf(localIndex, inputExprOrdinals);
        assert inputIndex >= 0;
        return new RexLocalRef(inputIndex, local.getType());
      } else {
        // It's a reference to what was a local expression at the
        // previous level, and was then projected.
        final int exprIndex = exprInverseOrdinals[localIndex];
        return new RexLocalRef(exprIndex, local.getType());
      }
    }
  }

  /** Finds the highest level used by any of the inputs of a given expression. */
  private static class MaxInputFinder extends RexVisitorImpl<Void> {
    int level;
    private final int[] exprLevels;

    MaxInputFinder(int[] exprLevels) {
      super(true);
      this.exprLevels = exprLevels;
    }

    @Override
    public Void visitLocalRef(RexLocalRef localRef) {
      int inputLevel = exprLevels[localRef.getIndex()];
      level = Math.max(level, inputLevel);
      return null;
    }

    /** Returns the highest level of any of the inputs of an expression. */
    public int maxInputFor(RexNode expr) {
      level = 0;
      expr.accept(this);
      return level;
    }
  }

  /**
   * Builds an array of the highest level which contains an expression which uses each expression as
   * an input.
   */
  private static class HighestUsageFinder extends RexVisitorImpl<Void> {
    private final int[] maxUsingLevelOrdinals;
    private int currentLevel;

    HighestUsageFinder(RexNode[] exprs, int[] exprLevels) {
      super(true);
      this.maxUsingLevelOrdinals = new int[exprs.length];
      Arrays.fill(maxUsingLevelOrdinals, -1);
      for (int i = 0; i < exprs.length; i++) {
        if (exprs[i] instanceof RexLiteral) {
          // Literals are always used directly. It never makes sense
          // to compute them at a lower level and project them to
          // where they are used.
          maxUsingLevelOrdinals[i] = -1;
          continue;
        }
        currentLevel = exprLevels[i];
        @SuppressWarnings("argument.type.incompatible")
        final Void unused = exprs[i].accept(this);
      }
    }

    public int[] getMaxUsingLevelOrdinals() {
      return maxUsingLevelOrdinals;
    }

    @Override
    public Void visitLocalRef(RexLocalRef ref) {
      final int index = ref.getIndex();
      maxUsingLevelOrdinals[index] = Math.max(maxUsingLevelOrdinals[index], currentLevel);
      return null;
    }
  }
}
