Implement alternative desugaring of for-if to filter.

Fallback to .filter if a .withFilter is not available, but do this
only for .withFilter calls generated from for expressions (this is
different from what scalac does; the latter can also rewrite
.withFilter calls given in source, which is not desirable.

2 files changed, 23 insertions, 5 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)
diff --git a/src/dotty/tools/dotc/typer/Typer.scala b/src/dotty/tools/dotc/typer/Typer.scala
index 934ca556d..c9a2e4b48 100644
--- a/src/dotty/tools/dotc/typer/Typer.scala
+++ b/src/dotty/tools/dotc/typer/Typer.scala
@@ -346,11 +346,17 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
       }
     }
 
-    if (ctx.compilationUnit.isJava && tree.name.isTypeName) {
+    def selectWithFallback(fallBack: => Tree) =
+      tryEither(tryCtx => asSelect(tryCtx))((_, _) => fallBack)
+
+    if (ctx.compilationUnit.isJava && tree.name.isTypeName)
       // SI-3120 Java uses the same syntax, A.B, to express selection from the
       // value A and from the type A. We have to try both.
-      tryEither(tryCtx => asSelect(tryCtx))((_, _) => asJavaSelectFromTypeTree(ctx))
-    } else asSelect(ctx)
+      selectWithFallback(asJavaSelectFromTypeTree(ctx))
+    else if (tree.name == nme.withFilter && tree.getAttachment(desugar.MaybeFilter).isDefined)
+      selectWithFallback(typedSelect(untpd.cpy.Select(tree)(tree.qualifier, nme.filter), pt))
+    else
+      asSelect(ctx)
   }
 
   def typedSelectFromTypeTree(tree: untpd.SelectFromTypeTree, pt: Type)(implicit ctx: Context): Tree = track("typedSelectFromTypeTree") {