Refactor name substitution out of ExprBuilder.

2 files changed, 17 insertions, 12 deletions

diff --git a/src/main/scala/scala/async/ExprBuilder.scala b/src/main/scala/scala/async/ExprBuilder.scala
index a3ba53e..415c47d 100644
--- a/src/main/scala/scala/async/ExprBuilder.scala
+++ b/src/main/scala/scala/async/ExprBuilder.scala
@@ -123,18 +123,11 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](va
     var nextState    : Int         = -1
     var nextJumpState: Option[Int] = None
 
-    private val renamer = new Transformer {
-      override def transform(tree: Tree) = tree match {
-        case Ident(_) if nameMap.keySet contains tree.symbol =>
-          Ident(nameMap(tree.symbol))
-        case _                                               =>
-          super.transform(tree)
-      }
-    }
+    def rename(tree: Tree) = substituteNames(tree, nameMap)
 
     def +=(stat: c.Tree): this.type = {
       assert(nextJumpState.isEmpty, s"statement appeared after a label jump: $stat")
-      def addStat() = stats += resetDuplicate(renamer.transform(stat))
+      def addStat() = stats += resetDuplicate(rename(stat))
       stat match {
         case Apply(fun, Nil) =>
           labelDefStates get fun.symbol match {
@@ -169,8 +162,7 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](va
      */
     def complete(awaitArg: c.Tree, awaitResultName: TermName, awaitResultType: Tree,
                  nextState: Int): this.type = {
-      val renamed = renamer.transform(awaitArg)
-      awaitable = resetDuplicate(renamed)
+      awaitable = resetDuplicate(rename(awaitArg))
       resultName = awaitResultName
       resultType = awaitResultType.tpe
       this.nextState = nextState
@@ -185,7 +177,7 @@ private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](va
     def resultWithIf(condTree: c.Tree, thenState: Int, elseState: Int): AsyncState = {
       // 1. build changed if-else tree
       // 2. insert that tree at the end of the current state
-      val cond = resetDuplicate(renamer.transform(condTree))
+      val cond = resetDuplicate(rename(condTree))
       def mkBranch(state: Int) = Block(mkStateTree(state), mkResumeApply)
       this += If(cond, mkBranch(thenState), mkBranch(elseState))
       new AsyncStateWithoutAwait(stats.toList, state)
diff --git a/src/main/scala/scala/async/TransformUtils.scala b/src/main/scala/scala/async/TransformUtils.scala
index b7bb2ce..03709ab 100644
--- a/src/main/scala/scala/async/TransformUtils.scala
+++ b/src/main/scala/scala/async/TransformUtils.scala
@@ -49,6 +49,19 @@ private[async] final case class TransformUtils[C <: Context](val c: C) {
   def isAwait(fun: Tree) =
     fun.symbol == defn.Async_await
 
+  /** Replace all `Ident` nodes referring to one of the keys n `renameMap` with a node
+    * referring to the corresponding new name
+    */
+  def substituteNames(tree: Tree, renameMap: Map[Symbol, Name]): Tree = {
+    val renamer = new Transformer {
+      override def transform(tree: Tree) = tree match {
+        case Ident(_) => (renameMap get tree.symbol).fold(tree)(Ident(_))
+        case _        => super.transform(tree)
+      }
+    }
+    renamer.transform(tree)
+  }
+
   /** Descends into the regions of the tree that are subject to the
     * translation to a state machine by `async`. When a nested template,
     * function, or by-name argument is encountered, the descend stops,