1 files changed, 5 insertions, 9 deletions
diff --git a/src/main/scala/scala/async/internal/TransformUtils.scala b/src/main/scala/scala/async/internal/TransformUtils.scala
index 7602905..18e994f 100644
--- a/src/main/scala/scala/async/internal/TransformUtils.scala
+++ b/src/main/scala/scala/async/internal/TransformUtils.scala
@@ -215,15 +215,11 @@ private[async] trait TransformUtils {
     localTyper = global.analyzer.newTyper(callSiteTyper.context.make(unit = callSiteTyper.context.unit))
   }
 
-  def transformAt(tree: Tree)(f: PartialFunction[Tree, (analyzer.Context => Tree)]) = {
-    object trans extends MacroTypingTransformer {
-      override def transform(tree: Tree): Tree = {
-        if (f.isDefinedAt(tree)) {
-          f(tree)(localTyper.context)
-        } else super.transform(tree)
-      }
-    }
-    trans.transform(tree)
+  def transformAt(tree: Tree)(f: PartialFunction[Tree, (TypingTransformApi => Tree)]) = {
+    typingTransform(tree)((tree, api) => {
+      if (f.isDefinedAt(tree)) f(tree)(api)
+      else api.default(tree)
+    })
   }
 
   def toMultiMap[A, B](as: Iterable[(A, B)]): Map[A, List[B]] =