1 files changed, 14 insertions, 2 deletions
diff --git a/src/dotty/tools/dotc/ast/Desugar.scala b/src/dotty/tools/dotc/ast/Desugar.scala
index 346af42b8..500b28233 100644
--- a/src/dotty/tools/dotc/ast/Desugar.scala
+++ b/src/dotty/tools/dotc/ast/Desugar.scala
@@ -8,6 +8,7 @@ import SymDenotations._, Symbols._, StdNames._, Annotations._, Trees._
 import Decorators._
 import language.higherKinds
 import collection.mutable.ListBuffer
+import util.Attachment
 import config.Printers._
 
 object desugar {
@@ -17,6 +18,11 @@ object desugar {
 
   import untpd._
 
+  /** Tags a .withFilter call generated by desugaring a for expression.
+   *  Such calls can alternatively be rewritten to use filter.
+   */
+  val MaybeFilter = new Attachment.Key[Unit]
+
   /** Info of a variable in a pattern: The named tree and its type */
   private type VarInfo = (NameTree, Tree)
 
@@ -773,6 +779,12 @@ object desugar {
           (Bind(name, pat), Ident(name))
       }
 
+      /** Add MaybeFilter attachment */
+      def orFilter(tree: Tree): tree.type = {
+        tree.putAttachment(MaybeFilter, ())
+        tree
+      }
+
       /** Make a pattern filter:
        *    rhs.withFilter { case pat => true case _ => false }
        *
@@ -803,7 +815,7 @@ object desugar {
         val cases = List(
           CaseDef(pat, EmptyTree, Literal(Constant(true))),
           CaseDef(Ident(nme.WILDCARD), EmptyTree, Literal(Constant(false))))
-        Apply(Select(rhs, nme.withFilter), makeCaseLambda(cases))
+        Apply(orFilter(Select(rhs, nme.withFilter)), makeCaseLambda(cases))
       }
 
       /** Is pattern `pat` irrefutable when matched against `rhs`?
@@ -858,7 +870,7 @@ object desugar {
           val vfrom1 = new IrrefutableGenFrom(makeTuple(allpats), rhs1)
           makeFor(mapName, flatMapName, vfrom1 :: rest1, body)
         case (gen: GenFrom) :: test :: rest =>
-          val filtered = Apply(rhsSelect(gen, nme.withFilter), makeLambda(gen.pat, test))
+          val filtered = Apply(orFilter(rhsSelect(gen, nme.withFilter)), makeLambda(gen.pat, test))
           val genFrom =
             if (isIrrefutableGenFrom(gen)) new IrrefutableGenFrom(gen.pat, filtered)
             else GenFrom(gen.pat, filtered)