Skip to content

Commit 5a0b191

Browse files
retronymretronym
committed
Allow await in applications with multiple argument lists
Before, we levied an implementation restriction to prevent this. As it turned out, that needlessly prevented use of `await` in the receiver of a multi-param-list application. This commit lifts the restriction altogether, and treats such applications holistically, being careful to preserve the left-to-right evaluation order of arguments in the translated code. - use `TreeInfo.Applied` and `Type#paramss` from `reflect.internal` to get the info we need - use the parameter name for the lifted argument val, rather than `argN` - encapsulate handling of by-name-ness and parameter names in `mapArgumentss` - test for evaluation order preservation
1 parent 74beb1b commit 5a0b191

File tree

4 files changed

+132
-70
lines changed

4 files changed

+132
-70
lines changed

‎src/main/scala/scala/async/AnfTransform.scala‎

Lines changed: 12 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -187,31 +187,31 @@ private[async] final case class AnfTransform[C <: Context](c: C) {
187187

188188
private[AnfTransform] def transformToList(tree: Tree): List[Tree] = trace("anf", tree) {
189189
def containsAwait = tree exists isAwait
190+
190191
tree match {
191192
case Select(qual, sel) if containsAwait =>
192193
val stats :+ expr = inline.transformToList(qual)
193194
stats :+ attachCopy(tree)(Select(expr, sel).setSymbol(tree.symbol))
194195

195-
case Apply(fun, args) if containsAwait =>
196-
checkForAwaitInNonPrimaryParamSection(fun, args)
197-
196+
case utils.Applied(fun, targs, argss @ (args :: rest)) if containsAwait =>
198197
// we an assume that no await call appears in a by-name argument position,
199198
// this has already been checked.
200-
val isByName: (Int) => Boolean = utils.isByName(fun)
201199
val funStats :+ simpleFun = inline.transformToList(fun)
202200
def isAwaitRef(name: Name) = name.toString.startsWith(utils.name.await + "$")
203-
val (argStats, argExprs): (List[List[Tree]], List[Tree]) =
204-
mapArguments[List[Tree]](args) {
205-
case (arg, i) if isByName(i) || isSafeToInline(arg) => (Nil, arg)
206-
case (arg@Ident(name), _) if isAwaitRef(name) => (Nil, arg) // not typed, so it eludes the check in `isSafeToInline`
207-
case (arg, i) =>
208-
inline.transformToList(arg) match {
201+
val (argStatss, argExprss): (List[List[List[Tree]]], List[List[Tree]]) =
202+
mapArgumentss[List[Tree]](fun, argss) {
203+
case arg if arg.isByName || isSafeToInline(arg.expr) => (Nil, arg.expr)
204+
case Arg(arg@Ident(name), _, _) if isAwaitRef(name) => (Nil, arg) // not typed, so it eludes the check in `isSafeToInline`
205+
case arg =>
206+
inline.transformToList(arg.expr) match {
209207
case stats :+ expr =>
210-
val valDef = defineVal(name.arg(i), expr, arg.pos)
208+
val valDef = defineVal(arg.argName, expr, arg.expr.pos)
211209
(stats :+ valDef, Ident(valDef.name))
212210
}
213211
}
214-
funStats ++ argStats.flatten :+ attachCopy(tree)(Apply(simpleFun, argExprs).setSymbol(tree.symbol))
212+
val core = if (targs.isEmpty) simpleFun else TypeApply(simpleFun, targs)
213+
val newApply = argExprss.foldLeft(core)(Apply(_, _).setSymbol(tree.symbol))
214+
funStats ++ argStatss.flatten.flatten :+ attachCopy(tree)(newApply)
215215
case Block(stats, expr) if containsAwait =>
216216
inline.transformToList(stats :+ expr)
217217

@@ -273,19 +273,4 @@ private[async] final case class AnfTransform[C <: Context](c: C) {
273273
}
274274
}
275275
}
276-
277-
def checkForAwaitInNonPrimaryParamSection(fun: Tree, args: List[Tree]) {
278-
// TODO treat the Apply(Apply(.., argsN), ...), args0) holistically, and rewrite
279-
// *all* argument lists in the correct order to preserve semantics.
280-
fun match {
281-
case Apply(fun1, _) =>
282-
fun1.tpe match {
283-
case MethodType(_, resultType: MethodType) if resultType =:= fun.tpe =>
284-
c.error(fun.pos, "implementation restriction: await may only be used in the first parameter list.")
285-
case _ =>
286-
}
287-
case _ =>
288-
}
289-
290-
}
291276
}

‎src/main/scala/scala/async/TransformUtils.scala‎

Lines changed: 55 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,6 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
3232
val await = "await"
3333
val bindSuffix = "$bind"
3434

35-
def arg(i: Int) = "arg" + i
36-
3735
def fresh(name: TermName): TermName = newTermName(fresh(name.toString))
3836

3937
def fresh(name: String): String = if (name.toString.contains("$")) name else c.fresh("" + name + "$")
@@ -102,11 +100,13 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
102100
case dd: DefDef => nestedMethod(dd)
103101
case fun: Function => function(fun)
104102
case m@Match(EmptyTree, _) => patMatFunction(m) // Pattern matching anonymous function under -Xoldpatmat of after `restorePatternMatchingFunctions`
105-
case Apply(fun, args) =>
103+
case Applied(fun, targs, argss @ (_ :: _)) =>
106104
val isInByName = isByName(fun)
107-
for ((arg, index) <- args.zipWithIndex) {
108-
if (!isInByName(index)) traverse(arg)
109-
else byNameArgument(arg)
105+
for ((args, i) <- argss.zipWithIndex) {
106+
for ((arg, j) <- args.zipWithIndex) {
107+
if (!isInByName(i, j)) traverse(arg)
108+
else byNameArgument(arg)
109+
}
110110
}
111111
traverse(fun)
112112
case _ => super.traverse(tree)
@@ -122,13 +122,31 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
122122
Set(Boolean_&&, Boolean_||)
123123
}
124124

125-
def isByName(fun: Tree): (Int => Boolean) = {
126-
if (Boolean_ShortCircuits contains fun.symbol) i => true
127-
else fun.tpe match {
128-
case MethodType(params, _) =>
129-
val isByNameParams = params.map(_.asTerm.isByNameParam)
130-
(i: Int) => isByNameParams.applyOrElse(i, (_: Int) => false)
131-
case _ => Map()
125+
def isByName(fun: Tree): ((Int, Int) => Boolean) = {
126+
if (Boolean_ShortCircuits contains fun.symbol) (i, j) => true
127+
else {
128+
val symtab = c.universe.asInstanceOf[reflect.internal.SymbolTable]
129+
val paramss = fun.tpe.asInstanceOf[symtab.Type].paramss
130+
val byNamess = paramss.map(_.map(_.isByNameParam))
131+
(i, j) => util.Try(byNamess(i)(j)).getOrElse(false)
132+
}
133+
}
134+
def argName(fun: Tree): ((Int, Int) => String) = {
135+
val symtab = c.universe.asInstanceOf[reflect.internal.SymbolTable]
136+
val paramss = fun.tpe.asInstanceOf[symtab.Type].paramss
137+
val namess = paramss.map(_.map(_.name.toString))
138+
(i, j) => util.Try(namess(i)(j)).getOrElse(s"arg_${i}_${j}")
139+
}
140+
141+
object Applied {
142+
val symtab = c.universe.asInstanceOf[scala.reflect.internal.SymbolTable]
143+
object treeInfo extends {
144+
val global: symtab.type = symtab
145+
} with reflect.internal.TreeInfo
146+
147+
def unapply(tree: Tree): Some[(Tree, List[Tree], List[List[Tree]])] = {
148+
val treeInfo.Applied(core, targs, argss) = tree.asInstanceOf[symtab.Tree]
149+
Some((core.asInstanceOf[Tree], targs.asInstanceOf[List[Tree]], argss.asInstanceOf[List[List[Tree]]]))
132150
}
133151
}
134152

@@ -302,7 +320,6 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
302320
}
303321
}
304322

305-
306323
def isSafeToInline(tree: Tree) = {
307324
val symtab = c.universe.asInstanceOf[scala.reflect.internal.SymbolTable]
308325
object treeInfo extends {
@@ -322,7 +339,7 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
322339
* @param f A function from argument (with '_*' unwrapped) and argument index to argument.
323340
* @tparam A The type of the auxillary result
324341
*/
325-
def mapArguments[A](args: List[Tree])(f: (Tree, Int) => (A, Tree)): (List[A], List[Tree]) = {
342+
private def mapArguments[A](args: List[Tree])(f: (Tree, Int) => (A, Tree)): (List[A], List[Tree]) = {
326343
args match {
327344
case args :+ Typed(tree, Ident(tpnme.WILDCARD_STAR)) =>
328345
val (a, argExprs :+ lastArgExpr) = (args :+ tree).zipWithIndex.map(f.tupled).unzip
@@ -332,4 +349,27 @@ private[async] final case class TransformUtils[C <: Context](c: C) {
332349
args.zipWithIndex.map(f.tupled).unzip
333350
}
334351
}
352+
353+
case class Arg(expr: Tree, isByName: Boolean, argName: String)
354+
355+
/**
356+
* Transform a list of argument lists, producing the transformed lists, and lists of auxillary
357+
* results.
358+
*
359+
* The function `f` need not concern itself with varargs arguments e.g (`xs : _*`). It will
360+
* receive `xs`, and it's result will be re-wrapped as `f(xs) : _*`.
361+
*
362+
* @param fun The function being applied
363+
* @param argss The argument lists
364+
* @return (auxillary results, mapped argument trees)
365+
*/
366+
def mapArgumentss[A](fun: Tree, argss: List[List[Tree]])(f: Arg => (A, Tree)): (List[List[A]], List[List[Tree]]) = {
367+
val isByNamess: (Int, Int) => Boolean = isByName(fun)
368+
val argNamess: (Int, Int) => String = argName(fun)
369+
argss.zipWithIndex.map { case (args, i) =>
370+
mapArguments[A](args) {
371+
(tree, j) => f(Arg(tree, isByNamess(i, j), argNamess(i, j)))
372+
}
373+
}.unzip
374+
}
335375
}

‎src/test/scala/scala/async/TreeInterrogation.scala‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -71,17 +71,14 @@ object TreeInterrogation extends App {
7171
val tb = mkToolbox("-cp target/scala-2.10/classes -Xprint:flatten")
7272
import scala.async.Async._
7373
val tree = tb.parse(
74-
""" import scala.async.AsyncId._
75-
| async {
76-
| val x = 1
77-
| val opt = Some("")
78-
| await(0)
79-
| val o @ Some(y) = opt
80-
|
81-
| {
82-
| val o @ Some(y) = Some(".")
83-
| }
74+
""" import _root_.scala.async.AsyncId.{async, await}
75+
| def foo[T](a0: Int)(b0: Int*) = s"a0 = $a0, b0 = ${b0.head}"
76+
| val res = async {
77+
| var i = 0
78+
| def get = async {i += 1; i}
79+
| foo[Int](await(get))(await(get) :: Nil : _*)
8480
| }
81+
| res
8582
| """.stripMargin)
8683
println(tree)
8784
val tree1 = tb.typeCheck(tree.duplicate)

‎src/test/scala/scala/async/run/anf/AnfTransformSpec.scala‎

Lines changed: 58 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -232,28 +232,68 @@ class AnfTransformSpec {
232232
}
233233

234234
@Test
235-
def awaitNotAllowedInNonPrimaryParamSection1() {
236-
expectError("implementation restriction: await may only be used in the first parameter list.") {
237-
"""
238-
| import _root_.scala.async.AsyncId.{async, await}
239-
| def foo(primary: Any)(i: Int) = i
240-
| async {
241-
| foo(???)(await(0))
242-
| }
243-
""".stripMargin
235+
def awaitInNonPrimaryParamSection1() {
236+
import _root_.scala.async.AsyncId.{async, await}
237+
def foo(a0: Int)(b0: Int) = s"a0 = $a0, b0 = $b0"
238+
val res = async {
239+
var i = 0
240+
def get = {i += 1; i}
241+
foo(get)(get)
242+
}
243+
res mustBe "a0 = 1, b0 = 2"
244+
}
245+
246+
@Test
247+
def awaitInNonPrimaryParamSection2() {
248+
import _root_.scala.async.AsyncId.{async, await}
249+
def foo[T](a0: Int)(b0: Int*) = s"a0 = $a0, b0 = ${b0.head}"
250+
val res = async {
251+
var i = 0
252+
def get = async {i += 1; i}
253+
foo[Int](await(get))(await(get) :: await(async(Nil)) : _*)
254+
}
255+
res mustBe "a0 = 1, b0 = 2"
256+
}
257+
258+
@Test
259+
def awaitInNonPrimaryParamSectionWithLazy1() {
260+
import _root_.scala.async.AsyncId.{async, await}
261+
def foo[T](a: => Int)(b: Int) = b
262+
val res = async {
263+
def get = async {0}
264+
foo[Int](???)(await(get))
244265
}
266+
res mustBe 0
245267
}
246268

247269
@Test
248-
def awaitNotAllowedInNonPrimaryParamSection2() {
249-
expectError("implementation restriction: await may only be used in the first parameter list.") {
250-
"""
251-
| import _root_.scala.async.AsyncId.{async, await}
252-
| def foo[T](primary: Any)(i: Int) = i
253-
| async {
254-
| foo[Int](???)(await(0))
255-
| }
256-
""".stripMargin
270+
def awaitInNonPrimaryParamSectionWithLazy2() {
271+
import _root_.scala.async.AsyncId.{async, await}
272+
def foo[T](a: Int)(b: => Int) = a
273+
val res = async {
274+
def get = async {0}
275+
foo[Int](await(get))(???)
276+
}
277+
res mustBe 0
278+
}
279+
280+
@Test
281+
def awaitWithLazy() {
282+
import _root_.scala.async.AsyncId.{async, await}
283+
def foo[T](a: Int, b: => Int) = a
284+
val res = async {
285+
def get = async {0}
286+
foo[Int](await(get), ???)
287+
}
288+
res mustBe 0
289+
}
290+
291+
@Test
292+
def awaitOkInReciever() {
293+
import scala.async.AsyncId.{async, await}
294+
class Foo { def bar(a: Int)(b: Int) = a + b }
295+
async {
296+
await(async(new Foo)).bar(1)(2)
257297
}
258298
}
259299

0 commit comments

Comments
 (0)