No need to extract vals for inlinable args.

We use `isExprSafeToInline` from the non-public reflection
API to check. In addtion, we now that an untyped Ident("await$N")
is also an inlinable expression.

3 files changed, 23 insertions, 15 deletions

diff --git a/src/main/scala/scala/async/AnfTransform.scala b/src/main/scala/scala/async/AnfTransform.scala
index 6b17b94..2fa96c9 100644
--- a/src/main/scala/scala/async/AnfTransform.scala
+++ b/src/main/scala/scala/async/AnfTransform.scala
@@ -194,13 +194,14 @@ private[async] final case class AnfTransform[C <: Context](c: C) {
           // this has already been checked.
           val isByName: (Int) => Boolean = utils.isByName(fun)
           val funStats :+ simpleFun = inline.transformToList(fun)
+          def isAwaitRef(name: Name) = name.toString.startsWith(utils.name.await + "$")
           val argLists: List[List[Tree]] = args.zipWithIndex map {
-            case (arg, i) if isByName(i) => List(arg)
-            case (arg, i)                => inline.transformToList(arg) match {
+            case (arg, i) if isByName(i) || isSafeToInline(arg) => List(arg)
+            case (arg@Ident(name), _) if isAwaitRef(name)       => List(arg) // not typed, so it eludes the check in `isSafeToInline`
+            case (arg, i)                                       => inline.transformToList(arg) match {
               case stats :+ expr =>
-                val valDef = defineVal(s"arg$i", expr, arg.pos)
+                val valDef = defineVal(name.arg(i), expr, arg.pos)
                 stats ::: List(valDef, Ident(valDef.name))
-              case xs            => xs
             }
           }
           val allArgStats = argLists flatMap (_.init)
diff --git a/src/main/scala/scala/async/TransformUtils.scala b/src/main/scala/scala/async/TransformUtils.scala
index 7571f88..23f39d2 100644
--- a/src/main/scala/scala/async/TransformUtils.scala
+++ b/src/main/scala/scala/async/TransformUtils.scala
@@ -30,6 +30,7 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
     val ifRes         = "ifres"
     val await         = "await"
     val bindSuffix    = "$bind"
+    def arg(i: Int)   = "arg" + i
 
     def fresh(name: TermName): TermName = newTermName(fresh(name.toString))
 
@@ -244,4 +245,12 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
     }
   }
 
+  def isSafeToInline(tree: Tree) = {
+    val symtab = c.universe.asInstanceOf[scala.reflect.internal.SymbolTable]
+    object treeInfo extends {
+      val global: symtab.type = symtab
+    } with reflect.internal.TreeInfo
+    val castTree = tree.asInstanceOf[symtab.Tree]
+    treeInfo.isExprSafeToInline(castTree)
+  }
 }
diff --git a/src/test/scala/scala/async/TreeInterrogation.scala b/src/test/scala/scala/async/TreeInterrogation.scala
index ecb1bca..b22faa9 100644
--- a/src/test/scala/scala/async/TreeInterrogation.scala
+++ b/src/test/scala/scala/async/TreeInterrogation.scala
@@ -70,17 +70,15 @@ object TreeInterrogation extends App {
     val cm = reflect.runtime.currentMirror
     val tb = mkToolbox("-cp target/scala-2.10/classes -Xprint:all")
     val tree = tb.parse(
-      """
-        | import scala.async.Async.{async, await}
-        | import scala.concurrent.{future, ExecutionContext, Await}
-        | import ExecutionContext.Implicits._
-        | import scala.concurrent.duration._
-        |
-        | try {
-        |   val f = async { throw new Exception("problem") }
-        |   Await.result(f, 1.second)
-        | } catch {
-        |   case ex: Exception if ex.getMessage == "problem" => // okay
+      """ import scala.async.AsyncId.{async, await}
+        | def foo(a: Int, b: Int) = (a, b)
+        | val result = async {
+        |   var i = 0
+        |   def next() = {
+        |     i += 1;
+        |     i
+        |   }
+        |   foo(next(), await(next()))
         | }
         | ()
         | """.stripMargin)