5 files changed, 35 insertions, 23 deletions
diff --git a/src/main/scala/scala/async/AnfTransform.scala b/src/main/scala/scala/async/AnfTransform.scala
index 2fa96c9..afcf6bd 100644
--- a/src/main/scala/scala/async/AnfTransform.scala
+++ b/src/main/scala/scala/async/AnfTransform.scala
@@ -195,18 +195,18 @@ private[async] final case class AnfTransform[C <: Context](c: C) {
           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) || 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(name.arg(i), expr, arg.pos)
-                stats ::: List(valDef, Ident(valDef.name))
+          val (argStats, argExprs): (List[List[Tree]], List[Tree]) =
+            mapArguments[List[Tree]](args) {
+              case (arg, i) if isByName(i) || isSafeToInline(arg) => (Nil, arg)
+              case (arg@Ident(name), _) if isAwaitRef(name)       => (Nil, arg) // not typed, so it eludes the check in `isSafeToInline`
+              case (arg, i)                                       =>
+                inline.transformToList(arg) match {
+                  case stats :+ expr =>
+                    val valDef = defineVal(name.arg(i), expr, arg.pos)
+                    (stats :+ valDef, Ident(valDef.name))
+                }
             }
-          }
-          val allArgStats = argLists flatMap (_.init)
-          val simpleArgs = argLists map (_.last)
-          funStats ++ allArgStats :+ attachCopy(tree)(Apply(simpleFun, simpleArgs).setSymbol(tree.symbol))
+          funStats ++ argStats.flatten :+ attachCopy(tree)(Apply(simpleFun, argExprs).setSymbol(tree.symbol))
         case Block(stats, expr) if containsAwait =>
           inline.transformToList(stats :+ expr)
 
diff --git a/src/main/scala/scala/async/Async.scala b/src/main/scala/scala/async/Async.scala
index 2ff0f07..6ad1441 100644
--- a/src/main/scala/scala/async/Async.scala
+++ b/src/main/scala/scala/async/Async.scala
@@ -7,9 +7,6 @@ package scala.async
 import scala.language.experimental.macros
 import scala.reflect.macros.Context
 
-/*
- * @author Philipp Haller
- */
 object Async extends AsyncBase {
 
   import scala.concurrent.Future
@@ -66,12 +63,12 @@ abstract class AsyncBase {
   def asyncImpl[T: c.WeakTypeTag](c: Context)(body: c.Expr[T]): c.Expr[futureSystem.Fut[T]] = {
     import c.universe._
 
-    val anaylzer = AsyncAnalysis[c.type](c)
+    val analyzer = AsyncAnalysis[c.type](c)
     val utils = TransformUtils[c.type](c)
     import utils.{name, defn}
     import builder.futureSystemOps
 
-    anaylzer.reportUnsupportedAwaits(body.tree)
+    analyzer.reportUnsupportedAwaits(body.tree)
 
     // Transform to A-normal form:
     //  - no await calls in qualifiers or arguments,
@@ -88,7 +85,7 @@ abstract class AsyncBase {
     // states of our generated state machine, e.g. a value assigned before
     // an `await` and read afterwards.
     val renameMap: Map[Symbol, TermName] = {
-      anaylzer.defTreesUsedInSubsequentStates(anfTree).map {
+      analyzer.defTreesUsedInSubsequentStates(anfTree).map {
         vd =>
           (vd.symbol, name.fresh(vd.name.toTermName))
       }.toMap
diff --git a/src/main/scala/scala/async/AsyncUtils.scala b/src/main/scala/scala/async/AsyncUtils.scala
index 0a54d2e..1ade5f0 100644
--- a/src/main/scala/scala/async/AsyncUtils.scala
+++ b/src/main/scala/scala/async/AsyncUtils.scala
@@ -3,9 +3,6 @@
  */
 package scala.async
 
-/*
- * @author Philipp Haller
- */
 object AsyncUtils {
 
   private def enabled(level: String) = sys.props.getOrElse(s"scala.async.$level", "false").equalsIgnoreCase("true")
diff --git a/src/main/scala/scala/async/ExprBuilder.scala b/src/main/scala/scala/async/ExprBuilder.scala
index ef27672..180e7b9 100644
--- a/src/main/scala/scala/async/ExprBuilder.scala
+++ b/src/main/scala/scala/async/ExprBuilder.scala
@@ -8,9 +8,6 @@ import scala.collection.mutable.ListBuffer
 import collection.mutable
 import language.existentials
 
-/*
- * @author Philipp Haller
- */
 private[async] final case class ExprBuilder[C <: Context, FS <: FutureSystem](c: C, futureSystem: FS, origTree: C#Tree) {
   builder =>
 
diff --git a/src/main/scala/scala/async/TransformUtils.scala b/src/main/scala/scala/async/TransformUtils.scala
index 22c1b90..5a4984f 100644
--- a/src/main/scala/scala/async/TransformUtils.scala
+++ b/src/main/scala/scala/async/TransformUtils.scala
@@ -293,4 +293,25 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
     val castTree = tree.asInstanceOf[symtab.Tree]
     treeInfo.isExprSafeToInline(castTree)
   }
+
+  /** Map a list of arguments to:
+    * - A list of argument Trees
+    * - A list of auxillary results.
+    *
+    * The function unwraps and rewraps the `arg :_*` construct.
+    *
+    * @param args The original argument trees
+    * @param f  A function from argument (with '_*' unwrapped) and argument index to argument.
+    * @tparam A The type of the auxillary result
+    */
+  def mapArguments[A](args: List[Tree])(f: (Tree, Int) => (A, Tree)): (List[A], List[Tree]) = {
+    args match {
+      case args :+ Typed(tree, Ident(tpnme.WILDCARD_STAR)) =>
+        val (a, argExprs :+ lastArgExpr) = (args :+ tree).zipWithIndex.map(f.tupled).unzip
+        val exprs = argExprs :+ Typed(lastArgExpr, Ident(tpnme.WILDCARD_STAR)).setPos(lastArgExpr.pos)
+        (a, exprs)
+      case args                                            =>
+        args.zipWithIndex.map(f.tupled).unzip
+    }
+  }
 }