Spark Sql 源码剖析（二）： TreeNode

// 所有 products 的基trait，至少包含 [[scala.Product1]] 至 [[scala.Product22]] 及 [[scala.Tuple1]] 至 [[scala.Tuple22]]
trait Product extends Any with Equals {
  // 第 n 个元素，从0开始
  def productElement(n: Int): Any

  // product size
  def productArity: Int

  // product 遍及所有元素的迭代器
  def productIterator: Iterator[Any] = new scala.collection.AbstractIterator[Any] {
    private var c: Int = 0
    private val cmax = productArity
    def hasNext = c < cmax
    def next() = { val result = productElement(c); c += 1; result }
  }
}

// Origin 表示第几行第几列
case class Origin(
  line: Option[Int] = None,
  startPosition: Option[Int] = None)

  def foreach(f: BaseType => Unit): Unit = {
    f(this)
    children.foreach(_.foreach(f))
  }

  def map[A](f: BaseType => A): Seq[A] = {
    val ret = new collection.mutable.ArrayBuffer[A]()
    foreach(ret += f(_))
    ret
  }

  def collect[B](pf: PartialFunction[BaseType, B]): Seq[B] = {
    val ret = new collection.mutable.ArrayBuffer[B]()
    val lifted = pf.lift
    foreach(node => lifted(node).foreach(ret.+=))
    ret
  }

  def transformDown(rule: PartialFunction[BaseType, BaseType]): BaseType = {
    val afterRule = CurrentOrigin.withOrigin(origin) {
        // 如果 this 是 BaseType 或其子类，则对 this 应用 rule 再返回应用 rule 后的结果，否则返回 this 
      rule.applyOrElse(this, identity[BaseType])
    }

    // Check if unchanged and then possibly return old copy to avoid gc churn.
    if (this fastEquals afterRule) {
        // 如果应用了 rule 后节点无变化，则递归将 rule 应用于 children
      mapChildren(_.transformDown(rule))
    } else {
        // 如果应用了 rule 后节点有变化，则本节点换成变化后的节点（children 不变），再将 rule 递归应用于子节点。也就是从根节点往下来应用 rule 替换节点
      afterRule.mapChildren(_.transformDown(rule))
    }
  }

case arg: TreeNode[_] if containsChild(arg) =>
  val newChild = f(arg.asInstanceOf[BaseType])
  if (!(newChild fastEquals arg)) {
    changed = true
    newChild
  } else {
    arg
  }
case Some(arg: TreeNode[_]) if containsChild(arg) =>
  val newChild = f(arg.asInstanceOf[BaseType])
  if (!(newChild fastEquals arg)) {
    changed = true
    Some(newChild)
  } else {
    Some(arg)
  }
case m: Map[_, _] => m.mapValues {
  case arg: TreeNode[_] if containsChild(arg) =>
    val newChild = f(arg.asInstanceOf[BaseType])
    if (!(newChild fastEquals arg)) {
      changed = true
      newChild
    } else {
      arg
    }
  case other => other
}.view.force // `mapValues` is lazy and we need to force it to materialize
case d: DataType => d // Avoid unpacking Structs
case args: Traversable[_] => args.map {
  case arg: TreeNode[_] if containsChild(arg) =>
    val newChild = f(arg.asInstanceOf[BaseType])
    if (!(newChild fastEquals arg)) {
      changed = true
      newChild
    } else {
      arg
    }
  case tuple@(arg1: TreeNode[_], arg2: TreeNode[_]) =>
    val newChild1 = if (containsChild(arg1)) {
      f(arg1.asInstanceOf[BaseType])
    } else {
      arg1.asInstanceOf[BaseType]
    }

    val newChild2 = if (containsChild(arg2)) {
      f(arg2.asInstanceOf[BaseType])
    } else {
      arg2.asInstanceOf[BaseType]
    }

    if (!(newChild1 fastEquals arg1) || !(newChild2 fastEquals arg2)) {
      changed = true
      (newChild1, newChild2)
    } else {
      tuple
    }
  case other => other
}
case nonChild: AnyRef => nonChild
case null => null

private def getNodeNumbered(number: MutableInt): Option[TreeNode[_]] = {
    if (number.i < 0) {
      None
    } else if (number.i == 0) {
      // 返回根节点
      Some(this)
    } else {
      number.i -= 1
      // 注意，此遍历顺序必须与numberedTreeString相同
      innerChildren.map(_.getNodeNumbered(number)).find(_ != None).getOrElse {
        children.map(_.getNodeNumbered(number)).find(_ != None).flatten
      }
    }
  }