Tasty Reflection: Tree Traversal

[anatoliykmetyuk]

// How to use TreeAccumulator with Unit
val acc = new TreeAccumulator[Unit] {
  def foldTree(x: Unit, tree: Tree)(implicit ctx: Context): Unit = {
    println(s"Tree ${tree.show}")
    foldOverTree(x, tree)
  }
  def foldPattern(x: Unit, tree: Pattern)(implicit ctx: Context): Unit = {
    println(s"Pattren ${tree.show}")
    foldOverPattern(x, tree)
  }
}
acc.foldTree((), tree)

// Equivalent of the above code
val trs = new TreeTraverser {
  override def traverseTree(tree: Tree)(implicit ctx: Context): Unit = {
    println(s"Tree ${tree.show}")
    traverseTreeChildren(tree)
  }

  override def traversePattern(tree: Pattern)(implicit ctx: Context): Unit = {
    println(s"Pattern ${tree.show}")
    traversePatternChildren(tree)
  }
}
trs.traverseTree(tree)

// How to replace the RHS of a variable called "name" to String "Dog"
val map = new TreeMap {
  override def transformStatement(tree: Statement)(implicit ctx: Context): Statement = {
    def localCtx(definition: Definition): Context = definition.symbol.localContext
    tree match {
      case IsValDef(tree) =>
        implicit val ctx = localCtx(tree)
        val tpt1 = transformTypeTree(tree.tpt)
        val rhs1 = tree.rhs.map { x =>
          if (tree.name == "name") Literal(Constant.String("Dog")) else transformTerm(x)
        }
        ValDef.copy(tree)(tree.name, tpt1, rhs1)
      case _ => super.transformStatement(tree)
    }
  }
}
map.transformTerm(tree.underlyingArgument).seal

Usage:

f {
  println("Hello World")
  val name = "Cat"
  println(s"I said hello world, $name")
}

Where f is the macro doing the transformations above.