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

Author: KacperFKorban

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

@@ -9,7 +9,7 @@ import Phases.Phase
 import transform.*
 import dotty.tools.backend
 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`.
@@ -91,7 +91,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

@@ -52,6 +52,11 @@ object desugar {
    */
   val ContextBoundParam: 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
@@ -1872,9 +1877,9 @@ object desugar {
      *    (Where empty for-comprehensions are excluded by the parser)
      *
      *   If the aliases are not followed by a guard, otherwise an error.
-     * 
+     *
      * With betterFors disabled, the translation is as follows:
-     * 
+     *
      * 1.
      *
      *    for (P <- G) E   ==>   G.foreach (P => E)
@@ -2044,14 +2049,22 @@ object desugar {
             if 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))
+            else
+              val aply = Apply(rhsSelect(gen, mapName), makeLambda(gen, body))
+              aply.putAttachment(TrailingForMap, ())
+              aply
           case (gen: GenFrom) :: rest
           if rest.dropWhile(_.isInstanceOf[GenAlias]).headOption.forall(e => e.isInstanceOf[GenFrom]) =>
             val cont = makeFor(mapName, flatMapName, rest, body)
             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, ())
+            else
+              aply
+            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)
+
+}