From 9f2280566e1a0804961d96e5937c11fc4be5a470 Mon Sep 17 00:00:00 2001
From: Tomasz Godzik <tgodzik@virtuslab.com>
Date: Sat, 15 Mar 2025 19:35:00 +0100
Subject: [PATCH] feature: Skip auto importing symbols we know are wrong in
 current context

---
 .../dotty/tools/pc/AutoImportsProvider.scala  | 30 ++++++++-
 .../pc/tests/edit/AutoImportsSuite.scala      | 63 +++++++++++++++++++
 2 files changed, 90 insertions(+), 3 deletions(-)

diff --git a/presentation-compiler/src/main/dotty/tools/pc/AutoImportsProvider.scala b/presentation-compiler/src/main/dotty/tools/pc/AutoImportsProvider.scala
index e35556ad11c9..d30ab813f9f2 100644
--- a/presentation-compiler/src/main/dotty/tools/pc/AutoImportsProvider.scala
+++ b/presentation-compiler/src/main/dotty/tools/pc/AutoImportsProvider.scala
@@ -10,6 +10,7 @@ import scala.meta.pc.*
 
 import dotty.tools.dotc.ast.tpd.*
 import dotty.tools.dotc.core.Symbols.*
+import dotty.tools.dotc.core.StdNames.*
 import dotty.tools.dotc.interactive.Interactive
 import dotty.tools.dotc.interactive.InteractiveDriver
 import dotty.tools.dotc.util.SourceFile
@@ -17,6 +18,7 @@ import dotty.tools.pc.completions.CompletionPos
 import dotty.tools.pc.utils.InteractiveEnrichments.*
 
 import org.eclipse.lsp4j as l
+import dotty.tools.dotc.core.Flags.Method
 
 final class AutoImportsProvider(
     search: SymbolSearch,
@@ -47,6 +49,17 @@ final class AutoImportsProvider(
     )
     import indexedContext.ctx
 
+
+    def correctInTreeContext(sym: Symbol) = path match
+      case (_: Ident) :: (sel: Select) :: _ =>
+        sym.info.allMembers.exists(_.name == sel.name)
+      case (_: Ident) :: (_: Apply) :: _ if !sym.is(Method) =>
+        def applyInObject =
+          sym.companionModule.info.allMembers.exists(_.name == nme.apply)
+        def applyInClass = sym.info.allMembers.exists(_.name == nme.apply)
+        applyInClass || applyInObject
+      case _ => true
+
     val isSeen = mutable.Set.empty[String]
     val symbols = List.newBuilder[Symbol]
     def visit(sym: Symbol): Boolean =
@@ -90,13 +103,24 @@ final class AutoImportsProvider(
         end match
       end mkEdit
 
-      for
+      val all = for
         sym <- results
         edits <- mkEdit(sym)
-      yield AutoImportsResultImpl(
+      yield (AutoImportsResultImpl(
         sym.owner.showFullName,
         edits.asJava
-      )
+      ), sym)
+
+      all match
+        case (onlyResult, _) :: Nil => List(onlyResult)
+        case Nil => Nil
+        case moreResults =>
+          val moreExact = moreResults.filter { case (_, sym) =>
+            correctInTreeContext(sym)
+          }
+          if moreExact.nonEmpty then moreExact.map(_._1)
+          else moreResults.map(_._1)
+
     else List.empty
     end if
   end autoImports
diff --git a/presentation-compiler/test/dotty/tools/pc/tests/edit/AutoImportsSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/edit/AutoImportsSuite.scala
index 3bb5bfea7bc0..bf568bc0636b 100644
--- a/presentation-compiler/test/dotty/tools/pc/tests/edit/AutoImportsSuite.scala
+++ b/presentation-compiler/test/dotty/tools/pc/tests/edit/AutoImportsSuite.scala
@@ -12,6 +12,69 @@ class AutoImportsSuite extends BaseAutoImportsSuite:
          |  <<Future>>.successful(2)
          |}
          |""".stripMargin,
+      """|scala.concurrent
+         |""".stripMargin
+    )
+
+  @Test def `basic-apply` =
+    check(
+      """|object A {
+          |  <<Future>>(2)
+          |}
+          |""".stripMargin,
+      """|scala.concurrent
+         |""".stripMargin,
+    )
+
+  @Test def `basic-function-apply` =
+    check(
+      """|
+          |object ForgeFor{
+          |  def importMe(): Int = ???
+          |}
+          |object ForgeFor2{
+          |  case class importMe()
+          |}
+          |
+          |
+          |object test2 {
+          |  <<importMe>>()
+          |}
+          |""".stripMargin,
+      """|ForgeFor2
+         |ForgeFor
+         |""".stripMargin
+    )
+
+  @Test def `basic-apply-wrong` =
+    check(
+      """|object A {
+          |  new <<Future>>(2)
+          |}
+          |""".stripMargin,
+      """|scala.concurrent
+         |java.util.concurrent
+         |""".stripMargin
+    )
+
+  @Test def `basic-fuzzy` =
+    check(
+      """|object A {
+          |  <<Future>>.thisMethodDoesntExist(2)
+          |}
+          |""".stripMargin,
+      """|scala.concurrent
+         |java.util.concurrent
+         |""".stripMargin
+    )
+
+  @Test def `typed-simple` =
+    check(
+      """|object A {
+          |  import scala.concurrent.Promise
+          |  val fut: <<Future>> = Promise[Unit]().future
+          |}
+          |""".stripMargin,
       """|scala.concurrent
          |java.util.concurrent
          |""".stripMargin