Add selective ANF transform

- Does not descend into class and module defs
- Adds several tests, including tests for if-else

1 files changed, 10 insertions, 2 deletions

diff --git a/src/main/scala/scala/async/Async.scala b/src/main/scala/scala/async/Async.scala
index b71ce74..8fc7ccf 100644
--- a/src/main/scala/scala/async/Async.scala
+++ b/src/main/scala/scala/async/Async.scala
@@ -74,7 +74,15 @@ abstract class AsyncBase {
     import builder.defn._
     import builder.name
     import builder.futureSystemOps
-    val (stats, expr) = body.tree match {
+
+    val transform = new AnfTransform[c.type](c)
+    val typedBody = c.typeCheck(body.tree)
+    val stats1 :+ expr1 = transform.anf.transformToList(typedBody)
+    val btree = c.typeCheck(Block(stats1, expr1))
+
+    AsyncUtils.vprintln(s"ANF transform expands to:\n $btree")
+
+    val (stats, expr) = btree match {
       case Block(stats, expr) => (stats, expr)
       case tree => (Nil, tree)
     }
@@ -86,7 +94,7 @@ abstract class AsyncBase {
     val handlerCases: List[CaseDef] = asyncBlockBuilder.mkCombinedHandlerCases[T]()
 
     val initStates = asyncBlockBuilder.asyncStates.init
-    val localVarTrees = initStates.flatMap(_.allVarDefs).toList
+    val localVarTrees = asyncBlockBuilder.asyncStates.flatMap(_.allVarDefs).toList
 
     /*
       lazy val onCompleteHandler = (tr: Try[Any]) => state match {