Use untpd.Tree instead of tpd.Tree for SelectionRangeProvider

Co-authored-by: NPCRUS <npcrus@gmail.com>
Co-authored-by: Yummy-Yums <sailewalderbs@gmail.com>

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,30 +82,38 @@ 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)
       selectionRange
 
     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(

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
+      )
+    )

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