Pre-SIP: improve for-comprehensions functionality

#1

Hello,

I’m the author of better-monadic-for compiler plugin (bm4 for short). It’s been brought to my attention that dotty would only allow pre-typer plugins in nightly builds, which means there cannot be a version of bm4 for production Scala 3.

Originally, the goal was to allow writing monadic code with destructuring. I.e. the code like

for {
  (a, b) <- computeThings
} yield a + b

could fail to compile depending on whether a value returned by computeThings has a withFilter method, and a common workaround is to use a temporary variable and destructure with a subsequent assignment:

for {
  temp <- computeThings
  (a, b) = temp
} yield a + b

Luckily, #2578 was fixed in dotty, so for-comprehensions don’t filter until asked to with a case keyword, greatly reducing the need for bm4 in the first place.

Two other features of the plugin that seems to be quite popular are ability to introduce values into implicit scope in the middle of a for-comprehension and the removal of map call as the last action. These are oriented towards users of effect monads (cats-effect IO/monix Task/ZIO) and related utilities (e.g. Resource and fs2.Stream).

These are things I propose Scala 3 includes directly. I’ll list my thoughs on that below.

The plugin also has removal of extraneous generation of TupleN from value bindings (reported as #2573), but unlike the two above it doesn’t affect how you write code with for-comprehensions, only its runtime characteristics, so it’s of lower priority.

Eliminating map from for-comprehension

Motivation

Scala for-comprehensions desugar to nested flatMap+map calls, for instance:

for {
  a <- f1
  b <- f2(a)
  c <- f3(a, b)
} yield c

becomes:

f1.flatMap { a =>
  f2(a).flatMap { b =>
    f3(a, b).map { c => c }
  }    
}

Notice that there is .map { c => c } call at the end. With any valid Functor from cats or scalaz it is guaranteed to not change anything in the result, and this is true without these libraries for most types from standard library (Option/Try/Either, immutable collections and Future). So, all it does here is allocate more objects for no good reason.

With effect types like IO, it gets worse, as this property makes for-comprehension unfit for describing infinite loops. I.e. code like such:

def loop: IO[Unit] =
  for {
    _ <- IO(println("Doing periodic things"))
    _ <- IO.sleep(1.second)
    _ <- loop
  } yield ()

will slowly but surely generate garbage elements with the .map((_: Unit) => ()) construct that desugaring generates, eventually crashing the application with OOME. better-monadic-for can detect such constructs and eliminate them, but without it the only option user has is to rewrite the definition entirely, not using for-comprehension.

Proposals

I see two possible solutions:

  1. Do this elimination in scalac
  2. Extend for-comprehension syntax to support use cases where you do NOT want any map.

Option #1 would provide quite an intuitive behavior, but would make desugaring itself a bit more complex. To handle cases such as loop above, it’s also required that it happens after typer (this is how bm4 plugin does it). It is also a double-edged sword for migration. Nothing would change for people who use bm4 already, and for immutable types it would bring performance improvements, but will break code that relies on that .map. I haven’t seen any code like this in the wild, but it could exist, and it will silently break if migrated as-is.

Option #2 would require people to know about it’s existence and also we would have more options for something we already have. For example, we would allow loops to be written as:

def loop: IO[Unit] =
  for {
    _ <- IO(println("Doing periodic things"))
    _ <- IO.sleep(1.second)
  } yield <- loop

that desugar into

IO(println("Doing periodic things")).flatMap { _ =>
  IO.sleep(1.second).flatMap { _ =>
    loop
  }
}

This doesn’t conflict with Scala 2.x, as <- is a reserved character sequence and could not be used as an identifier. All existing code (that is not affected by #2578 fix already) would retain current behavior.

Declaring implicits in the middle of for-comprehension

Implicits (givens? I’m not sure what’s the current name) are the defining feature of Scala, but limited syntax of for-comprehension makes it impossible to define them within one. Every binding is effectively a val, with no other options. defs can be simulated with function values, vars - with AtomicReference or something alike, but nothing in Scala of today would let you define an equivalent of implicit val. For effect types in particular, there are things that have to be created in IO/Resource context, such as clients, thread pools and tagless final algebras, but using them in 2.x requires either passing implicit parameters explicitly:

def createIOPool: IO[ExecutionContext]
class HttpClient[F[_]]
object HttpClient {
  def create[F[_]: Async](implicit pool: ExecutionContext): F[HttpClient]
}
def doPolling[F[_]: Monad: HttpClient](url: String): F[Unit] = ???


for {
  blockingPool <- createIOPool
  httpClient   <- HttpClient.create[IO](implicitly, blockingPool)
  _            <- doPolling("http://example.com")(implicitly, httpClient)
  _            <- doPolling("http://scala-lang.org")(implicitly, httpClient)
} yield ()

or ditching for-comprehensions so that you can use a lambda with implicit:

createIOPool.flatMap { implicit blockingPool =>
  HttpClient.create[IO].flatMap { implicit client =>
    doPolling("http://example.com").flatMap { _ =>
      doPolling("http://scala-lang.org")
    }
  }
}

better-monadic-for adds implicit0 construct which can be used as a pattern inside for-comprehension directly:

for {
  implicit0(blockingPool: ExecutionContext) <- createIOPool
  implicit0(httpClient: HttpClient[IO])     <- HttpClient.create[IO]

  _ <- doPolling("http://example.com")
  _ <- doPolling("http://scala-lang.org")
} yield ()

I’ve found this to be very helpful occasionally. Particularly at the application entry point, where everything is initialized and dependencies are wired together, rewriting to version without for-comprehensions is extremely tedious and creates quite a bit of nesting. It’s also supported by Intellij IDEA since 2019.1.

Proposal

Allow alias given definitions at the left hand side of a generator or a = binding:

for {
  given blockingPool as ExecutionContext <- createIOPool
  given as HttpClient[IO]                <- HttpClient.create[IO]

  _ <- doPolling("http://example.com")
  _ <- doPolling("http://scala-lang.org")
} yield ()

This can actually compile a specific case (given as Type where Type has no type parameters) with 2.12 syntax if somebody were to write e.g.

object as {
  def unappply(a: (Any, Any)): Option[(Any, Any)] = Some(a)
}

but due to given being a keyword, it would require rewrite to update already.


P.S. While we’re talking about for-comprehensions, maybe we should also address scala bug #907 too? I can’t do it in the plugin without ridiculous parser hacking.

81 Likes
#2

This is one of compiler plugins that almost always end up in my build. It would be awesome if at least final .map removal got into Scala 3 and it would be epic if we could get implicit declaration too.

7 Likes
#3

This is one of the things that makes programming with Monads hard. Remove map from “for” comprehensions or maybe create “do” comprehensions for foldMap only that resolve all Monad implicits correctly.

#4

kind projector, macro-paradise and bm4 are default plugins in every project I start.

I use implicit0(...) <- right now as very powerful DI - replacement.
It removes the need for large ReaderT-like monad contexts, or any DI implementation.
Simple, pure, predictable, ordered, effectful initialization of application, all with a single syntactic extension, that will become so much prettier when the anonymous form given as Foo[F] <- Foo[I, F] is a thing.

10 Likes
#5

I haven’t used bm4 yet, so I have less of a horse in this race, but I gotta say, this syntax:

for {
  given blockingPool as ExecutionContext <- createIOPool
  given as HttpClient[IO] <- HttpClient.create[IO]
  ...

makes intuitive sense to me in the new world, and seems like an appropriate and consistent enhancement. I like it…

2 Likes
#6

Ah I see you beat me to it @oleg-py :smiley:

I have been looking at this as time allowed over the last couple of weeks.

Here’s a quick implementation of the elimination of the final map call in scalac.

Here’s a quick change which implements a simplification of value bindings in scalac.

However, while I was implementing these I ran into a couple of things which are worth discussing here.

Eliminating map from for-comprehension

This is a more complicated change than it might seem in terms of downstream breakage.

Consider the following code:

val xs = for { x <- List(1, 2, 3); if x > 3 } yield x

xs.length

Once the map call has been eliminated, we have a problem. withFilter returns a FilterMonadic in Scala <= 2.12 and a WithFilter in Scala 2.13. The terminal map call is usually used to convert the WithFilter back into the target collection type, so I suspect this change breaks lots of existing code.

We need to decide what to do about that - disable the optimization when there is a filter? Change the filter to a strict filter call when it is the final method call in the for comprehension? What if the user code has not implemented filter for that type?

The other thing to consider here is whether users will have written code which follows the monad right identity law, i.e. such map(x => x) actually is a no-op. I’m sure that there is user code out there which would be affected by this change.

I suspect we would need to run a community build to estimate the amount of breakage.

Simplifying desugaring of value bindings

As @oleg-py says this is probably less of a priority but nevertheless there is a problem with this change too. The idea of this change is to alter the desugaring of code like the following

for {
  x <- List(1, 2, 3)
  y = x + 1
  z <- List(4, 5, 6)
} yield (x + y + z)

from this:

List(1, 2, 3).map { x =>
  val y = x + 1
  (x, y)
}.flatMap {
  case (x, y) => 
  List(4, 5, 6).map { z =>
    x + y + z
  }
}

to this:

List(1, 2, 3).flatMap { x =>
  val y = x + 1
  List(4, 5, 6).map { z =>
    x + y + z
  }
}

In other words, we are trying to change value bindings from a tuple that gets threaded through each call to map or flatMap to a simple val that is in scope for nested map and flatMap calls.

However, this causes problems for code like the following:

for {
  x <- List(1, 2, 3)
  y = x + 1
  if y > 0
  z <- List(4, 5, 6)
} yield (x + y + z)

The new desugaring causes problems like this:

List(1, 2, 3).withFilter( /* what is y here? how do we write the filter condition? */ ).flatMap { x =>
  val y = x + 1
  List(4, 5, 6).map { z =>
    x + y + z
  }
}

We would need to decide what to do with if after value bindings if we proceed with that change.

Personally I would suggest that if should only be allowed immediately after a <- generator, and I think @odersky has mentioned this solution previously, but once again this seems like it would break a lot of existing code.

Alternatives

  • Should we change the for comprehension at all? Perhaps we should investigate whether we can repurpose do as it seems to see little use (my perception about this might be totally wrong)?
  • Should we investigate whether we could add a completely different kind of comprehension syntax? There are lots of interesting ideas in the FP space now, such as F# computation expressions and OCaml’s new comprehension syntax.
4 Likes
#7

What bm4 does is optimize it after typer, and only if the type would not change. So in filter calls it would not be triggered. The code is here.

But yes, there is a potential for breakage. This is why I think new syntax is a better solution - you would not be able to write that if after yield <- outside for-comprehension.

I’m quite in favor of adding unit/pure/guard on the instance, so the desugaring would be something like:

val lifted$1 = List(1, 2, 3)
lifted$1.flatMap {
  val y = x + 1
  // Typeclasses could work nicely here, but the concept is very simple - this, but with cached instances:
  // def guard: List[Unit] =
  //   if (b) List(()) else List.empty[Unit] 
  lifted$1.companion.guard(y > 0).flatMap { _ =>
      List(4, 5, 6).map { z =>
        x + y + z
      }
  }
}

bm4 does a bit of lookahead and doesn’t alter chains that are followed by withFilter to avoid breakage. Not a simple rewrite by any means, however, and changing the strategy is probably the way to go there.

I think it’s of lower priority because the end user shouldn’t see the difference with that change except in allocations, not that it’s less desirable.


As for other comprehensions, I don’t think there’s something that would address the need of community significantly better. Besides, for is already getting an upgrade with case pattern syntax.

3 Likes
#8

I think there is a case for investigating whether there is a more flexible syntax that we could use that might permit applicative desugaring or monoid comprehensions, but I agree that it seems out of scope for this discussion.

1 Like
#9

+1 for embracing fully F#-like computation expressions

But the problem with the current for comprehension and redundant map could be solved by a relatively small change:

What we need is a for comprehension without the yield part. So instead of

for {
  a <- f1
  b <- f2(a)
  c <- f3(a, b)
} yield c

it would be just

for {
  a <- f1
  b <- f2(a)
  f3(a, b)
}

That’s how Haskell does it, btw. Allowing that seem easily doable, doesn’t it? Or are there problems with it?

#10

The latter already desugars to foreach in scala

#11

oh, i must admit i didn’t know abot the foreach desugaring.

can we change that? so that it does what we want it to do?

#12

That would be nice but it seems unlikely - for without yield is ubiquitous in Scala code especially for new users migrating from Java and other imperative languages.

I think @oleg-py’s approach of using a different keyword is more likely to be accepted as it can’t change the behaviour of existing code.

2 Likes
#13

I think there is a less obtrusive way of mitigating the performance problems caused by the ending map.

The idea is to have the compiler detect when the mapped function is an identity, and use some id instance in that case, as in: {...}.map(id) where:

private val _id: Any => Any = x => x
def id[A]: A => A = _id.asInstanceOf[A => A]

Then, each implementation can special-case that specific instance, avoiding allocations. For example, List could have:

abstract class List[+A] {
  ...
  def map[B](f: A => B): List[B] =
    if (f == id) this.asInstanceOf[List[B]]
    else ... // normal implementation
  ...
}

(We could also probably use some GADT tricks to make the compiler see that when f == id then A =:= B so we can avoid the cast.)

1 Like
#14

That seems great within the Scala collections library, but for everyone else it is no better than the status quo, i.e. each library maintainer will have to add their own special-case code to avoid these allocations. The idea of having a single id instance for the compiler to use is pretty compelling for enabling that optimisation though!

1 Like
#15

How is it no better than the status quo, when it allows anyone to make their map definitions faster for better for comprehension performance? Currently, there is no way of achieving that at all.

Yes, it is not ideal. But no one has come up with an ideal solution so far, and I don’t believe there is one. You’ll either have to add syntax to the language or break older code. This is just the least obtrusive way of achieving the stated goal.

#16

If you want to do that now then you can encourage your users to use better-monadic-for, which does the optimization without special code in every library.
I’m not saying it’s a bad idea and I’m certain that library authors would use it.

1 Like
#17

One change to this latter proposal, which seems interesting to me, is to make:

def id[A]: A => A = implicitly[A =:= A]

Which is okay currently because A =:= A extends A <:< A which extends A => A.

then:

def map[B](fn: A => B): List[B] =
  fn match {
    case ident: A <:< B => ident.substituteCo[List](this)
    case notId =>
       ...
  }

This can work without any casts, and since substituteCo is ultimately identity, the JVM should inline this away. This trick works for any covariant type. If you have an invariant type, you can do the same except with case ident: A =:= B =>

So, this proposal gets all the benefits of _id: Any => Any except without the need to sprinkle casts throughout our code.

3 Likes
#18

As I mentioned, a GADT also does the trick — but thanks to Dotty’s recently improved GADT reasoning capabilities, we can write the nicer:

sealed abstract class <:<[-A,+B] extends (A => B)
case class Refl[A] extends (A <:< A) {
  def apply(a: A): A = a
}
def id[A]: A => A = new Refl

def map[A,B](ls: List[A]): (A => B) => List[B] = {
  case _: Refl[a] => ls
  case f => ls.map(f)
}
1 Like
#19

Sure. But note that, as was mentioned above, better-monadic-for does not handle chains ended by withFilter, so it’s not a general solution either. The map(id) trick would still save a list/whatever construction in this case.

#20

I’m more in favor of new syntax over magic optimization.

  • People writing map would have to know that there is this optimization option
  • People maintaining scalac would have to check function bodies to synthesize in a proper =:= or <:< instead. It’s actually a bit tricky - has to be done post-typer to account for yield () becoming _ => (), which might or might not be an identity and subtyping, and that means that map isn’t .map(a => b) anymore, it could be desugared to new FunctorOps[F](fa).map[A, B](a => b)(the[Functor[F]])
  • People using for-comprehensions would have to know if their container has it implemented AND when compiler triggers it. Library authors which abstract over F, effect type or just container, actually cannot possibly know the former - I expect it to go mainstream, but there just might be edge cases spanning across multiple libs causing memory leaks in production.

OTOH, you get the syntax which directly expresses the intent (flatMap into that thing directly, no actions afterwards) and if it compiles, it will work.

I never actually bothered to support optimization for withFilter in better-monadic-for, both due to time constraints and the fact that withFilter on types that support it is already lazy, so ending with .withFilter(f).map(a => a) or just .filter(f) would make very small difference in allocs.

3 Likes