/*
 * 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.spark.sql.catalyst.expressions.codegen

import java.io.ByteArrayInputStream

import scala.annotation.tailrec
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
import scala.util.control.NonFatal

import com.google.common.util.concurrent.{ExecutionError, UncheckedExecutionException}
import org.codehaus.commons.compiler.{CompileException, InternalCompilerException}
import org.codehaus.janino.ClassBodyEvaluator
import org.codehaus.janino.util.ClassFile
import org.codehaus.janino.util.ClassFile.CodeAttribute

import org.apache.spark.{SparkException, SparkIllegalArgumentException, TaskContext, TaskKilledException}
import org.apache.spark.executor.InputMetrics
import org.apache.spark.internal.{Logging, LogKeys, MDC}
import org.apache.spark.internal.LogKeys._
import org.apache.spark.metrics.source.CodegenMetrics
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.encoders.HashableWeakReference
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.codegen.Block._
import org.apache.spark.sql.catalyst.types._
import org.apache.spark.sql.catalyst.util.{ArrayData, CollationAwareUTF8String, CollationFactory, CollationSupport, MapData, SQLOrderingUtil, UnsafeRowUtils}
import org.apache.spark.sql.catalyst.util.DateTimeConstants.NANOS_PER_MILLIS
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.Platform
import org.apache.spark.unsafe.types._
import org.apache.spark.util.{LongAccumulator, NonFateSharingCache, ParentClassLoader, Utils}

/**
 * Java source for evaluating an [[Expression]] given a [[InternalRow]] of input.
 *
 * @param code The sequence of statements required to evaluate the expression.
 *             It should be empty string, if `isNull` and `value` are already existed, or no code
 *             needed to evaluate them (literals).
 * @param isNull A term that holds a boolean value representing whether the expression evaluated
 *                 to null.
 * @param value A term for a (possibly primitive) value of the result of the evaluation. Not
 *              valid if `isNull` is set to `true`.
 */
case class ExprCode(var code: Block, var isNull: ExprValue, var value: ExprValue)

object ExprCode {
  def apply(isNull: ExprValue, value: ExprValue): ExprCode = {
    ExprCode(code = EmptyBlock, isNull, value)
  }

  def forNullValue(dataType: DataType): ExprCode = {
    ExprCode(code = EmptyBlock, isNull = TrueLiteral, JavaCode.defaultLiteral(dataType))
  }

  def forNonNullValue(value: ExprValue): ExprCode = {
    ExprCode(code = EmptyBlock, isNull = FalseLiteral, value = value)
  }
}

/**
 * State used for subexpression elimination.
 *
 * @param eval The source code for evaluating the subexpression.
 * @param children The sequence of subexpressions as the children expressions. Before
 *                 evaluating this subexpression, we should evaluate all children
 *                 subexpressions first. This is used if we want to selectively evaluate
 *                 particular subexpressions, instead of all at once. In the case, we need
 *                 to make sure we evaluate all children subexpressions too.
 */
case class SubExprEliminationState(eval: ExprCode, children: Seq[SubExprEliminationState])

object SubExprEliminationState {
  def apply(eval: ExprCode): SubExprEliminationState = {
    new SubExprEliminationState(eval, Seq.empty)
  }

  def apply(
      eval: ExprCode,
      children: Seq[SubExprEliminationState]): SubExprEliminationState = {
    new SubExprEliminationState(eval, children.reverse)
  }
}

/**
 * Codes and common subexpressions mapping used for subexpression elimination.
 *
 * @param states Foreach expression that is participating in subexpression elimination,
 *               the state to use.
 * @param exprCodesNeedEvaluate Some expression codes that need to be evaluated before
 *                              calling common subexpressions.
 */
case class SubExprCodes(
    states: Map[ExpressionEquals, SubExprEliminationState],
    exprCodesNeedEvaluate: Seq[ExprCode])

/**
 * The main information about a new added function.
 *
 * @param functionName String representing the name of the function
 * @param innerClassName Optional value which is empty if the function is added to
 *                       the outer class, otherwise it contains the name of the
 *                       inner class in which the function has been added.
 * @param innerClassInstance Optional value which is empty if the function is added to
 *                           the outer class, otherwise it contains the name of the
 *                           instance of the inner class in the outer class.
 */
private[codegen] case class NewFunctionSpec(
    functionName: String,
    innerClassName: Option[String],
    innerClassInstance: Option[String])

/**
 * A context for codegen, tracking a list of objects that could be passed into generated Java
 * function.
 */
class CodegenContext extends Logging {

  import CodeGenerator._

  /**
   * Holding a list of objects that could be used passed into generated class.
   */
  val references: mutable.ArrayBuffer[Any] = new mutable.ArrayBuffer[Any]()

  /**
   * Add an object to `references`.
   *
   * Returns the code to access it.
   *
   * This does not to store the object into field but refer it from the references field at the
   * time of use because number of fields in class is limited so we should reduce it.
   */
  def addReferenceObj(objName: String, obj: Any, className: String = null): String = {
    val idx = references.length
    references += obj
    val clsName = Option(className).getOrElse(CodeGenerator.typeName(obj.getClass))
    s"(($clsName) references[$idx] /* $objName */)"
  }

  /**
   * Holding the variable name of the input row of the current operator, will be used by
   * `BoundReference` to generate code.
   *
   * Note that if `currentVars` is not null, `BoundReference` prefers `currentVars` over `INPUT_ROW`
   * to generate code. If you want to make sure the generated code use `INPUT_ROW`, you need to set
   * `currentVars` to null, or set `currentVars(i)` to null for certain columns, before calling
   * `Expression.genCode`.
   */
  var INPUT_ROW = "i"

  /**
   * Holding a list of generated columns as input of current operator, will be used by
   * BoundReference to generate code.
   */
  var currentVars: Seq[ExprCode] = null

  /**
   * Holding expressions' inlined mutable states like `MonotonicallyIncreasingID.count` as a
   * 2-tuple: java type, variable name.
   * As an example, ("int", "count") will produce code:
   * {{{
   *   private int count;
   * }}}
   * as a member variable
   *
   * They will be kept as member variables in generated classes like `SpecificProjection`.
   *
   * Exposed for tests only.
   */
  private[catalyst] val inlinedMutableStates: mutable.ArrayBuffer[(String, String)] =
    mutable.ArrayBuffer.empty[(String, String)]

  /**
   * The mapping between mutable state types and corresponding compacted arrays.
   * The keys are java type string. The values are [[MutableStateArrays]] which encapsulates
   * the compacted arrays for the mutable states with the same java type.
   *
   * Exposed for tests only.
   */
  private[catalyst] val arrayCompactedMutableStates: mutable.Map[String, MutableStateArrays] =
    mutable.Map.empty[String, MutableStateArrays]

  // An array holds the code that will initialize each state
  // Exposed for tests only.
  private[catalyst] val mutableStateInitCode: mutable.ArrayBuffer[String] =
    mutable.ArrayBuffer.empty[String]

  // Tracks the names of all the mutable states.
  private val mutableStateNames: mutable.HashSet[String] = mutable.HashSet.empty

  /**
   * This class holds a set of names of mutableStateArrays that is used for compacting mutable
   * states for a certain type, and holds the next available slot of the current compacted array.
   */
  class MutableStateArrays {
    val arrayNames = mutable.ListBuffer.empty[String]
    createNewArray()

    private[this] var currentIndex = 0

    private def createNewArray() = {
      val newArrayName = freshName("mutableStateArray")
      mutableStateNames += newArrayName
      arrayNames.append(newArrayName)
    }

    def getCurrentIndex: Int = currentIndex

    /**
     * Returns the reference of next available slot in current compacted array. The size of each
     * compacted array is controlled by the constant `MUTABLESTATEARRAY_SIZE_LIMIT`.
     * Once reaching the threshold, new compacted array is created.
     */
    def getNextSlot(): String = {
      if (currentIndex < MUTABLESTATEARRAY_SIZE_LIMIT) {
        val res = s"${arrayNames.last}[$currentIndex]"
        currentIndex += 1
        res
      } else {
        createNewArray()
        currentIndex = 1
        s"${arrayNames.last}[0]"
      }
    }

  }

  /**
   * A map containing the mutable states which have been defined so far using
   * `addImmutableStateIfNotExists`. Each entry contains the name of the mutable state as key and
   * its Java type and init code as value.
   */
  private val immutableStates: mutable.Map[String, (String, String)] =
    mutable.Map.empty[String, (String, String)]

  /**
   * Add a mutable state as a field to the generated class. c.f. the comments above.
   *
   * @param javaType Java type of the field. Note that short names can be used for some types,
   *                 e.g. InternalRow, UnsafeRow, UnsafeArrayData, etc. Other types will have to
   *                 specify the fully-qualified Java type name. See the code in doCompile() for
   *                 the list of default imports available.
   *                 Also, generic type arguments are accepted but ignored.
   * @param variableName Name of the field.
   * @param initFunc Function includes statement(s) to put into the init() method to initialize
   *                 this field. The argument is the name of the mutable state variable.
   *                 If left blank, the field will be default-initialized.
   * @param forceInline whether the declaration and initialization code may be inlined rather than
   *                    compacted. Please set `true` into forceInline for one of the followings:
   *                    1. use the original name of the status
   *                    2. expect to non-frequently generate the status
   *                       (e.g. not much sort operators in one stage)
   * @param useFreshName If this is false and the mutable state ends up inlining in the outer
   *                     class, the name is not changed
   * @return the name of the mutable state variable, which is the original name or fresh name if
   *         the variable is inlined to the outer class, or an array access if the variable is to
   *         be stored in an array of variables of the same type.
   *         A variable will be inlined into the outer class when one of the following conditions
   *         are satisfied:
   *         1. forceInline is true
   *         2. its type is primitive type and the total number of the inlined mutable variables
   *            is less than `OUTER_CLASS_VARIABLES_THRESHOLD`
   *         3. its type is multi-dimensional array
   *         When a variable is compacted into an array, the max size of the array for compaction
   *         is given by `MUTABLESTATEARRAY_SIZE_LIMIT`.
   */
  def addMutableState(
      javaType: String,
      variableName: String,
      initFunc: String => String = _ => "",
      forceInline: Boolean = false,
      useFreshName: Boolean = true): String = {

    // want to put a primitive type variable at outerClass for performance
    val canInlinePrimitive = isPrimitiveType(javaType) &&
      (inlinedMutableStates.length < OUTER_CLASS_VARIABLES_THRESHOLD)
    if (forceInline || canInlinePrimitive || javaType.contains("[][]")) {
      val varName = if (useFreshName) freshName(variableName) else variableName
      val initCode = initFunc(varName)
      inlinedMutableStates += ((javaType, varName))
      mutableStateInitCode += initCode
      mutableStateNames += varName
      varName
    } else {
      val arrays = arrayCompactedMutableStates.getOrElseUpdate(javaType, new MutableStateArrays)
      val element = arrays.getNextSlot()

      val initCode = initFunc(element)
      mutableStateInitCode += initCode
      element
    }
  }

  /**
   * Add an immutable state as a field to the generated class only if it does not exist yet a field
   * with that name. This helps reducing the number of the generated class' fields, since the same
   * variable can be reused by many functions.
   *
   * Even though the added variables are not declared as final, they should never be reassigned in
   * the generated code to prevent errors and unexpected behaviors.
   *
   * Internally, this method calls `addMutableState`.
   *
   * @param javaType Java type of the field.
   * @param variableName Name of the field.
   * @param initFunc Function includes statement(s) to put into the init() method to initialize
   *                 this field. The argument is the name of the mutable state variable.
   */
  def addImmutableStateIfNotExists(
      javaType: String,
      variableName: String,
      initFunc: String => String = _ => ""): Unit = {
    val existingImmutableState = immutableStates.get(variableName)
    if (existingImmutableState.isEmpty) {
      addMutableState(javaType, variableName, initFunc, useFreshName = false, forceInline = true)
      immutableStates(variableName) = (javaType, initFunc(variableName))
    } else {
      val (prevJavaType, prevInitCode) = existingImmutableState.get
      assert(prevJavaType == javaType, s"$variableName has already been defined with type " +
        s"$prevJavaType and now it is tried to define again with type $javaType.")
      assert(prevInitCode == initFunc(variableName), s"$variableName has already been defined " +
        s"with different initialization statements.")
    }
  }

  /**
   * Add buffer variable which stores data coming from an [[InternalRow]]. This methods guarantees
   * that the variable is safely stored, which is important for (potentially) byte array backed
   * data types like: UTF8String, ArrayData, MapData & InternalRow.
   */
  def addBufferedState(dataType: DataType, variableName: String, initCode: String): ExprCode = {
    val value = addMutableState(javaType(dataType), variableName)
    val code = UserDefinedType.sqlType(dataType) match {
      case _: StringType => code"$value = $initCode.clone();"
      case _: StructType | _: ArrayType | _: MapType => code"$value = $initCode.copy();"
      case _ => code"$value = $initCode;"
    }
    ExprCode(code, FalseLiteral, JavaCode.global(value, dataType))
  }

  def declareMutableStates(): String = {
    // It's possible that we add same mutable state twice, e.g. the `mergeExpressions` in
    // `TypedAggregateExpression`, we should call `distinct` here to remove the duplicated ones.
    val inlinedStates = inlinedMutableStates.distinct.map { case (javaType, variableName) =>
      s"private $javaType $variableName;"
    }

    val arrayStates = arrayCompactedMutableStates.flatMap { case (javaType, mutableStateArrays) =>
      val numArrays = mutableStateArrays.arrayNames.size
      mutableStateArrays.arrayNames.zipWithIndex.map { case (arrayName, index) =>
        val length = if (index + 1 == numArrays) {
          mutableStateArrays.getCurrentIndex
        } else {
          MUTABLESTATEARRAY_SIZE_LIMIT
        }
        if (javaType.contains("[]")) {
          // initializer had an one-dimensional array variable
          val baseType = javaType.substring(0, javaType.length - 2)
          s"private $javaType[] $arrayName = new $baseType[$length][];"
        } else {
          // initializer had a scalar variable
          s"private $javaType[] $arrayName = new $javaType[$length];"
        }
      }
    }

    (inlinedStates ++ arrayStates).mkString("\n")
  }

  def initMutableStates(): String = {
    // It's possible that we add same mutable state twice, e.g. the `mergeExpressions` in
    // `TypedAggregateExpression`, we should call `distinct` here to remove the duplicated ones.
    val initCodes = mutableStateInitCode.distinct.map(_ + "\n")

    // The generated initialization code may exceed 64kb function size limit in JVM if there are too
    // many mutable states, so split it into multiple functions.
    splitExpressions(expressions = initCodes.toSeq, funcName = "init", arguments = Nil)
  }

  /**
   * Code statements to initialize states that depend on the partition index.
   * An integer `partitionIndex` will be made available within the scope.
   */
  val partitionInitializationStatements: mutable.ArrayBuffer[String] = mutable.ArrayBuffer.empty

  def addPartitionInitializationStatement(statement: String): Unit = {
    partitionInitializationStatements += statement
  }

  def initPartition(): String = {
    partitionInitializationStatements.mkString("\n")
  }

  /**
   * Holds expressions that are equivalent. Used to perform subexpression elimination
   * during codegen.
   *
   * For expressions that appear more than once, generate additional code to prevent
   * recomputing the value.
   *
   * For example, consider two expression generated from this SQL statement:
   *  SELECT (col1 + col2), (col1 + col2) / col3.
   *
   *  equivalentExpressions will match the tree containing `col1 + col2` and it will only
   *  be evaluated once.
   */
  private val equivalentExpressions: EquivalentExpressions = new EquivalentExpressions

  // Foreach expression that is participating in subexpression elimination, the state to use.
  // Visible for testing.
  private[expressions] var subExprEliminationExprs =
    Map.empty[ExpressionEquals, SubExprEliminationState]

  // The collection of sub-expression result resetting methods that need to be called on each row.
  private val subexprFunctions = mutable.ArrayBuffer.empty[String]

  val outerClassName = "OuterClass"

  /**
   * Holds the class and instance names to be generated, where `OuterClass` is a placeholder
   * standing for whichever class is generated as the outermost class and which will contain any
   * inner sub-classes. All other classes and instance names in this list will represent private,
   * inner sub-classes.
   */
  private val classes: mutable.ListBuffer[(String, String)] =
    mutable.ListBuffer[(String, String)](outerClassName -> null)

  // A map holding the current size in bytes of each class to be generated.
  private val classSize: mutable.Map[String, Int] =
    mutable.Map[String, Int](outerClassName -> 0)

  // Nested maps holding function names and their code belonging to each class.
  private val classFunctions: mutable.Map[String, mutable.Map[String, String]] =
    mutable.Map(outerClassName -> mutable.Map.empty[String, String])

  // Verbatim extra code to be added to the OuterClass.
  private val extraClasses: mutable.ListBuffer[String] = mutable.ListBuffer[String]()

  // Returns the size of the most recently added class.
  private def currClassSize(): Int = classSize(classes.head._1)

  // Returns the class name and instance name for the most recently added class.
  private def currClass(): (String, String) = classes.head

  // Adds a new class. Requires the class' name, and its instance name.
  private def addClass(className: String, classInstance: String): Unit = {
    classes.prepend(className -> classInstance)
    classSize += className -> 0
    classFunctions += className -> mutable.Map.empty[String, String]
  }

  /**
   * Adds a function to the generated class. If the code for the `OuterClass` grows too large, the
   * function will be inlined into a new private, inner class, and a class-qualified name for the
   * function will be returned. Otherwise, the function will be inlined to the `OuterClass` the
   * simple `funcName` will be returned.
   *
   * @param funcName the class-unqualified name of the function
   * @param funcCode the body of the function
   * @param inlineToOuterClass whether the given code must be inlined to the `OuterClass`. This
   *                           can be necessary when a function is declared outside of the context
   *                           it is eventually referenced and a returned qualified function name
   *                           cannot otherwise be accessed.
   * @return the name of the function, qualified by class if it will be inlined to a private,
   *         inner class
   */
  def addNewFunction(
      funcName: String,
      funcCode: String,
      inlineToOuterClass: Boolean = false): String = {
    val newFunction = addNewFunctionInternal(funcName, funcCode, inlineToOuterClass)
    newFunction match {
      case NewFunctionSpec(functionName, None, None) => functionName
      case NewFunctionSpec(functionName, Some(_), Some(innerClassInstance)) =>
        innerClassInstance + "." + functionName
      case _ =>
        throw QueryExecutionErrors.addNewFunctionMismatchedWithFunctionError(funcName)
    }
  }

  private[this] def addNewFunctionInternal(
      funcName: String,
      funcCode: String,
      inlineToOuterClass: Boolean): NewFunctionSpec = {
    val (className, classInstance) = if (inlineToOuterClass) {
      outerClassName -> ""
    } else if (currClassSize() > GENERATED_CLASS_SIZE_THRESHOLD) {
      val className = freshName("NestedClass")
      val classInstance = freshName("nestedClassInstance")

      addClass(className, classInstance)

      className -> classInstance
    } else {
      currClass()
    }

    addNewFunctionToClass(funcName, funcCode, className)

    if (className == outerClassName) {
      NewFunctionSpec(funcName, None, None)
    } else {
      NewFunctionSpec(funcName, Some(className), Some(classInstance))
    }
  }

  private[this] def addNewFunctionToClass(
      funcName: String,
      funcCode: String,
      className: String) = {
    classSize(className) += funcCode.length
    classFunctions(className) += funcName -> funcCode
  }

  /**
   * Declares all function code. If the added functions are too many, split them into nested
   * sub-classes to avoid hitting Java compiler constant pool limitation.
   */
  def declareAddedFunctions(): String = {
    val inlinedFunctions = classFunctions(outerClassName).values

    // Nested, private sub-classes have no mutable state (though they do reference the outer class'
    // mutable state), so we declare and initialize them inline to the OuterClass.
    val initNestedClasses = classes.filter(_._1 != outerClassName).map {
      case (className, classInstance) =>
        s"private $className $classInstance = new $className();"
    }

    val declareNestedClasses = classFunctions.filter { case (k, _) => k != outerClassName }.map {
      case (className, functions) =>
        s"""
           |private class $className {
           |  ${functions.values.mkString("\n")}
           |}
           """.stripMargin
    }

    (inlinedFunctions ++ initNestedClasses ++ declareNestedClasses).mkString("\n")
  }

  /**
   * Emits extra inner classes added with addExtraCode
   */
  def emitExtraCode(): String = {
    extraClasses.mkString("\n")
  }

  /**
   * Add extra source code to the outermost generated class.
   * @param code verbatim source code of the inner class to be added.
   */
  def addInnerClass(code: String): Unit = {
    extraClasses.append(code)
  }

  /**
   * The map from a variable name to it's next ID.
   */
  private val freshNameIds = new mutable.HashMap[String, Int]
  freshNameIds += INPUT_ROW -> 1

  /**
   * A prefix used to generate fresh name.
   */
  var freshNamePrefix = ""

  /**
   * The map from a place holder to a corresponding comment
   */
  private val placeHolderToComments = new mutable.HashMap[String, String]

  /**
   * Returns a term name that is unique within this instance of a `CodegenContext`.
   */
  def freshName(name: String): String = synchronized {
    val fullName = if (freshNamePrefix == "") {
      name
    } else {
      s"${freshNamePrefix}_$name"
    }
    val id = freshNameIds.getOrElse(fullName, 0)
    freshNameIds(fullName) = id + 1
    s"${fullName}_$id"
  }

  /**
   * Creates an `ExprValue` representing a local java variable of required data type.
   */
  def freshVariable(name: String, dt: DataType): VariableValue =
    JavaCode.variable(freshName(name), dt)

  /**
   * Creates an `ExprValue` representing a local java variable of required Java class.
   */
  def freshVariable(name: String, javaClass: Class[_]): VariableValue =
    JavaCode.variable(freshName(name), javaClass)

  /**
   * Generates code for equal expression in Java.
   */
  def genEqual(dataType: DataType, c1: String, c2: String): String = dataType match {
    case BinaryType => s"java.util.Arrays.equals($c1, $c2)"
    case FloatType =>
      s"((java.lang.Float.isNaN($c1) && java.lang.Float.isNaN($c2)) || $c1 == $c2)"
    case DoubleType =>
      s"((java.lang.Double.isNaN($c1) && java.lang.Double.isNaN($c2)) || $c1 == $c2)"
    case st: StringType if st.supportsBinaryOrdering => s"$c1.binaryEquals($c2)"
    case st: StringType => s"$c1.semanticEquals($c2, ${st.collationId})"
    case dt: DataType if isPrimitiveType(dt) => s"$c1 == $c2"
    case dt: DataType if dt.isInstanceOf[AtomicType] => s"$c1.equals($c2)"
    case array: ArrayType => genComp(array, c1, c2) + " == 0"
    case struct: StructType => genComp(struct, c1, c2) + " == 0"
    case udt: UserDefinedType[_] => genEqual(udt.sqlType, c1, c2)
    case CalendarIntervalType => s"$c1.equals($c2)"
    case NullType => "false"
    case _ =>
      throw QueryExecutionErrors.cannotGenerateCodeForIncomparableTypeError(
        "equality", dataType)
  }

  /**
   * Generates code for comparing two expressions.
   *
   * @param dataType data type of the expressions
   * @param c1 name of the variable of expression 1's output
   * @param c2 name of the variable of expression 2's output
   */
  def genComp(dataType: DataType, c1: String, c2: String): String = dataType match {
    // java boolean doesn't support > or < operator
    case BooleanType => s"($c1 == $c2 ? 0 : ($c1 ? 1 : -1))"
    case DoubleType =>
      val clsName = SQLOrderingUtil.getClass.getName.stripSuffix("$")
      s"$clsName.compareDoubles($c1, $c2)"
    case FloatType =>
      val clsName = SQLOrderingUtil.getClass.getName.stripSuffix("$")
      s"$clsName.compareFloats($c1, $c2)"
    case st: StringType if st.supportsBinaryOrdering => s"$c1.binaryCompare($c2)"
    case st: StringType => s"$c1.semanticCompare($c2, ${st.collationId})"
    case dt: DataType if isPrimitiveType(dt) => s"($c1 > $c2 ? 1 : $c1 < $c2 ? -1 : 0)"
    case BinaryType => s"org.apache.spark.unsafe.types.ByteArray.compareBinary($c1, $c2)"
    case CalendarIntervalType => s"$c1.compareTo($c2)"
    case NullType => "0"
    case array: ArrayType =>
      val elementType = array.elementType
      val compareFunc = freshName("compareArray")
      val minLength = freshName("minLength")
      val funcCode: String =
        s"""
          public int $compareFunc(ArrayData a, ArrayData b) {
            // when comparing unsafe arrays, try equals first as it compares the binary directly
            // which is very fast.
            if (a instanceof UnsafeArrayData && b instanceof UnsafeArrayData && a.equals(b)) {
              return 0;
            }
            int lengthA = a.numElements();
            int lengthB = b.numElements();
            int $minLength = (lengthA > lengthB) ? lengthB : lengthA;
            for (int i = 0; i < $minLength; i++) {
              ${genCompElementsAt("a", "b", "i", elementType)}
            }

            if (lengthA < lengthB) {
              return -1;
            } else if (lengthA > lengthB) {
              return 1;
            }
            return 0;
          }
        """
      s"${addNewFunction(compareFunc, funcCode)}($c1, $c2)"
    case schema: StructType =>
      val comparisons = GenerateOrdering.genComparisons(this, schema)
      val compareFunc = freshName("compareStruct")
      val funcCode: String =
        s"""
          public int $compareFunc(InternalRow a, InternalRow b) {
            // when comparing unsafe rows, try equals first as it compares the binary directly
            // which is very fast.
            if (a instanceof UnsafeRow && b instanceof UnsafeRow && a.equals(b)) {
              return 0;
            }
            $comparisons
            return 0;
          }
        """
      s"${addNewFunction(compareFunc, funcCode)}($c1, $c2)"
    case map: MapType =>
      val compareFunc = freshName("compareMapData")
      val funcCode = genCompMapData(map.keyType, map.valueType, compareFunc)
      s"${addNewFunction(compareFunc, funcCode)}($c1, $c2)"
    case other if other.isInstanceOf[AtomicType] => s"$c1.compare($c2)"
    case udt: UserDefinedType[_] => genComp(udt.sqlType, c1, c2)
    case _ =>
      throw QueryExecutionErrors.cannotGenerateCodeForIncomparableTypeError("compare", dataType)
  }

  private def genCompMapData(
      keyType: DataType,
      valueType: DataType,
      compareFunc: String): String = {
    s"""
       |public int $compareFunc(MapData a, MapData b) {
       |  int lengthA = a.numElements();
       |  int lengthB = b.numElements();
       |  ArrayData keyArrayA = a.keyArray();
       |  ArrayData valueArrayA = a.valueArray();
       |  ArrayData keyArrayB = b.keyArray();
       |  ArrayData valueArrayB = b.valueArray();
       |  int minLength = (lengthA > lengthB) ? lengthB : lengthA;
       |  for (int i = 0; i < minLength; i++) {
       |    ${genCompElementsAt("keyArrayA", "keyArrayB", "i", keyType)}
       |    ${genCompElementsAt("valueArrayA", "valueArrayB", "i", valueType)}
       |  }
       |
       |  if (lengthA < lengthB) {
       |    return -1;
       |  } else if (lengthA > lengthB) {
       |    return 1;
       |  }
       |  return 0;
       |}
     """.stripMargin
  }

  private def genCompElementsAt(arrayA: String, arrayB: String, i: String,
    elementType : DataType): String = {
    val elementA = freshName("elementA")
    val isNullA = freshName("isNullA")
    val elementB = freshName("elementB")
    val isNullB = freshName("isNullB")
    val jt = javaType(elementType);
    s"""
       |boolean $isNullA = $arrayA.isNullAt($i);
       |boolean $isNullB = $arrayB.isNullAt($i);
       |if ($isNullA && $isNullB) {
       |  // Nothing
       |} else if ($isNullA) {
       |  return -1;
       |} else if ($isNullB) {
       |  return 1;
       |} else {
       |  $jt $elementA = ${getValue(arrayA, elementType, i)};
       |  $jt $elementB = ${getValue(arrayB, elementType, i)};
       |  int comp = ${genComp(elementType, elementA, elementB)};
       |  if (comp != 0) {
       |    return comp;
       |  }
       |}
     """.stripMargin
  }

  /**
   * Generates code for greater of two expressions.
   *
   * @param dataType data type of the expressions
   * @param c1 name of the variable of expression 1's output
   * @param c2 name of the variable of expression 2's output
   */
  def genGreater(dataType: DataType, c1: String, c2: String): String = javaType(dataType) match {
    case JAVA_BYTE | JAVA_SHORT | JAVA_INT | JAVA_LONG => s"$c1 > $c2"
    case _ => s"(${genComp(dataType, c1, c2)}) > 0"
  }

  /**
   * Generates code for updating `partialResult` if `item` is smaller than it.
   *
   * @param dataType data type of the expressions
   * @param partialResult `ExprCode` representing the partial result which has to be updated
   * @param item `ExprCode` representing the new expression to evaluate for the result
   */
  def reassignIfSmaller(dataType: DataType, partialResult: ExprCode, item: ExprCode): String = {
    s"""
       |if (!${item.isNull} && (${partialResult.isNull} ||
       |  ${genGreater(dataType, partialResult.value, item.value)})) {
       |  ${partialResult.isNull} = false;
       |  ${partialResult.value} = ${item.value};
       |}
      """.stripMargin
  }

  /**
   * Generates code for updating `partialResult` if `item` is greater than it.
   *
   * @param dataType data type of the expressions
   * @param partialResult `ExprCode` representing the partial result which has to be updated
   * @param item `ExprCode` representing the new expression to evaluate for the result
   */
  def reassignIfGreater(dataType: DataType, partialResult: ExprCode, item: ExprCode): String = {
    s"""
       |if (!${item.isNull} && (${partialResult.isNull} ||
       |  ${genGreater(dataType, item.value, partialResult.value)})) {
       |  ${partialResult.isNull} = false;
       |  ${partialResult.value} = ${item.value};
       |}
      """.stripMargin
  }

  /**
   * Generates code to do null safe execution, i.e. only execute the code when the input is not
   * null by adding null check if necessary.
   *
   * @param nullable used to decide whether we should add null check or not.
   * @param isNull the code to check if the input is null.
   * @param execute the code that should only be executed when the input is not null.
   */
  def nullSafeExec(nullable: Boolean, isNull: String)(execute: String): String = {
    if (nullable) {
      s"""
        if (!$isNull) {
          $execute
        }
      """
    } else {
      "\n" + execute
    }
  }

  /**
   * Generates code to do null safe execution when accessing properties of complex
   * ArrayData elements.
   *
   * @param nullElements used to decide whether the ArrayData might contain null or not.
   * @param isNull a variable indicating whether the result will be evaluated to null or not.
   * @param arrayData a variable name representing the ArrayData.
   * @param execute the code that should be executed only if the ArrayData doesn't contain
   *                any null.
   */
  def nullArrayElementsSaveExec(
      nullElements: Boolean,
      isNull: String,
      arrayData: String)(
      execute: String): String = {
    val i = freshName("idx")
    if (nullElements) {
      s"""
         |for (int $i = 0; !$isNull && $i < $arrayData.numElements(); $i++) {
         |  $isNull |= $arrayData.isNullAt($i);
         |}
         |if (!$isNull) {
         |  $execute
         |}
       """.stripMargin
    } else {
      execute
    }
  }

  /**
   * Splits the generated code of expressions into multiple functions, because function has
   * 64kb code size limit in JVM. If the class to which the function would be inlined would grow
   * beyond 1000kb, we declare a private, inner sub-class, and the function is inlined to it
   * instead, because classes have a constant pool limit of 65,536 named values.
   *
   * Note that different from `splitExpressions`, we will extract the current inputs of this
   * context and pass them to the generated functions. The input is `INPUT_ROW` for normal codegen
   * path, and `currentVars` for whole stage codegen path. Whole stage codegen path is not
   * supported yet.
   *
   * @param expressions the codes to evaluate expressions.
   * @param funcName the split function name base.
   * @param extraArguments the list of (type, name) of the arguments of the split function,
   *                       except for the current inputs like `ctx.INPUT_ROW`.
   * @param returnType the return type of the split function.
   * @param makeSplitFunction makes split function body, e.g. add preparation or cleanup.
   * @param foldFunctions folds the split function calls.
   */
  def splitExpressionsWithCurrentInputs(
      expressions: Seq[String],
      funcName: String = "apply",
      extraArguments: Seq[(String, String)] = Nil,
      returnType: String = "void",
      makeSplitFunction: String => String = identity,
      foldFunctions: Seq[String] => String = _.mkString("", ";\n", ";")): String = {
    // TODO: support whole stage codegen
    if (INPUT_ROW == null || currentVars != null) {
      expressions.mkString("\n")
    } else {
      splitExpressions(
        expressions,
        funcName,
        ("InternalRow", INPUT_ROW) +: extraArguments,
        returnType,
        makeSplitFunction,
        foldFunctions)
    }
  }

  /**
   * Splits the generated code of expressions into multiple functions, because function has
   * 64kb code size limit in JVM. If the class to which the function would be inlined would grow
   * beyond 1000kb, we declare a private, inner sub-class, and the function is inlined to it
   * instead, because classes have a constant pool limit of 65,536 named values.
   *
   * @param expressions the codes to evaluate expressions.
   * @param funcName the split function name base.
   * @param arguments the list of (type, name) of the arguments of the split function.
   * @param returnType the return type of the split function.
   * @param makeSplitFunction makes split function body, e.g. add preparation or cleanup.
   * @param foldFunctions folds the split function calls.
   */
  def splitExpressions(
      expressions: Seq[String],
      funcName: String,
      arguments: Seq[(String, String)],
      returnType: String = "void",
      makeSplitFunction: String => String = identity,
      foldFunctions: Seq[String] => String = _.mkString("", ";\n", ";")): String = {
    val blocks = buildCodeBlocks(expressions)

    if (blocks.length == 1) {
      // inline execution if only one block
      blocks.head
    } else {
      if (Utils.isTesting) {
        // Passing global variables to the split method is dangerous, as any mutating to it is
        // ignored and may lead to unexpected behavior.
        arguments.foreach { case (_, name) =>
          assert(!mutableStateNames.contains(name),
            s"split function argument $name cannot be a global variable.")
        }
      }

      val func = freshName(funcName)
      val argString = arguments.map { case (t, name) => s"$t $name" }.mkString(", ")
      val functions = blocks.zipWithIndex.map { case (body, i) =>
        val name = s"${func}_$i"
        val code = s"""
           |private $returnType $name($argString) {
           |  ${makeSplitFunction(body)}
           |}
         """.stripMargin
        addNewFunctionInternal(name, code, inlineToOuterClass = false)
      }

      val (outerClassFunctions, innerClassFunctions) = functions.partition(_.innerClassName.isEmpty)

      val argsString = arguments.map(_._2).mkString(", ")
      val outerClassFunctionCalls = outerClassFunctions.map(f => s"${f.functionName}($argsString)")

      val innerClassFunctionCalls = generateInnerClassesFunctionCalls(
        innerClassFunctions,
        func,
        arguments,
        returnType,
        makeSplitFunction,
        foldFunctions)

      foldFunctions(outerClassFunctionCalls ++ innerClassFunctionCalls)
    }
  }

  /**
   * Splits the generated code of expressions into multiple sequences of String
   * based on a threshold of length of a String
   *
   * @param expressions the codes to evaluate expressions.
   */
  private def buildCodeBlocks(expressions: Seq[String]): Seq[String] = {
    val blocks = new ArrayBuffer[String]()
    val blockBuilder = new StringBuilder()
    var length = 0
    val splitThreshold = SQLConf.get.methodSplitThreshold
    for (code <- expressions) {
      // We can't know how many bytecode will be generated, so use the length of source code
      // as metric. A method should not go beyond 8K, otherwise it will not be JITted, should
      // also not be too small, or it will have many function calls (for wide table), see the
      // results in BenchmarkWideTable.
      if (length > splitThreshold) {
        blocks += blockBuilder.toString()
        blockBuilder.clear()
        length = 0
      }
      blockBuilder.append(code)
      length += CodeFormatter.stripExtraNewLinesAndComments(code).length
    }
    blocks += blockBuilder.toString()
    blocks.toSeq
  }

  /**
   * Here we handle all the methods which have been added to the inner classes and
   * not to the outer class.
   * Since they can be many, their direct invocation in the outer class adds many entries
   * to the outer class' constant pool. This can cause the constant pool to past JVM limit.
   * Moreover, this can cause also the outer class method where all the invocations are
   * performed to grow beyond the 64k limit.
   * To avoid these problems, we group them and we call only the grouping methods in the
   * outer class.
   *
   * @param functions a [[Seq]] of [[NewFunctionSpec]] defined in the inner classes
   * @param funcName the split function name base.
   * @param arguments the list of (type, name) of the arguments of the split function.
   * @param returnType the return type of the split function.
   * @param makeSplitFunction makes split function body, e.g. add preparation or cleanup.
   * @param foldFunctions folds the split function calls.
   * @return an [[Iterable]] containing the methods' invocations
   */
  private def generateInnerClassesFunctionCalls(
      functions: Seq[NewFunctionSpec],
      funcName: String,
      arguments: Seq[(String, String)],
      returnType: String,
      makeSplitFunction: String => String,
      foldFunctions: Seq[String] => String): Iterable[String] = {
    val innerClassToFunctions = mutable.LinkedHashMap.empty[(String, String), Seq[String]]
    functions.foreach(f => {
      val key = (f.innerClassName.get, f.innerClassInstance.get)
      val value = f.functionName +: innerClassToFunctions.getOrElse(key, Seq.empty[String])
      innerClassToFunctions.put(key, value)
    })

    val argDefinitionString = arguments.map { case (t, name) => s"$t $name" }.mkString(", ")
    val argInvocationString = arguments.map(_._2).mkString(", ")

    innerClassToFunctions.flatMap {
      case ((innerClassName, innerClassInstance), innerClassFunctions) =>
        // for performance reasons, the functions are prepended, instead of appended,
        // thus here they are in reversed order
        val orderedFunctions = innerClassFunctions.reverse
        if (orderedFunctions.size > MERGE_SPLIT_METHODS_THRESHOLD) {
          // Adding a new function to each inner class which contains the invocation of all the
          // ones which have been added to that inner class. For example,
          //   private class NestedClass {
          //     private void apply_862(InternalRow i) { ... }
          //     private void apply_863(InternalRow i) { ... }
          //       ...
          //     private void apply(InternalRow i) {
          //       apply_862(i);
          //       apply_863(i);
          //       ...
          //     }
          //   }
          val body = foldFunctions(orderedFunctions.map(name => s"$name($argInvocationString)"))
          val code = s"""
              |private $returnType $funcName($argDefinitionString) {
              |  ${makeSplitFunction(body)}
              |}
            """.stripMargin
          addNewFunctionToClass(funcName, code, innerClassName)
          Seq(s"$innerClassInstance.$funcName($argInvocationString)")
        } else {
          orderedFunctions.map(f => s"$innerClassInstance.$f($argInvocationString)")
        }
    }
  }

  /**
   * Returns the code for subexpression elimination after splitting it if necessary.
   */
  def subexprFunctionsCode: String = {
    // Whole-stage codegen's subexpression elimination is handled in another code path
    assert(currentVars == null || subexprFunctions.isEmpty)
    splitExpressions(subexprFunctions.toSeq, "subexprFunc_split", Seq("InternalRow" -> INPUT_ROW))
  }

  /**
   * Perform a function which generates a sequence of ExprCodes with a given mapping between
   * expressions and common expressions, instead of using the mapping in current context.
   */
  def withSubExprEliminationExprs(
      newSubExprEliminationExprs: Map[ExpressionEquals, SubExprEliminationState])(
      f: => Seq[ExprCode]): Seq[ExprCode] = {
    val oldsubExprEliminationExprs = subExprEliminationExprs
    subExprEliminationExprs = newSubExprEliminationExprs

    val genCodes = f

    // Restore previous subExprEliminationExprs
    subExprEliminationExprs = oldsubExprEliminationExprs
    genCodes
  }

  /**
   * Evaluates a sequence of `SubExprEliminationState` which represent subexpressions. After
   * evaluating a subexpression, this method will clean up the code block to avoid duplicate
   * evaluation.
   */
  def evaluateSubExprEliminationState(subExprStates: Iterable[SubExprEliminationState]): String = {
    val code = new StringBuilder()

    subExprStates.foreach { state =>
      val currentCode = evaluateSubExprEliminationState(state.children) + "\n" + state.eval.code
      code.append(currentCode + "\n")
      state.eval.code = EmptyBlock
    }

    code.toString()
  }

  /**
   * Checks and sets up the state and codegen for subexpression elimination in whole-stage codegen.
   *
   * This finds the common subexpressions, generates the code snippets that evaluate those
   * expressions and populates the mapping of common subexpressions to the generated code snippets.
   *
   * The generated code snippet for subexpression is wrapped in `SubExprEliminationState`, which
   * contains an `ExprCode` and the children `SubExprEliminationState` if any. The `ExprCode`
   * includes java source code, result variable name and is-null variable name of the subexpression.
   *
   * Besides, this also returns a sequences of `ExprCode` which are expression codes that need to
   * be evaluated (as their input parameters) before evaluating subexpressions.
   *
   * To evaluate the returned subexpressions, please call `evaluateSubExprEliminationState` with
   * the `SubExprEliminationState`s to be evaluated. During generating the code, it will cleanup
   * the states to avoid duplicate evaluation.
   *
   * The details of subexpression generation:
   *   1. Gets subexpression set. See `EquivalentExpressions`.
   *   2. Generate code of subexpressions as a whole block of code (non-split case)
   *   3. Check if the total length of the above block is larger than the split-threshold. If so,
   *      try to split it in step 4, otherwise returning the non-split code block.
   *   4. Check if parameter lengths of all subexpressions satisfy the JVM limitation, if so,
   *      try to split, otherwise returning the non-split code block.
   *   5. For each subexpression, generating a function and put the code into it. To evaluate the
   *      subexpression, just call the function.
   *
   * The explanation of subexpression codegen:
   *   1. Wrapping in `withSubExprEliminationExprs` call with current subexpression map. Each
   *      subexpression may depends on other subexpressions (children). So when generating code
   *      for subexpressions, we iterate over each subexpression and put the mapping between
   *      (subexpression -> `SubExprEliminationState`) into the map. So in next subexpression
   *      evaluation, we can look for generated subexpressions and do replacement.
   */
  def subexpressionEliminationForWholeStageCodegen(expressions: Seq[Expression]): SubExprCodes = {
    // Create a clear EquivalentExpressions and SubExprEliminationState mapping
    val equivalentExpressions: EquivalentExpressions = new EquivalentExpressions
    val localSubExprEliminationExprsForNonSplit =
      mutable.HashMap.empty[ExpressionEquals, SubExprEliminationState]

    // Add each expression tree and compute the common subexpressions.
    expressions.foreach(equivalentExpressions.addExprTree(_))

    // Get all the expressions that appear at least twice and set up the state for subexpression
    // elimination.
    val commonExprs = equivalentExpressions.getCommonSubexpressions

    val nonSplitCode = {
      val allStates = mutable.ArrayBuffer.empty[SubExprEliminationState]
      commonExprs.map { expr =>
        withSubExprEliminationExprs(localSubExprEliminationExprsForNonSplit.toMap) {
          val eval = expr.genCode(this)
          // Collects other subexpressions from the children.
          val childrenSubExprs = mutable.ArrayBuffer.empty[SubExprEliminationState]
          expr.foreach { e =>
            subExprEliminationExprs.get(ExpressionEquals(e)) match {
              case Some(state) => childrenSubExprs += state
              case _ =>
            }
          }
          val state = SubExprEliminationState(eval, childrenSubExprs.toSeq)
          localSubExprEliminationExprsForNonSplit.put(ExpressionEquals(expr), state)
          allStates += state
          Seq(eval)
        }
      }
      allStates.toSeq
    }

    // For some operators, they do not require all its child's outputs to be evaluated in advance.
    // Instead it only early evaluates part of outputs, for example, `ProjectExec` only early
    // evaluate the outputs used more than twice. So we need to extract these variables used by
    // subexpressions and evaluate them before subexpressions.
    val (inputVarsForAllFuncs, exprCodesNeedEvaluate) = commonExprs.map { expr =>
      val (inputVars, exprCodes) = getLocalInputVariableValues(this, expr)
      (inputVars.toSeq, exprCodes.toSeq)
    }.unzip

    val needSplit = nonSplitCode.map(_.eval.code.length).sum > SQLConf.get.methodSplitThreshold
    val (subExprsMap, exprCodes) = if (needSplit) {
      if (inputVarsForAllFuncs.map(calculateParamLengthFromExprValues).forall(isValidParamLength)) {
        val localSubExprEliminationExprs =
          mutable.HashMap.empty[ExpressionEquals, SubExprEliminationState]

        commonExprs.zipWithIndex.foreach { case (expr, i) =>
          val eval = withSubExprEliminationExprs(localSubExprEliminationExprs.toMap) {
            Seq(expr.genCode(this))
          }.head

          val value = addMutableState(javaType(expr.dataType), "subExprValue")

          val isNullLiteral = eval.isNull match {
            case TrueLiteral | FalseLiteral => true
            case _ => false
          }
          val (isNull, isNullEvalCode) = if (!isNullLiteral) {
            val v = addMutableState(JAVA_BOOLEAN, "subExprIsNull")
            (JavaCode.isNullGlobal(v), s"$v = ${eval.isNull};")
          } else {
            (eval.isNull, "")
          }

          // Generate the code for this expression tree and wrap it in a function.
          val fnName = freshName("subExpr")
          val inputVars = inputVarsForAllFuncs(i)
          val argList =
            inputVars.map(v => s"${CodeGenerator.typeName(v.javaType)} ${v.variableName}")
          val fn =
            s"""
               |private void $fnName(${argList.mkString(", ")}) {
               |  ${eval.code}
               |  $isNullEvalCode
               |  $value = ${eval.value};
               |}
               """.stripMargin

          // Collects other subexpressions from the children.
          val childrenSubExprs = mutable.ArrayBuffer.empty[SubExprEliminationState]
          expr.foreach { e =>
            localSubExprEliminationExprs.get(ExpressionEquals(e)) match {
              case Some(state) => childrenSubExprs += state
              case _ =>
            }
          }

          val inputVariables = inputVars.map(_.variableName).mkString(", ")
          val code = code"${addNewFunction(fnName, fn)}($inputVariables);"
          val state = SubExprEliminationState(
            ExprCode(code, isNull, JavaCode.global(value, expr.dataType)),
            childrenSubExprs.toSeq)
          localSubExprEliminationExprs.put(ExpressionEquals(expr), state)
        }
        (localSubExprEliminationExprs, exprCodesNeedEvaluate)
      } else {
        val errMsg = "Failed to split subexpression code into small functions because " +
          "the parameter length of at least one split function went over the JVM limit: " +
          MAX_JVM_METHOD_PARAMS_LENGTH
        if (Utils.isTesting) {
          throw SparkException.internalError(errMsg)
        } else {
          logInfo(errMsg)
          (localSubExprEliminationExprsForNonSplit, Seq.empty)
        }
      }
    } else {
      (localSubExprEliminationExprsForNonSplit, Seq.empty)
    }
    SubExprCodes(subExprsMap.toMap, exprCodes.flatten)
  }

  /**
   * Checks and sets up the state and codegen for subexpression elimination. This finds the
   * common subexpressions, generates the functions that evaluate those expressions and populates
   * the mapping of common subexpressions to the generated functions.
   */
  private def subexpressionElimination(expressions: Seq[Expression]): Unit = {
    // Add each expression tree and compute the common subexpressions.
    expressions.foreach(equivalentExpressions.addExprTree(_))

    // Get all the expressions that appear at least twice and set up the state for subexpression
    // elimination.
    val commonExprs = equivalentExpressions.getCommonSubexpressions
    commonExprs.foreach { expr =>
      val fnName = freshName("subExpr")
      val isNull = addMutableState(JAVA_BOOLEAN, "subExprIsNull")
      val value = addMutableState(javaType(expr.dataType), "subExprValue")

      // Generate the code for this expression tree and wrap it in a function.
      val eval = expr.genCode(this)
      val fn =
        s"""
           |private void $fnName(InternalRow $INPUT_ROW) {
           |  ${eval.code}
           |  $isNull = ${eval.isNull};
           |  $value = ${eval.value};
           |}
           """.stripMargin

      // Add a state and a mapping of the common subexpressions that are associate with this
      // state. Adding this expression to subExprEliminationExprMap means it will call `fn`
      // when it is code generated. This decision should be a cost based one.
      //
      // The cost of doing subexpression elimination is:
      //   1. Extra function call, although this is probably *good* as the JIT can decide to
      //      inline or not.
      // The benefit doing subexpression elimination is:
      //   1. Running the expression logic. Even for a simple expression, it is likely more than 3
      //      above.
      //   2. Less code.
      // Currently, we will do this for all non-leaf only expression trees (i.e. expr trees with
      // at least two nodes) as the cost of doing it is expected to be low.

      val subExprCode = s"${addNewFunction(fnName, fn)}($INPUT_ROW);"
      subexprFunctions += subExprCode
      val state = SubExprEliminationState(
        ExprCode(code"$subExprCode",
          JavaCode.isNullGlobal(isNull),
          JavaCode.global(value, expr.dataType)))
      subExprEliminationExprs += ExpressionEquals(expr) -> state
    }
  }

  /**
   * Generates code for expressions. If doSubexpressionElimination is true, subexpression
   * elimination will be performed. Subexpression elimination assumes that the code for each
   * expression will be combined in the `expressions` order.
   */
  def generateExpressions(
      expressions: Seq[Expression],
      doSubexpressionElimination: Boolean = false): Seq[ExprCode] = {
    // We need to make sure that we do not reuse stateful expressions. This is needed for codegen
    // as well because some expressions may implement `CodegenFallback`.
    val cleanedExpressions = expressions.map(_.freshCopyIfContainsStatefulExpression())
    if (doSubexpressionElimination) subexpressionElimination(cleanedExpressions)
    cleanedExpressions.map(e => e.genCode(this))
  }

  /**
   * get a map of the pair of a place holder and a corresponding comment
   */
  def getPlaceHolderToComments(): collection.Map[String, String] = placeHolderToComments

  /**
   * Register a comment and return the corresponding place holder
   *
   * @param placeholderId an optionally specified identifier for the comment's placeholder.
   *                      The caller should make sure this identifier is unique within the
   *                      compilation unit. If this argument is not specified, a fresh identifier
   *                      will be automatically created and used as the placeholder.
   * @param force whether to force registering the comments
   */
   def registerComment(
       text: => String,
       placeholderId: String = "",
       force: Boolean = false): Block = {
    if (force || SQLConf.get.codegenComments) {
      val name = if (placeholderId != "") {
        assert(!placeHolderToComments.contains(placeholderId))
        placeholderId
      } else {
        freshName("c")
      }
      val comment = if (text.contains("\n") || text.contains("\r")) {
        text.split("(\r\n)|\r|\n").mkString("/**\n * ", "\n * ", "\n */")
      } else {
        s"// $text"
      }
      placeHolderToComments += (name -> comment)
      code"/*$name*/"
    } else {
      EmptyBlock
    }
  }
}

/**
 * A wrapper for generated class, defines a `generate` method so that we can pass extra objects
 * into generated class.
 */
abstract class GeneratedClass {
  def generate(references: Array[Any]): Any
}

/**
 * A wrapper for the source code to be compiled by [[CodeGenerator]].
 */
class CodeAndComment(val body: String, val comment: collection.Map[String, String])
  extends Serializable {
  override def equals(that: Any): Boolean = that match {
    case t: CodeAndComment if t.body == body => true
    case _ => false
  }

  override def hashCode(): Int = body.hashCode
}

/**
 * A base class for generators of byte code to perform expression evaluation.  Includes a set of
 * helpers for referring to Catalyst types and building trees that perform evaluation of individual
 * expressions.
 */
abstract class CodeGenerator[InType <: AnyRef, OutType <: AnyRef] extends Logging {

  protected val genericMutableRowType: String = classOf[GenericInternalRow].getName

  /**
   * Generates a class for a given input expression.  Called when there is not cached code
   * already available.
   */
  protected def create(in: InType): OutType

  /**
   * Canonicalizes an input expression. Used to avoid double caching expressions that differ only
   * cosmetically.
   */
  protected def canonicalize(in: InType): InType

  /** Binds an input expression to a given input schema */
  protected def bind(in: InType, inputSchema: Seq[Attribute]): InType

  /** Generates the requested evaluator binding the given expression(s) to the inputSchema. */
  def generate(expressions: InType, inputSchema: Seq[Attribute]): OutType =
    generate(bind(expressions, inputSchema))

  /** Generates the requested evaluator given already bound expression(s). */
  def generate(expressions: InType): OutType = create(canonicalize(expressions))

  /**
   * Create a new codegen context for expression evaluator, used to store those
   * expressions that don't support codegen
   */
  def newCodeGenContext(): CodegenContext = {
    new CodegenContext
  }
}

/**
 * Java bytecode statistics of a compiled class by Janino.
 */
case class ByteCodeStats(maxMethodCodeSize: Int, maxConstPoolSize: Int, numInnerClasses: Int)

object ByteCodeStats {
  val UNAVAILABLE = ByteCodeStats(-1, -1, -1)
}

object CodeGenerator extends Logging {

  // This is the default value of HugeMethodLimit in the OpenJDK HotSpot JVM,
  // beyond which methods will be rejected from JIT compilation
  final val DEFAULT_JVM_HUGE_METHOD_LIMIT = 8000

  // The max valid length of method parameters in JVM.
  final val MAX_JVM_METHOD_PARAMS_LENGTH = 255

  // The max number of constant pool entries in JVM.
  final val MAX_JVM_CONSTANT_POOL_SIZE = 65535

  // This is the threshold over which the methods in an inner class are grouped in a single
  // method which is going to be called by the outer class instead of the many small ones
  final val MERGE_SPLIT_METHODS_THRESHOLD = 3

  // The number of named constants that can exist in the class is limited by the Constant Pool
  // limit, 65,536. We cannot know how many constants will be inserted for a class, so we use a
  // threshold of 1000k bytes to determine when a function should be inlined to a private, inner
  // class.
  final val GENERATED_CLASS_SIZE_THRESHOLD = 1000000

  // This is the threshold for the number of global variables, whose types are primitive type or
  // complex type (e.g. more than one-dimensional array), that will be placed at the outer class
  final val OUTER_CLASS_VARIABLES_THRESHOLD = 10000

  // This is the maximum number of array elements to keep global variables in one Java array
  // 32767 is the maximum integer value that does not require a constant pool entry in a Java
  // bytecode instruction
  final val MUTABLESTATEARRAY_SIZE_LIMIT = 32768

  // The Java source code generated by whole-stage codegen on the Driver side is sent to each
  // Executor for compilation and data processing. This is very effective in processing large
  // amounts of data in a distributed environment. However, in the test environment,
  // because the amount of data is not large or not executed in parallel, the compilation time
  // of these Java source code will become a major part of the entire test runtime. When
  // running test cases, we summarize the total compilation time and output it to the execution
  // log for easy analysis and view.
  private val _compileTime = new LongAccumulator

  // Returns the total compile time of Java source code in nanoseconds.
  // Visible for testing
  def compileTime: Long = _compileTime.sum

  // Reset compile time.
  // Visible for testing
  def resetCompileTime(): Unit = _compileTime.reset()

  /**
   * Compile the Java source code into a Java class, using Janino.
   *
   * @return a pair of a generated class and the bytecode statistics of generated functions.
   */
  def compile(code: CodeAndComment): (GeneratedClass, ByteCodeStats) = try {
    val classLoaderRef = new HashableWeakReference(Utils.getContextOrSparkClassLoader)
    cache.get((classLoaderRef, code))
  } catch {
    // Cache.get() may wrap the original exception. See the following URL
    // https://guava.dev/releases/14.0.1/api/docs/com/google/common/cache/
    //   Cache.html#get(K,%20java.util.concurrent.Callable)
    case e @ (_: UncheckedExecutionException | _: ExecutionError) =>
      throw e.getCause
  }

  /**
   * Compile the Java source code into a Java class, using Janino.
   */
  private[this] def doCompile(code: CodeAndComment): (GeneratedClass, ByteCodeStats) = {
    val evaluator = new ClassBodyEvaluator()

    // A special classloader used to wrap the actual parent classloader of
    // [[org.codehaus.janino.ClassBodyEvaluator]] (see CodeGenerator.doCompile). This classloader
    // does not throw a ClassNotFoundException with a cause set (i.e. exception.getCause returns
    // a null). This classloader is needed because janino will throw the exception directly if
    // the parent classloader throws a ClassNotFoundException with cause set instead of trying to
    // find other possible classes (see org.codehaus.janinoClassLoaderIClassLoader's
    // findIClass method). Please also see https://issues.apache.org/jira/browse/SPARK-15622 and
    // https://issues.apache.org/jira/browse/SPARK-11636.
    val parentClassLoader = new ParentClassLoader(Utils.getContextOrSparkClassLoader)
    evaluator.setParentClassLoader(parentClassLoader)
    // Cannot be under package codegen, or fail with java.lang.InstantiationException
    evaluator.setClassName("org.apache.spark.sql.catalyst.expressions.GeneratedClass")
    evaluator.setDefaultImports(
      classOf[Platform].getName,
      classOf[InternalRow].getName,
      classOf[UnsafeRow].getName,
      classOf[UTF8String].getName,
      classOf[Decimal].getName,
      classOf[CalendarInterval].getName,
      classOf[VariantVal].getName,
      classOf[ArrayData].getName,
      classOf[UnsafeArrayData].getName,
      classOf[MapData].getName,
      classOf[UnsafeMapData].getName,
      classOf[Expression].getName,
      classOf[TaskContext].getName,
      classOf[TaskKilledException].getName,
      classOf[InputMetrics].getName,
      classOf[CollationAwareUTF8String].getName,
      classOf[CollationFactory].getName,
      classOf[CollationSupport].getName,
      QueryExecutionErrors.getClass.getName.stripSuffix("$")
    )
    evaluator.setExtendedClass(classOf[GeneratedClass])

    logDebug({
      // Only add extra debugging info to byte code when we are going to print the source code.
      evaluator.setDebuggingInformation(true, true, false)
      s"\n${CodeFormatter.format(code)}"
    })

    val codeStats = try {
      evaluator.cook("generated.java", code.body)
      updateAndGetCompilationStats(evaluator)
    } catch {
      case e: InternalCompilerException =>
        logError("Failed to compile the generated Java code.", e)
        logGeneratedCode(code)
        throw QueryExecutionErrors.internalCompilerError(e)
      case e: CompileException =>
        logError("Failed to compile the generated Java code.", e)
        logGeneratedCode(code)
        throw QueryExecutionErrors.compilerError(e)
    }

    (evaluator.getClazz().getConstructor().newInstance().asInstanceOf[GeneratedClass], codeStats)
  }

  private def logGeneratedCode(code: CodeAndComment): Unit = {
    val maxLines = SQLConf.get.loggingMaxLinesForCodegen
    if (Utils.isTesting) {
      logError(s"\n${CodeFormatter.format(code, maxLines)}")
    } else {
      logInfo(s"\n${CodeFormatter.format(code, maxLines)}")
    }
  }

  /**
   * Returns the bytecode statistics (max method bytecode size, max constant pool size, and
   * # of inner classes) of generated classes by inspecting Janino classes.
   * Also, this method updates the metrics information.
   */
  private def updateAndGetCompilationStats(evaluator: ClassBodyEvaluator): ByteCodeStats = {
    // First retrieve the generated classes.
    val classes = evaluator.getBytecodes.asScala

    // Then walk the classes to get at the method bytecode.
    val codeStats = classes.map { case (_, classBytes) =>
      val classCodeSize = classBytes.length
      CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(classCodeSize)
      try {
        val cf = new ClassFile(new ByteArrayInputStream(classBytes))
        val constPoolSize = cf.getConstantPoolSize
        val methodCodeSizes = cf.methodInfos.asScala.flatMap { method =>
          method.getAttributes.collect { case attr: CodeAttribute =>
            val byteCodeSize = attr.code.length
            CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize)

            if (byteCodeSize > DEFAULT_JVM_HUGE_METHOD_LIMIT) {
              logInfo(log"Generated method too long to be JIT compiled: " +
                log"${MDC(LogKeys.CLASS_NAME, cf.getThisClassName)}." +
                log"${MDC(LogKeys.METHOD_NAME, method.getName)} is " +
                log"${MDC(LogKeys.BYTECODE_SIZE, byteCodeSize)} bytes")
            }

            byteCodeSize
          }
        }
        (methodCodeSizes.max, constPoolSize)
      } catch {
        case NonFatal(e) =>
          logWarning("Error calculating stats of compiled class.", e)
          (-1, -1)
      }
    }

    val (maxMethodSizes, constPoolSize) = codeStats.unzip
    ByteCodeStats(
      maxMethodCodeSize = maxMethodSizes.max,
      maxConstPoolSize = constPoolSize.max,
      // Minus 2 for `GeneratedClass` and an outer-most generated class
      numInnerClasses = classes.size - 2)
  }

  /**
   * A cache of generated classes.
   *
   * From the Guava Docs: A Cache is similar to ConcurrentMap, but not quite the same. The most
   * fundamental difference is that a ConcurrentMap persists all elements that are added to it until
   * they are explicitly removed. A Cache on the other hand is generally configured to evict entries
   * automatically, in order to constrain its memory footprint.  Note that this cache does not use
   * weak keys/values and thus does not respond to memory pressure.
   *
   * Codegen can be slow. Use a non fate sharing cache in case a query gets canceled during codegen
   * while other queries wait on the same code, so that those other queries don't get wrongly
   * aborted. See [[NonFateSharingCache]] for more details.
   */
  private val cache = {
    val loadFunc: ((HashableWeakReference, CodeAndComment)) => (GeneratedClass, ByteCodeStats) = {
      case (_, code) =>
        val startTime = System.nanoTime()
        val result = doCompile(code)
        val endTime = System.nanoTime()
        val duration = endTime - startTime
        val timeMs: Double = duration.toDouble / NANOS_PER_MILLIS
        CodegenMetrics.METRIC_SOURCE_CODE_SIZE.update(code.body.length)
        CodegenMetrics.METRIC_COMPILATION_TIME.update(timeMs.toLong)
        logInfo(log"Code generated in ${MDC(TOTAL_TIME, timeMs)} ms")
        _compileTime.add(duration)
        result
    }
    NonFateSharingCache(loadFunc, SQLConf.get.codegenCacheMaxEntries)
  }

  /**
   * Name of Java primitive data type
   */
  final val JAVA_BOOLEAN = "boolean"
  final val JAVA_BYTE = "byte"
  final val JAVA_SHORT = "short"
  final val JAVA_INT = "int"
  final val JAVA_LONG = "long"
  final val JAVA_FLOAT = "float"
  final val JAVA_DOUBLE = "double"

  /**
   * List of java primitive data types
   */
  val primitiveTypes =
    Seq(JAVA_BOOLEAN, JAVA_BYTE, JAVA_SHORT, JAVA_INT, JAVA_LONG, JAVA_FLOAT, JAVA_DOUBLE)

  /**
   * Returns true if a Java type is Java primitive primitive type
   */
  def isPrimitiveType(jt: String): Boolean = primitiveTypes.contains(jt)

  def isPrimitiveType(dt: DataType): Boolean = isPrimitiveType(javaType(dt))

  /**
   * Returns the specialized code to access a value from `inputRow` at `ordinal`.
   */
  @tailrec
  def getValue(input: String, dataType: DataType, ordinal: String): String = {
    val jt = javaType(dataType)
    dataType match {
      case udt: UserDefinedType[_] => getValue(input, udt.sqlType, ordinal)
      case _ if isPrimitiveType(jt) => s"$input.get${primitiveTypeName(jt)}($ordinal)"
      case _ => PhysicalDataType(dataType) match {
        case _: PhysicalArrayType => s"$input.getArray($ordinal)"
        case PhysicalBinaryType => s"$input.getBinary($ordinal)"
        case PhysicalCalendarIntervalType => s"$input.getInterval($ordinal)"
        case t: PhysicalDecimalType => s"$input.getDecimal($ordinal, ${t.precision}, ${t.scale})"
        case _: PhysicalMapType => s"$input.getMap($ordinal)"
        case PhysicalNullType => "null"
        case _: PhysicalStringType => s"$input.getUTF8String($ordinal)"
        case t: PhysicalStructType => s"$input.getStruct($ordinal, ${t.fields.length})"
        case PhysicalVariantType => s"$input.getVariant($ordinal)"
        case _ => s"($jt)$input.get($ordinal, null)"
      }
    }
  }

  /**
   * Generates code creating a [[UnsafeArrayData]] or
   * [[org.apache.spark.sql.catalyst.util.GenericArrayData]] based on given parameters.
   *
   * @param arrayName name of the array to create
   * @param elementType data type of the elements in source array
   * @param numElements code representing the number of elements the array should contain
   * @param additionalErrorMessage string to include in the error message
   *
   * @return code representing the allocation of [[ArrayData]]
   */
  def createArrayData(
      arrayName: String,
      elementType: DataType,
      numElements: String,
      additionalErrorMessage: String): String = {
    val elementSize = if (CodeGenerator.isPrimitiveType(elementType)) {
      elementType.defaultSize
    } else {
      -1
    }
    s"""
       |ArrayData $arrayName = ArrayData.allocateArrayData(
       |  $elementSize, $numElements, "$additionalErrorMessage");
     """.stripMargin
  }

  /**
   * Generates assignment code for an [[ArrayData]]
   *
   * @param dstArray name of the array to be assigned
   * @param elementType data type of the elements in destination and source arrays
   * @param srcArray name of the array to be read
   * @param needNullCheck value which shows whether a nullcheck is required for the returning
   *                      assignment
   * @param dstArrayIndex an index variable to access each element of destination array
   * @param srcArrayIndex an index variable to access each element of source array
   *
   * @return code representing an assignment to each element of the [[ArrayData]], which requires
   *         a pair of destination and source loop index variables
   */
  def createArrayAssignment(
      dstArray: String,
      elementType: DataType,
      srcArray: String,
      dstArrayIndex: String,
      srcArrayIndex: String,
      needNullCheck: Boolean): String = {
    CodeGenerator.setArrayElement(dstArray, elementType, dstArrayIndex,
      CodeGenerator.getValue(srcArray, elementType, srcArrayIndex),
      if (needNullCheck) Some(s"$srcArray.isNullAt($srcArrayIndex)") else None)
  }

  /**
   * Returns the code to update a column in Row for a given DataType.
   */
  @tailrec
  def setColumn(row: String, dataType: DataType, ordinal: Int, value: String): String = {
    val jt = javaType(dataType)
    dataType match {
      case _ if isPrimitiveType(jt) => s"$row.set${primitiveTypeName(jt)}($ordinal, $value)"
      case CalendarIntervalType => s"$row.setInterval($ordinal, $value)"
      case t: DecimalType => s"$row.setDecimal($ordinal, $value, ${t.precision})"
      case udt: UserDefinedType[_] => setColumn(row, udt.sqlType, ordinal, value)
      // The UTF8String, InternalRow, ArrayData and MapData may came from UnsafeRow, we should copy
      // it to avoid keeping a "pointer" to a memory region which may get updated afterwards.
      case _: StringType | _: StructType | _: ArrayType | _: MapType =>
        s"$row.update($ordinal, $value.copy())"
      case _ => s"$row.update($ordinal, $value)"
    }
  }

  /**
   * Update a column in MutableRow from ExprCode.
   *
   * @param isVectorized True if the underlying row is of type `ColumnarBatch.Row`, false otherwise
   */
  def updateColumn(
      row: String,
      dataType: DataType,
      ordinal: Int,
      ev: ExprCode,
      nullable: Boolean,
      isVectorized: Boolean = false): String = {
    if (nullable) {
      // Can't call setNullAt on DecimalType/CalendarIntervalType, because we need to keep the
      // offset
      if (!isVectorized && UnsafeRowUtils.avoidSetNullAt(dataType)) {
        s"""
           |if (!${ev.isNull}) {
           |  ${setColumn(row, dataType, ordinal, ev.value)};
           |} else {
           |  ${setColumn(row, dataType, ordinal, "null")};
           |}
         """.stripMargin
      } else {
        s"""
           |if (!${ev.isNull}) {
           |  ${setColumn(row, dataType, ordinal, ev.value)};
           |} else {
           |  $row.setNullAt($ordinal);
           |}
         """.stripMargin
      }
    } else {
      s"""${setColumn(row, dataType, ordinal, ev.value)};"""
    }
  }

  /**
   * Returns the specialized code to set a given value in a column vector for a given `DataType`.
   */
  def setValue(vector: String, rowId: String, dataType: DataType, value: String): String = {
    val jt = javaType(dataType)
    dataType match {
      case _ if isPrimitiveType(jt) =>
        s"$vector.put${primitiveTypeName(jt)}($rowId, $value);"
      case t: DecimalType => s"$vector.putDecimal($rowId, $value, ${t.precision});"
      case CalendarIntervalType => s"$vector.putInterval($rowId, $value);"
      case _: StringType => s"$vector.putByteArray($rowId, $value.getBytes());"
      case _ =>
        throw new SparkIllegalArgumentException(
          errorClass = "_LEGACY_ERROR_TEMP_3233",
          messageParameters = Map("dataType" -> dataType.toString))
    }
  }

  /**
   * Generates code of setter for an [[ArrayData]].
   */
  def setArrayElement(
      array: String,
      elementType: DataType,
      i: String,
      value: String,
      isNull: Option[String] = None): String = {
    val isPrimitiveType = CodeGenerator.isPrimitiveType(elementType)
    val setFunc = if (isPrimitiveType) {
      s"set${CodeGenerator.primitiveTypeName(elementType)}"
    } else {
      "update"
    }
    if (isNull.isDefined && isPrimitiveType) {
      s"""
         |if (${isNull.get}) {
         |  $array.setNullAt($i);
         |} else {
         |  $array.$setFunc($i, $value);
         |}
       """.stripMargin
    } else {
      s"$array.$setFunc($i, $value);"
    }
  }

  /**
   * Returns the specialized code to set a given value in a column vector for a given `DataType`
   * that could potentially be nullable.
   */
  def updateColumn(
      vector: String,
      rowId: String,
      dataType: DataType,
      ev: ExprCode,
      nullable: Boolean): String = {
    if (nullable) {
      s"""
         |if (!${ev.isNull}) {
         |  ${setValue(vector, rowId, dataType, ev.value)}
         |} else {
         |  $vector.putNull($rowId);
         |}
       """.stripMargin
    } else {
      s"""${setValue(vector, rowId, dataType, ev.value)};"""
    }
  }

  /**
   * Returns the specialized code to access a value from a column vector for a given `DataType`.
   */
  def getValueFromVector(vector: String, dataType: DataType, rowId: String): String = {
    val sqlDataType = dataType match {
      case udt: UserDefinedType[_] => udt.sqlType
      case _ => dataType
    }

    if (sqlDataType.isInstanceOf[StructType]) {
      // `ColumnVector.getStruct` is different from `InternalRow.getStruct`, it only takes an
      // `ordinal` parameter.
      s"$vector.getStruct($rowId)"
    } else {
      getValue(vector, sqlDataType, rowId)
    }
  }

  /**
   * This methods returns two values in a Tuple.
   *
   * First value: Extracts all the input variables from references and subexpression
   * elimination states for a given `expr`. This result will be used to split the
   * generated code of expressions into multiple functions.
   *
   * Second value: Returns the set of `ExprCodes`s which are necessary codes before
   * evaluating subexpressions.
   */
  def getLocalInputVariableValues(
      ctx: CodegenContext,
      expr: Expression,
      subExprs: Map[ExpressionEquals, SubExprEliminationState] = Map.empty)
      : (Set[VariableValue], Set[ExprCode]) = {
    val argSet = mutable.Set[VariableValue]()
    val exprCodesNeedEvaluate = mutable.Set[ExprCode]()

    if (ctx.INPUT_ROW != null) {
      argSet += JavaCode.variable(ctx.INPUT_ROW, classOf[InternalRow])
    }

    // Collects local variables from a given `expr` tree
    val collectLocalVariable = (ev: ExprValue) => ev match {
      case vv: VariableValue => argSet += vv
      case _ =>
    }

    val stack = mutable.Stack[Expression](expr)
    while (stack.nonEmpty) {
      stack.pop() match {
        case ref: BoundReference if ctx.currentVars != null &&
            ctx.currentVars(ref.ordinal) != null =>
          val exprCode = ctx.currentVars(ref.ordinal)
          // If the referred variable is not evaluated yet.
          if (exprCode.code != EmptyBlock) {
            exprCodesNeedEvaluate += exprCode.copy()
            exprCode.code = EmptyBlock
          }
          collectLocalVariable(exprCode.value)
          collectLocalVariable(exprCode.isNull)

        case e =>
          subExprs.get(ExpressionEquals(e)) match {
            case Some(state) =>
              collectLocalVariable(state.eval.value)
              collectLocalVariable(state.eval.isNull)
            case None =>
              stack.pushAll(e.children)
          }
      }
    }

    (argSet.toSet, exprCodesNeedEvaluate.toSet)
  }

  /**
   * Returns the name used in accessor and setter for a Java primitive type.
   */
  def primitiveTypeName(jt: String): String = jt match {
    case JAVA_INT => "Int"
    case _ => boxedType(jt)
  }

  def primitiveTypeName(dt: DataType): String = primitiveTypeName(javaType(dt))

  /**
   * Returns the Java type for a DataType.
   */
  def javaType(dt: DataType): String = dt match {
    case udt: UserDefinedType[_] => javaType(udt.sqlType)
    case ObjectType(cls) if cls.isArray => s"${javaType(ObjectType(cls.getComponentType))}[]"
    case ObjectType(cls) => cls.getName
    case _ => PhysicalDataType(dt) match {
      case _: PhysicalArrayType => "ArrayData"
      case PhysicalBinaryType => "byte[]"
      case PhysicalBooleanType => JAVA_BOOLEAN
      case PhysicalByteType => JAVA_BYTE
      case PhysicalCalendarIntervalType => "CalendarInterval"
      case PhysicalIntegerType => JAVA_INT
      case _: PhysicalDecimalType => "Decimal"
      case PhysicalDoubleType => JAVA_DOUBLE
      case PhysicalFloatType => JAVA_FLOAT
      case PhysicalLongType => JAVA_LONG
      case _: PhysicalMapType => "MapData"
      case PhysicalShortType => JAVA_SHORT
      case _: PhysicalStringType => "UTF8String"
      case _: PhysicalStructType => "InternalRow"
      case _: PhysicalVariantType => "VariantVal"
      case _ => "Object"
    }
  }

  @tailrec
  def javaClass(dt: DataType): Class[_] = dt match {
    case BooleanType => java.lang.Boolean.TYPE
    case ByteType => java.lang.Byte.TYPE
    case ShortType => java.lang.Short.TYPE
    case IntegerType | DateType | _: YearMonthIntervalType => java.lang.Integer.TYPE
    case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType =>
      java.lang.Long.TYPE
    case FloatType => java.lang.Float.TYPE
    case DoubleType => java.lang.Double.TYPE
    case _: DecimalType => classOf[Decimal]
    case BinaryType => classOf[Array[Byte]]
    case _: StringType => classOf[UTF8String]
    case CalendarIntervalType => classOf[CalendarInterval]
    case _: StructType => classOf[InternalRow]
    case _: ArrayType => classOf[ArrayData]
    case _: MapType => classOf[MapData]
    case udt: UserDefinedType[_] => javaClass(udt.sqlType)
    case ObjectType(cls) => cls
    case VariantType => classOf[VariantVal]
    case _ => classOf[Object]
  }

  /**
   * Returns the boxed type in Java.
   */
  def boxedType(jt: String): String = jt match {
    case JAVA_BOOLEAN => "Boolean"
    case JAVA_BYTE => "Byte"
    case JAVA_SHORT => "Short"
    case JAVA_INT => "Integer"
    case JAVA_LONG => "Long"
    case JAVA_FLOAT => "Float"
    case JAVA_DOUBLE => "Double"
    case other => other
  }

  def boxedType(dt: DataType): String = boxedType(javaType(dt))

  def typeName(clazz: Class[_]): String = {
    if (clazz.isArray) {
      typeName(clazz.getComponentType) + "[]"
    } else {
      clazz.getName
    }
  }

  /**
   * Returns the representation of default value for a given Java Type.
   * @param jt the string name of the Java type
   * @param typedNull if true, for null literals, return a typed (with a cast) version
   */
  def defaultValue(jt: String, typedNull: Boolean): String = jt match {
    case JAVA_BOOLEAN => "false"
    case JAVA_BYTE => "(byte)-1"
    case JAVA_SHORT => "(short)-1"
    case JAVA_INT => "-1"
    case JAVA_LONG => "-1L"
    case JAVA_FLOAT => "-1.0f"
    case JAVA_DOUBLE => "-1.0"
    case _ => if (typedNull) s"(($jt)null)" else "null"
  }

  def defaultValue(dt: DataType, typedNull: Boolean = false): String =
    defaultValue(javaType(dt), typedNull)

  /**
   * Returns the length of parameters for a Java method descriptor. `this` contributes one unit
   * and a parameter of type long or double contributes two units. Besides, for nullable parameter,
   * we also need to pass a boolean parameter for the null status.
   */
  def calculateParamLength(params: Seq[Expression]): Int = {
    def paramLengthForExpr(input: Expression): Int = {
      val javaParamLength = javaType(input.dataType) match {
        case JAVA_LONG | JAVA_DOUBLE => 2
        case _ => 1
      }
      // For a nullable expression, we need to pass in an extra boolean parameter.
      (if (input.nullable) 1 else 0) + javaParamLength
    }
    // Initial value is 1 for `this`.
    1 + params.map(paramLengthForExpr).sum
  }

  def calculateParamLengthFromExprValues(params: Seq[ExprValue]): Int = {
    def paramLengthForExpr(input: ExprValue): Int = input.javaType match {
      case java.lang.Long.TYPE | java.lang.Double.TYPE => 2
      case _ => 1
    }
    // Initial value is 1 for `this`.
    1 + params.map(paramLengthForExpr).sum
  }

  /**
   * In Java, a method descriptor is valid only if it represents method parameters with a total
   * length less than a pre-defined constant.
   */
  def isValidParamLength(paramLength: Int): Boolean = {
    // This config is only for testing
    SQLConf.get.getConfString("spark.sql.CodeGenerator.validParamLength", null) match {
      case null | "" => paramLength <= MAX_JVM_METHOD_PARAMS_LENGTH
      case validLength => paramLength <= validLength.toInt
    }
  }
}
