scala.meta.contrib.TreeOps.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of scalameta_2.11 Show documentation

Scalameta umbrella module that includes all public APIs

There is a newer version: 4.12.2

package scala.meta
package contrib

import scala.meta.contrib.equality.Equal
import scala.meta.contrib.equality.TreeEquality

import scala.annotation.tailrec
import scala.language.higherKinds

object TreeOps {

  class SimpleTraverser {
    def apply(tree: Tree): Unit = tree.children.foreach(apply)
  }

  def contains[F[x <: Tree] <: TreeEquality[x]](
      tree: Tree
  )(toFind: Tree)(implicit conv: Tree => F[Tree], eqEv: Equal[F[Tree]]): Boolean =
    find(tree)(_.isEqual[F](toFind)).nonEmpty

  def find(tree: Tree)(f: Tree => Boolean): Option[Tree] =
    collectFirst(tree) { case x if f(x) => x }

  def forall(tree: Tree)(f: Tree => Boolean): Boolean = find(tree)(t => !f(t)).isEmpty

  def exists(tree: Tree)(f: Tree => Boolean): Boolean = find(tree)(t => f(t)).isDefined

  def collectFirst[B](tree: Tree)(pf: PartialFunction[Tree, B]): Option[B] = {
    var result = Option.empty[B]
    object traverser extends SimpleTraverser {
      override def apply(t: Tree): Unit =
        if (result.isEmpty && pf.isDefinedAt(t)) result = Some(pf(t))
        else if (result.isEmpty) super.apply(t)
    }
    traverser(tree)
    result
  }

  def foreach(tree: Tree)(f: Tree => Unit): Unit = {
    object traverser extends SimpleTraverser {
      override def apply(t: Tree): Unit = {
        f(t)
        super.apply(t)
      }
    }
    traverser(tree)
  }

  def descendants(tree: Tree): List[Tree] = {
    val builder = List.newBuilder[Tree]
    object traverser extends SimpleTraverser {
      override def apply(t: Tree): Unit = {
        if (t ne tree) builder += t
        super.apply(t)
      }
    }
    traverser(tree)
    builder.result()
  }

  def collect[B](tree: Tree)(pf: PartialFunction[Tree, B]): List[B] = {
    val builder = List.newBuilder[B]
    object traverser extends SimpleTraverser {
      override def apply(t: Tree): Unit = {
        if (pf.isDefinedAt(t)) builder += pf(t)
        super.apply(t)
      }
    }
    traverser(tree)
    builder.result()
  }

  @tailrec
  final def ancestors(tree: Tree, accum: List[Tree] = Nil): List[Tree] = tree.parent match {
    case Some(parent) => ancestors(parent, parent +: accum)
    case _ => accum
  }
}