From 41c735e5b48d565210980538690ed3a180594393 Mon Sep 17 00:00:00 2001
From: Katarzyna Marek <kasia@marek.net>
Date: Wed, 24 Apr 2024 12:42:26 +0200
Subject: [PATCH] fix: prefer non-export definition locations

---
 .../dotty/tools/pc/PcDefinitionProvider.scala   | 11 +++++++----
 .../pc/tests/definition/PcDefinitionSuite.scala | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/presentation-compiler/src/main/dotty/tools/pc/PcDefinitionProvider.scala b/presentation-compiler/src/main/dotty/tools/pc/PcDefinitionProvider.scala
index 0de81ec39711..536ddae4203c 100644
--- a/presentation-compiler/src/main/dotty/tools/pc/PcDefinitionProvider.scala
+++ b/presentation-compiler/src/main/dotty/tools/pc/PcDefinitionProvider.scala
@@ -13,7 +13,7 @@ import dotty.tools.dotc.ast.NavigateAST
 import dotty.tools.dotc.ast.tpd.*
 import dotty.tools.dotc.ast.untpd
 import dotty.tools.dotc.core.Contexts.Context
-import dotty.tools.dotc.core.Flags.ModuleClass
+import dotty.tools.dotc.core.Flags.{Exported, ModuleClass}
 import dotty.tools.dotc.core.Symbols.*
 import dotty.tools.dotc.interactive.Interactive
 import dotty.tools.dotc.interactive.InteractiveDriver
@@ -123,9 +123,12 @@ class PcDefinitionProvider(
       case symbols @ (sym :: other) =>
         val isLocal = sym.source == pos.source
         if isLocal then
-          val defs =
-            Interactive.findDefinitions(List(sym), driver, false, false).filter(_.source == sym.source)
-          defs.headOption match
+          val (exportedDefs, otherDefs) =
+            Interactive.findDefinitions(List(sym), driver, false, false)
+              .filter(_.source == sym.source)
+              .partition(_.tree.symbol.is(Exported))
+
+          otherDefs.headOption.orElse(exportedDefs.headOption)  match
             case Some(srcTree) =>
               val pos = srcTree.namePos
               pos.toLocation match
diff --git a/presentation-compiler/test/dotty/tools/pc/tests/definition/PcDefinitionSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/definition/PcDefinitionSuite.scala
index 9636aea77c2e..c7c9b9979404 100644
--- a/presentation-compiler/test/dotty/tools/pc/tests/definition/PcDefinitionSuite.scala
+++ b/presentation-compiler/test/dotty/tools/pc/tests/definition/PcDefinitionSuite.scala
@@ -274,6 +274,23 @@ class PcDefinitionSuite extends BasePcDefinitionSuite:
         |""".stripMargin
     )
 
+  @Test def exportTermExtension =
+    check(
+      """|package a
+         |class Test extends A {
+         |  assert("Hello".fo@@o == "HelloFoo")
+         |}
+         |
+         |trait A {
+         |  export B.*
+         |}
+         |
+         |object B {
+         |  extension (value: String) def <<foo>>: String = s"${value}Foo"
+         |}
+         |""".stripMargin
+    )
+
   @Test def `named-arg-local` =
     check(
       """|