Drop trailing implicit args to dropped map

Author: som-snytt

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

@@ -12,42 +12,62 @@ import dotty.tools.dotc.transform.MegaPhase.MiniPhase
 import dotty.tools.dotc.ast.desugar
 
 /** Drop unused trailing map calls in for comprehensions.
-  * We can drop the map call if:
-  * - it won't change the type of the expression, and
-  * - the function is an identity function or a const function to unit.
-  *
-  * The latter condition is checked in [[Desugar.scala#makeFor]]
-  */
+ *
+ *  We can drop the map call if:
+ *  - it won't change the type of the expression, and
+ *  - the function is an identity function or a const function to unit.
+ *
+ *  The latter condition is checked in [[Desugar.scala#makeFor]]
+ */
 class DropForMap extends MiniPhase:
   import DropForMap.*
 
   override def phaseName: String = DropForMap.name
 
   override def description: String = DropForMap.description
 
-  override def transformApply(tree: Apply)(using Context): Tree =
-    tree.removeAttachment(desugar.TrailingForMap) match
-    case Some(_) =>
+  /** r.map(x => x)(using y) --> r
+   *       ^ TrailingForMap
+   */
+  override def transformApply(tree: Apply)(using Context): Tree = tree match
+    case Unmapped(f) =>
+      if f.tpe =:= tree.tpe then // make sure that the type of the expression won't change
+        f // drop the map call
+      else
+        f match
+        case Converted(r) if r.tpe =:= tree.tpe => r // drop the map call and the conversion
+        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] =
       tree match
-      case aply @ Apply(MapCall(f), List(Lambda(List(param), body))) =>
-        if f.tpe =:= aply.tpe then // make sure that the type of the expression won't change
-          return f // drop the map call
-        else
-          f match
-          case Converted(r) if r.tpe =:= aply.tpe =>
-            return r // drop the map call and the conversion
+      case Apply(fun, Lambda(_ :: Nil, _) :: Nil) =>
+        tree.removeAttachment(desugar.TrailingForMap) match
+        case Some(_) =>
+          fun match
+          case MapCall(f) => return Some(f)
           case _ =>
+        case _ =>
+      case Apply(fun, _) =>
+        fun.tpe match
+        case mt: MethodType if mt.isImplicitMethod => return loop(fun)
+        case _ =>
       case _ =>
-    case _ =>
-    tree
+      None
+    end loop
+    def unapply(tree: Apply)(using Context): Option[Tree] =
+      tree.tpe match
+      case _: MethodOrPoly => None
+      case _ => loop(tree)
 
   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
+    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

tests/run/i23409b.scala

@@ -0,0 +1,26 @@
+//> using options -preview
+
+final class Implicit()
+
+final class Id[+A, -U](val value: A):
+  def map[B](f: A => B)(using Implicit): Id[B, U] = ??? //Id(f(value))
+  def flatMap[B, V <: U](f: A => Id[B, V]): Id[B, V] = f(value)
+  def run: A = value
+
+type Foo = Foo.type
+case object Foo:
+  def get: Id[Int, Foo] = Id(42)
+
+type Bar = Bar.type
+case object Bar:
+  def inc(i: Int): Id[Int, Bar] = Id(i * 10)
+
+def program(using Implicit) =
+  for
+    a <- Foo.get
+    x <- Bar.inc(a)
+  yield x
+
+@main def Test = println:
+  given Implicit = Implicit()
+  program.run