Crude implementation of removing trailing unit-literal maps from for-comprehensions

Author: KacperFKorban

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -8,7 +8,7 @@ import parsing.Parser
 import Phases.Phase
 import transform.*
 import backend.jvm.{CollectSuperCalls, GenBCode}
-import localopt.StringInterpolatorOpt
+import localopt.{StringInterpolatorOpt, DropForMap}
 
 /** The central class of the dotc compiler. The job of a compiler is to create
  *  runs, which process given `phases` in a given `rootContext`.
@@ -90,7 +90,8 @@ class Compiler {
          new ExplicitOuter,          // Add accessors to outer classes from nested ones.
          new ExplicitSelf,           // Make references to non-trivial self types explicit as casts
          new StringInterpolatorOpt,  // Optimizes raw and s and f string interpolators by rewriting them to string concatenations or formats
-         new DropBreaks) ::          // Optimize local Break throws by rewriting them
+         new DropBreaks,             // Optimize local Break throws by rewriting them
+         new DropForMap) ::          // Drop unused trailing map calls in for comprehensions
     List(new PruneErasedDefs,        // Drop erased definitions from scopes and simplify erased expressions
          new UninitializedDefs,      // Replaces `compiletime.uninitialized` by `_`
          new InlinePatterns,         // Remove placeholders of inlined patterns

compiler/src/dotty/tools/dotc/ast/Desugar.scala

@@ -64,6 +64,11 @@ object desugar {
    */
   val PolyFunctionApply: Property.Key[Unit] = Property.StickyKey()
 
+  /** An attachment key to indicate that an Apply is created as a last `map`
+   *  scall in a for-comprehension.
+   */
+  val TrailingForMap: Property.Key[Unit] = Property.StickyKey()
+
   /** What static check should be applied to a Match? */
   enum MatchCheck {
     case None, Exhaustive, IrrefutablePatDef, IrrefutableGenFrom
@@ -2149,11 +2154,13 @@ object desugar {
       enums match {
         case Nil if betterForsEnabled => body
         case (gen: GenFrom) :: Nil =>
+          val aply = Apply(rhsSelect(gen, mapName), makeLambda(gen, body))
           if betterForsEnabled
-            && gen.checkMode != GenCheckMode.Filtered // results of withFilter have the wrong type
-            && deepEquals(gen.pat, body)
-          then gen.expr  // avoid a redundant map with identity
-          else Apply(rhsSelect(gen, mapName), makeLambda(gen, body))
+          && gen.checkMode != GenCheckMode.Filtered // results of withFilter have the wrong type
+          // && deepEquals(gen.pat, body)
+          then
+            aply.putAttachment(TrailingForMap, ())
+          aply
         case (gen: GenFrom) :: (rest @ (GenFrom(_, _, _) :: _)) =>
           val cont = makeFor(mapName, flatMapName, rest, body)
           Apply(rhsSelect(gen, flatMapName), makeLambda(gen, cont))
@@ -2164,7 +2171,10 @@ object desugar {
           val selectName =
             if rest.exists(_.isInstanceOf[GenFrom]) then flatMapName
             else mapName
-          Apply(rhsSelect(gen, selectName), makeLambda(gen, cont))
+          val aply = Apply(rhsSelect(gen, selectName), makeLambda(gen, cont))
+          if selectName == mapName then
+            aply.pushAttachment(TrailingForMap, ())
+          aply
         case (gen: GenFrom) :: (rest @ GenAlias(_, _) :: _) =>
           val (valeqs, rest1) = rest.span(_.isInstanceOf[GenAlias])
           val pats = valeqs map { case GenAlias(pat, _) => pat }

compiler/src/dotty/tools/dotc/transform/localopt/DropForMap.scala

@@ -0,0 +1,56 @@
+package dotty.tools.dotc
+package transform.localopt
+
+import scala.language.unsafeNulls
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Decorators.*
+import dotty.tools.dotc.core.Constants.Constant
+import dotty.tools.dotc.core.Contexts.*
+import dotty.tools.dotc.core.StdNames.*
+import dotty.tools.dotc.core.Symbols.*
+import dotty.tools.dotc.core.Types.*
+import dotty.tools.dotc.transform.MegaPhase.MiniPhase
+import dotty.tools.dotc.typer.ConstFold
+import dotty.tools.dotc.ast.desugar
+import scala.util.chaining.*
+
+class DropForMap extends MiniPhase:
+  import tpd.*
+
+  override def phaseName: String = DropForMap.name
+
+  override def description: String = DropForMap.description
+
+  override def transformApply(tree: tpd.Apply)(using Context): tpd.Tree =
+    if !tree.hasAttachment(desugar.TrailingForMap) then tree.tap(_.removeAttachment(desugar.TrailingForMap))
+    else tree match
+      case Apply(MapCall(f), List(Lambda(List(param), body)))
+      if isEssentiallyUnitLiteral(param, body) && param.tpt.tpe.isRef(defn.UnitClass) =>
+        f
+      case _ =>
+        tree.tap(_.removeAttachment(desugar.TrailingForMap))
+
+  private object Lambda:
+    def unapply(tree: Tree)(using Context): Option[(List[ValDef], Tree)] =
+      tree match
+        case Block(List(defdef: DefDef), Closure(Nil, ref, _)) if ref.symbol == defdef.symbol && !defdef.paramss.exists(_.forall(_.isType)) =>
+          Some((defdef.termParamss.flatten, defdef.rhs))
+        case _ => None
+
+  private object MapCall:
+    def unapply(tree: Tree)(using Context): Option[Tree] = tree match
+      case Select(f, nme.map) => Some(f)
+      case Apply(fn, _) => unapply(fn)
+      case TypeApply(fn, _) => unapply(fn)
+      case _ => None
+
+  def isEssentiallyUnitLiteral(param: ValDef, tree: Tree)(using Context): Boolean = tree match
+    case Literal(Constant(())) => true
+    case Match(scrutinee, List(CaseDef(_, EmptyTree, body))) => isEssentiallyUnitLiteral(param, body)
+    case Block(Nil, expr) => isEssentiallyUnitLiteral(param, expr)
+    case _ => false
+
+object DropForMap:
+  val name: String = "dropForMap"
+  val description: String = "Drop unused trailing map calls in for comprehensions"

tests/run/map-unit-elim.check

@@ -0,0 +1 @@
+MySome(())

tests/run/map-unit-elim.scala

@@ -0,0 +1,34 @@
+import scala.language.experimental.betterFors
+
+class myOptionPackage(doOnMap: => Unit) {
+  sealed trait MyOption[+A] {
+    def map[B](f: A => B): MyOption[B] = this match {
+      case MySome(x) => {
+        doOnMap
+        MySome(f(x))
+      }
+      case MyNone => MyNone
+    }
+    def flatMap[B](f: A => MyOption[B]): MyOption[B] = this match {
+      case MySome(x) => f(x)
+      case MyNone => MyNone
+    }
+  }
+  case class MySome[A](x: A) extends MyOption[A]
+  case object MyNone extends MyOption[Nothing]
+}
+
+object Test extends App {
+
+  val myOption = new myOptionPackage(println("map called"))
+
+  import myOption.*
+
+  val z = for {
+    a <- MySome(1)
+    b <- MySome(())
+  } yield ()
+
+  println(z)
+
+}