From 90ebd6f7f0b456859d5c40cf8bd6aaccb98a70d0 Mon Sep 17 00:00:00 2001
From: rochala <rochala@tutanota.com>
Date: Tue, 4 Mar 2025 02:32:10 +0100
Subject: [PATCH 1/2] Use untpd.Tree instead of tpd.Tree for
 SelectionRangeProvider

Co-authored-by: NPCRUS <npcrus@gmail.com>
Co-authored-by: Yummy-Yums <sailewalderbs@gmail.com>
---
 .../tools/pc/SelectionRangeProvider.scala     | 48 ++++++++++------
 .../tools/pc/tests/SelectionRangeSuite.scala  | 56 ++++++++++++++++++-
 project/Build.scala                           |  4 ++
 3 files changed, 89 insertions(+), 19 deletions(-)

diff --git a/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala b/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala
index 7973f4103ff6..6bf7052caf5f 100644
--- a/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala
+++ b/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala
@@ -6,7 +6,8 @@ import java.util as ju
 import scala.jdk.CollectionConverters._
 import scala.meta.pc.OffsetParams
 
-import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.ast.untpd.*
+import dotty.tools.dotc.ast.NavigateAST
 import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.interactive.Interactive
 import dotty.tools.dotc.interactive.InteractiveDriver
@@ -44,10 +45,13 @@ class SelectionRangeProvider(
       val source = SourceFile.virtual(filePath.toString, text)
       driver.run(uri, source)
       val pos = driver.sourcePosition(param)
-      val path =
-        Interactive.pathTo(driver.openedTrees(uri), pos)(using ctx)
+      val unit = driver.compilationUnits(uri)
 
-      val bareRanges = path
+      val untpdPath: List[Tree] = NavigateAST
+        .pathTo(pos.span, List(unit.untpdTree), true).collect:
+          case untpdTree: Tree => untpdTree
+
+      val bareRanges = untpdPath
         .flatMap(selectionRangesFromTree(pos))
 
       val comments =
@@ -78,7 +82,7 @@ class SelectionRangeProvider(
   end selectionRange
 
   /** Given a tree, create a seq of [[SelectionRange]]s corresponding to that tree. */
-  private def selectionRangesFromTree(pos: SourcePosition)(tree: tpd.Tree)(using Context) =
+  private def selectionRangesFromTree(pos: SourcePosition)(tree: Tree)(using Context) =
     def toSelectionRange(srcPos: SourcePosition) =
       val selectionRange = new SelectionRange()
       selectionRange.setRange(srcPos.toLsp)
@@ -86,22 +90,30 @@ class SelectionRangeProvider(
 
     val treeSelectionRange = toSelectionRange(tree.sourcePos)
 
+    def getArgsSpan(args: List[Tree]): Option[SourcePosition] =
+      args match
+        case Seq(param) => Some(param.sourcePos)
+        case params @ Seq(head, tail*) =>
+          val srcPos = head.sourcePos
+          val lastSpan = tail.last.span
+          Some(SourcePosition(srcPos.source, srcPos.span union lastSpan, srcPos.outer))
+        case Seq() => None
+
     tree match
-      case tpd.DefDef(name, paramss, tpt, rhs) =>
+      case DefDef(_, paramss, _, _) =>
         // If source position is within a parameter list, add a selection range covering that whole list.
-        val selectedParams =
-          paramss
-            .iterator
-            .flatMap: // parameter list to a sourcePosition covering the whole list
-              case Seq(param) => Some(param.sourcePos)
-              case params @ Seq(head, tail*) =>
-                val srcPos = head.sourcePos
-                val lastSpan = tail.last.span
-                Some(SourcePosition(srcPos.source, srcPos.span union lastSpan, srcPos.outer))
-              case Seq() => None
-            .find(_.contains(pos))
-            .map(toSelectionRange)
+        val selectedParams = paramss
+          .flatMap(getArgsSpan) // parameter list to a sourcePosition covering the whole list
+          .find(_.contains(pos))
+          .map(toSelectionRange)
         selectedParams ++ Seq(treeSelectionRange)
+
+      case Function(args, body) =>
+        val allArgs = getArgsSpan(args)
+          .find(_.contains(pos))
+          .map(toSelectionRange)
+
+        allArgs ++ Seq(treeSelectionRange)
       case _ => Seq(treeSelectionRange)
 
   private def setParent(
diff --git a/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala
index 143d998a0ec1..4f8740996bb8 100644
--- a/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala
+++ b/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala
@@ -75,6 +75,12 @@ class SelectionRangeSuite extends BaseSelectionRangeSuite:
            |    b <- Some(2)
            |  } yield a + b
            |}""".stripMargin,
+        """|object Main extends App {
+           |  val total = for {
+           |    >>region>>a <- Some(1)<<region<<
+           |    b <- Some(2)
+           |  } yield a + b
+           |}""".stripMargin,
         """|object Main extends App {
            |  val total = >>region>>for {
            |    a <- Some(1)
@@ -102,7 +108,7 @@ class SelectionRangeSuite extends BaseSelectionRangeSuite:
       )
     )
 
-  @Test def `function params` =
+  @Test def `function-params-1` =
     check(
       """|object Main extends App {
          |  def func(a@@: Int, b: Int) =
@@ -124,6 +130,32 @@ class SelectionRangeSuite extends BaseSelectionRangeSuite:
       )
     )
 
+  @Test def `function-params-2` =
+    check(
+      """|object Main extends App {
+         |  val func = (a@@: Int, b: Int) =>
+         |    a + b
+         |}""".stripMargin,
+      List[String](
+        """|object Main extends App {
+           |  val func = (>>region>>a: Int<<region<<, b: Int) =>
+           |    a + b
+           |}""".stripMargin,
+        """|object Main extends App {
+           |  val func = (>>region>>a: Int, b: Int<<region<<) =>
+           |    a + b
+           |}""".stripMargin,
+        """|object Main extends App {
+           |  val func = >>region>>(a: Int, b: Int) =>
+           |    a + b<<region<<
+           |}""".stripMargin,
+        """|object Main extends App {
+           |  >>region>>val func = (a: Int, b: Int) =>
+           |    a + b<<region<<
+           |}""".stripMargin
+      )
+  )
+
   @Test def `def - type params` =
     check(
       "object Main extends App { def foo[Type@@ <: T1, B](hi: Int, b: Int, c:Int) = ??? }",
@@ -133,3 +165,25 @@ class SelectionRangeSuite extends BaseSelectionRangeSuite:
         "object Main extends App { >>region>>def foo[Type <: T1, B](hi: Int, b: Int, c:Int) = ???<<region<< }"
       )
     )
+
+
+  @Test def `arithmetic` =
+    check(
+      """|object Main extends App {
+         |  def x = 12 * (34 + 5@@6)
+         |}""".stripMargin,
+      List[String](
+        """|object Main extends App {
+           |  def x = 12 * (34 + >>region>>56<<region<<)
+           |}""".stripMargin,
+        """|object Main extends App {
+           |  def x = 12 * (>>region>>34 + 56<<region<<)
+           |}""".stripMargin,
+        """|object Main extends App {
+           |  def x = 12 * >>region>>(34 + 56)<<region<<
+           |}""".stripMargin,
+        """|object Main extends App {
+           |  def x = >>region>>12 * (34 + 56)<<region<<
+           |}""".stripMargin
+      )
+    )
diff --git a/project/Build.scala b/project/Build.scala
index 3e5ef3f43050..f047bc0ad65a 100644
--- a/project/Build.scala
+++ b/project/Build.scala
@@ -1514,6 +1514,10 @@ object Build {
       ivyConfigurations += SourceDeps.hide,
       transitiveClassifiers := Seq("sources"),
       scalacOptions ++= Seq("-source", "3.3"), // To avoid fatal migration warnings
+      publishLocal := publishLocal.dependsOn( // It is best to publish all together. It is not rare to make changes in both compiler / presentation compiler and it can get misaligned
+        `scala3-compiler-bootstrapped` / publishLocal,
+        `scala3-library-bootstrapped` / publishLocal,
+      ).value,
       Compile / scalacOptions ++= Seq("-Yexplicit-nulls", "-Wsafe-init"),
       Compile / sourceGenerators += Def.task {
         val s = streams.value

From af12fac51b0f505224440d72b9f5336b1a3c8d12 Mon Sep 17 00:00:00 2001
From: rochala <rochala@tutanota.com>
Date: Mon, 5 May 2025 18:38:37 +0200
Subject: [PATCH 2/2] Refactor

---
 .../tools/pc/SelectionRangeProvider.scala     | 43 +++++-------
 .../tools/pc/tests/SelectionRangeSuite.scala  | 67 ++++++++++++++++++-
 2 files changed, 82 insertions(+), 28 deletions(-)

diff --git a/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala b/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala
index 6bf7052caf5f..42760f38d9af 100644
--- a/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala
+++ b/presentation-compiler/src/main/dotty/tools/pc/SelectionRangeProvider.scala
@@ -24,10 +24,7 @@ import org.eclipse.lsp4j.SelectionRange
  * @param compiler Metals Global presentation compiler wrapper.
  * @param params offset params converted from the selectionRange params.
  */
-class SelectionRangeProvider(
-    driver: InteractiveDriver,
-    params: ju.List[OffsetParams]
-):
+class SelectionRangeProvider(driver: InteractiveDriver, params: ju.List[OffsetParams]):
 
   /**
    * Get the seletion ranges for the provider params
@@ -88,33 +85,25 @@ class SelectionRangeProvider(
       selectionRange.setRange(srcPos.toLsp)
       selectionRange
 
-    val treeSelectionRange = toSelectionRange(tree.sourcePos)
+    val treeSelectionRange = Seq(toSelectionRange(tree.sourcePos))
 
-    def getArgsSpan(args: List[Tree]): Option[SourcePosition] =
+    def allArgsSelectionRange(args: List[Tree]): Option[SelectionRange] =
       args match
-        case Seq(param) => Some(param.sourcePos)
-        case params @ Seq(head, tail*) =>
-          val srcPos = head.sourcePos
-          val lastSpan = tail.last.span
-          Some(SourcePosition(srcPos.source, srcPos.span union lastSpan, srcPos.outer))
-        case Seq() => None
+        case Nil => None
+        case list =>
+          val srcPos = list.head.sourcePos
+          val lastSpan = list.last.span
+          val allArgsSrcPos = SourcePosition(srcPos.source, srcPos.span union lastSpan, srcPos.outer)
+          if allArgsSrcPos.contains(pos) then Some(toSelectionRange(allArgsSrcPos))
+          else None
 
     tree match
-      case DefDef(_, paramss, _, _) =>
-        // If source position is within a parameter list, add a selection range covering that whole list.
-        val selectedParams = paramss
-          .flatMap(getArgsSpan) // parameter list to a sourcePosition covering the whole list
-          .find(_.contains(pos))
-          .map(toSelectionRange)
-        selectedParams ++ Seq(treeSelectionRange)
-
-      case Function(args, body) =>
-        val allArgs = getArgsSpan(args)
-          .find(_.contains(pos))
-          .map(toSelectionRange)
-
-        allArgs ++ Seq(treeSelectionRange)
-      case _ => Seq(treeSelectionRange)
+      case DefDef(_, paramss, _, _) => paramss.flatMap(allArgsSelectionRange) ++ treeSelectionRange
+      case Apply(_, args) => allArgsSelectionRange(args) ++ treeSelectionRange
+      case TypeApply(_, args) => allArgsSelectionRange(args) ++ treeSelectionRange
+      case UnApply(_, _, pattern) => allArgsSelectionRange(pattern) ++ treeSelectionRange
+      case Function(args, body) => allArgsSelectionRange(args) ++ treeSelectionRange
+      case _ => treeSelectionRange
 
   private def setParent(
       child: SelectionRange,
diff --git a/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala
index 4f8740996bb8..578c8e8bede4 100644
--- a/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala
+++ b/presentation-compiler/test/dotty/tools/pc/tests/SelectionRangeSuite.scala
@@ -172,7 +172,7 @@ class SelectionRangeSuite extends BaseSelectionRangeSuite:
       """|object Main extends App {
          |  def x = 12 * (34 + 5@@6)
          |}""".stripMargin,
-      List[String](
+      List(
         """|object Main extends App {
            |  def x = 12 * (34 + >>region>>56<<region<<)
            |}""".stripMargin,
@@ -187,3 +187,68 @@ class SelectionRangeSuite extends BaseSelectionRangeSuite:
            |}""".stripMargin
       )
     )
+
+  @Test def `function` =
+    check(
+      "val hello = (aaa: Int, bb@@b: Int, ccc: Int) => ???",
+      List(
+        "val hello = (aaa: Int, >>region>>bbb: Int<<region<<, ccc: Int) => ???",
+        "val hello = (>>region>>aaa: Int, bbb: Int, ccc: Int<<region<<) => ???",
+        "val hello = >>region>>(aaa: Int, bbb: Int, ccc: Int) => ???<<region<<",
+        ">>region>>val hello = (aaa: Int, bbb: Int, ccc: Int) => ???<<region<<",
+      )
+    )
+
+  @Test def `defdef` =
+    check(
+      "def hello(aaa: Int, bb@@b: Int, ccc: Int) = ???",
+      List(
+        "def hello(aaa: Int, >>region>>bbb: Int<<region<<, ccc: Int) = ???",
+        "def hello(>>region>>aaa: Int, bbb: Int, ccc: Int<<region<<) = ???",
+        ">>region>>def hello(aaa: Int, bbb: Int, ccc: Int) = ???<<region<<",
+      )
+    )
+
+  @Test def `apply` =
+    check(
+      "def hello = List(111, 2@@22, 333)",
+      List(
+        "def hello = List(111, >>region>>222<<region<<, 333)",
+        "def hello = List(>>region>>111, 222, 333<<region<<)",
+        "def hello = >>region>>List(111, 222, 333)<<region<<",
+        ">>region>>def hello = List(111, 222, 333)<<region<<",
+      )
+    )
+
+  @Test def `type-apply` =
+    check(
+      "def hello = Map[String, I@@nt]()",
+      List(
+        "def hello = Map[String, >>region>>Int<<region<<]()",
+        "def hello = Map[>>region>>String, Int<<region<<]()",
+        "def hello = >>region>>Map[String, Int]<<region<<()",
+        "def hello = >>region>>Map[String, Int]()<<region<<",
+        ">>region>>def hello = Map[String, Int]()<<region<<",
+      )
+    )
+
+  @Test def `unapply` =
+    check(
+      "val List(aaa, b@@bb, ccc) = List(111, 222, 333)",
+      List(
+        "val List(aaa, >>region>>bbb<<region<<, ccc) = List(111, 222, 333)",
+        "val List(>>region>>aaa, bbb, ccc<<region<<) = List(111, 222, 333)",
+        "val >>region>>List(aaa, bbb, ccc)<<region<< = List(111, 222, 333)",
+        ">>region>>val List(aaa, bbb, ccc) = List(111, 222, 333)<<region<<",
+      )
+    )
+
+  @Test def `single` =
+    check(
+      "def hello = List(2@@22)",
+      List(
+        "def hello = List(>>region>>222<<region<<)",
+        "def hello = >>region>>List(222)<<region<<",
+        ">>region>>def hello = List(222)<<region<<",
+      )
+    )