Drop inlined map

Author: som-snytt

compiler/src/dotty/tools/dotc/inlines/Inlines.scala

@@ -571,7 +571,7 @@ object Inlines:
 
       // Take care that only argument bindings go into `bindings`, since positions are
       // different for bindings from arguments and bindings from body.
-      val inlined = tpd.Inlined(call, bindings, expansion)
+      val inlined = tpd.Inlined(call, bindings, expansion.withAttachmentsFrom(call))
 
       if !hasOpaqueProxies then inlined
       else

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

@@ -20,7 +20,6 @@ import dotty.tools.dotc.ast.desugar
  *  The latter condition is checked in [[Desugar.scala#makeFor]]
  */
 class DropForMap extends MiniPhase:
-  import DropForMap.*
 
   override def phaseName: String = DropForMap.name
 
@@ -39,6 +38,22 @@ class DropForMap extends MiniPhase:
         case _ => tree
     case tree => tree
 
+  override def transformInlined(tree: Inlined)(using Context): Tree = tree match
+    case Inlined(call, bindings, expansion) if expansion.hasAttachment(desugar.TrailingForMap) =>
+      call match
+      case Unmapped(f) =>
+        bindings.collectFirst:
+          case vd: ValDef if f.sameTree(vd.rhs) =>
+            expansion.find:
+              case Inlined(Thicket(Nil), Nil, Ident(ident)) => ident == vd.name
+              case _ => false
+            .match
+              case Some(ref) => cpy.Inlined(tree)(call, bindings, ref)
+              case _ => tree
+        .getOrElse(tree)
+      case _ => tree
+    case tree => tree
+
   // Extracts a fun from a possibly nested Apply with lambda and arbitrary implicit args.
   private object Unmapped:
     private def loop(tree: Tree)(using Context): Option[Tree] =

tests/run/better-fors-map-inlined.check

@@ -0,0 +1,4 @@
+MySome(())
+MySome(2)
+MySome((2,3))
+MySome((2,(3,4)))

tests/run/better-fors-map-inlined.scala

@@ -0,0 +1,58 @@
+//> using options -preview
+
+class myOptionModule(doOnMap: => Unit):
+  sealed trait MyOption[+A]:
+    inline 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 MyOption:
+    def apply[A](x: A): MyOption[A] = MySome(x)
+
+@main def Test =
+
+  val myOption = myOptionModule(???)
+
+  import myOption.*
+
+  def portablePrintMyOption(opt: MyOption[Any]): Unit = println:
+    opt match
+    case MySome(()) => "MySome(())"
+    case opt => opt
+
+  val z = for {
+    a <- MyOption(1)
+    b <- MyOption(())
+  } yield ()
+
+  portablePrintMyOption(z)
+
+  val z2 = for {
+    a <- MyOption(1)
+    b <- MyOption(2)
+  } yield b
+
+  portablePrintMyOption(z2)
+
+  val z3 = for {
+    a <- MyOption(1)
+    (b, c) <- MyOption((2, 3))
+  } yield (b, c)
+
+  portablePrintMyOption(z3)
+
+  val z4 = for {
+    a <- MyOption(1)
+    (b, (c, d)) <- MyOption((2, (3, 4)))
+  } yield (b, (c, d))
+
+  portablePrintMyOption(z4)
+end Test