Merge pull request #351 from sjrd/auto-positions

Automatically fill in positions of trees based on their children.

build.sbt

@@ -121,6 +121,11 @@ lazy val tastyQuery =
           ProblemFilters.exclude[DirectMissingMethodProblem]("tastyquery.Types#PolyType.fromParamsSymbols"),
           ProblemFilters.exclude[DirectMissingMethodProblem]("tastyquery.Types#TypeLambda.fromParamsSymbols"),
           ProblemFilters.exclude[DirectMissingMethodProblem]("tastyquery.Types#TypeLambdaTypeCompanion.fromParamsSymbols"),
+
+          // New abstract methods in a completely sealed hierarchy, not an issue
+          ProblemFilters.exclude[InheritedNewAbstractMethodProblem]("tastyquery.Trees#*.canEqual"),
+          ProblemFilters.exclude[InheritedNewAbstractMethodProblem]("tastyquery.Trees#*.productArity"),
+          ProblemFilters.exclude[InheritedNewAbstractMethodProblem]("tastyquery.Trees#*.productElement"),
         )
       },
     )

tasty-query/shared/src/main/scala/tastyquery/SourcePosition.scala

@@ -3,9 +3,11 @@ package tastyquery
 import tastyquery.Spans.Span
 import tastyquery.Spans.NoSpan
 
-final class SourcePosition private[tastyquery] (val sourceFile: SourceFile, private val span: Span):
+final class SourcePosition private[tastyquery] (val sourceFile: SourceFile, private[tastyquery] val span: Span):
   override def toString(): String = s"$sourceFile:$span"
 
+  private[tastyquery] def isAuto: Boolean = !span.exists && sourceFile != SourceFile.NoSource
+
   /** True if this source position is unknown. */
   def isUnknown: Boolean = !span.exists
 
@@ -76,4 +78,7 @@ end SourcePosition
 
 object SourcePosition:
   val NoPosition: SourcePosition = new SourcePosition(SourceFile.NoSource, NoSpan)
+
+  private[tastyquery] def auto(source: SourceFile, span: Span): SourcePosition =
+    new SourcePosition(source, span)
 end SourcePosition

tasty-query/shared/src/main/scala/tastyquery/Trees.scala

@@ -1,5 +1,6 @@
 package tastyquery
 
+import scala.annotation.constructorOnly
 import scala.annotation.tailrec
 
 import scala.collection.mutable
@@ -16,7 +17,40 @@ import tastyquery.Types.*
 
 object Trees {
 
-  sealed abstract class Tree(val pos: SourcePosition) {
+  sealed abstract class Tree(@constructorOnly posInit: SourcePosition) extends Product {
+    val pos: SourcePosition =
+      if posInit.isAuto then (this: @unchecked).computeAutoPos(posInit.sourceFile)
+      else posInit
+
+    private def computeAutoPos(source: SourceFile): SourcePosition = this match
+      case Inlined(call, _, _) =>
+        call.pos
+
+      case _ =>
+        def rec(span: Span, item: Any): Span = item.asInstanceOf[Matchable] match
+          case item: Tree =>
+            if item.pos.sourceFile == source then span.union(item.pos.span)
+            else span
+          case x :: xs =>
+            rec(rec(span, x), xs)
+          case Left(x) =>
+            rec(span, x)
+          case Right(x) =>
+            rec(span, x)
+          case _ =>
+            span
+        end rec
+
+        var span = NoSpan
+        val len = productArity
+        var i = 0
+        while i != len do
+          span = rec(span, productElement(i))
+          i += 1
+
+        SourcePosition(source, span)
+    end computeAutoPos
+
     def withPos(pos: SourcePosition): Tree
 
     private def subtrees: List[Tree] =

tasty-query/shared/src/main/scala/tastyquery/reader/tasties/TreeUnpickler.scala

@@ -201,18 +201,16 @@ private[tasties] class TreeUnpickler private (
   private def posErrorMsg: String = s"at address ${reader.currentAddr} in file $filename"
   private def posErrorMsg(atAddr: Addr): String = s"at address $atAddr in file $filename"
 
-  def spanAt(addr: Addr)(using SourceFile): SourcePosition =
+  def spanAt(addr: Addr): Span =
     posUnpicklerOpt match {
       case Some(posUnpickler) =>
-        val span = posUnpickler.spanAt(addr)
-        if span.exists then SourcePosition(summon[SourceFile], span)
-        else SourcePosition.NoPosition
+        posUnpickler.spanAt(addr)
       case _ =>
-        SourcePosition.NoPosition
+        NoSpan
     }
 
-  def span(using SourceFile): SourcePosition =
-    spanAt(reader.currentAddr)
+  def span(using source: SourceFile): SourcePosition =
+    SourcePosition.auto(source, spanAt(reader.currentAddr))
 
   private inline def maybeAdjustSourceFileIn[A](inline op: SourceFile ?=> A)(using source: SourceFile): A =
     val newSourceFile = posUnpicklerOpt match
@@ -805,11 +803,12 @@ private[tasties] class TreeUnpickler private (
       if name == nme.Wildcard || name == nme.WildcardSequence then WildcardPattern(typ.requireType)(spn)
       else ExprPattern(makeIdent(name, typ, spn))(spn)
     case TYPED =>
+      val spn = span
       reader.readByte()
       reader.readEnd()
       val body = readPattern
       val tpt = readTypeTree
-      TypeTest(body, tpt)(body.pos.union(tpt.pos))
+      TypeTest(body, tpt)(spn)
     case BIND =>
       val spn = span
       val start = reader.currentAddr
@@ -823,10 +822,10 @@ private[tasties] class TreeUnpickler private (
       symbol.withDeclaredType(typ)
       definingTree(symbol, Bind(name, body, symbol)(spn))
     case ALTERNATIVE =>
+      val spn = span
       reader.readByte()
       val end = reader.readEnd()
       val alts = reader.until(end)(readPattern)
-      val spn = alts.map(_.pos).reduce(_.union(_))
       Alternative(alts)(spn)
     case UNAPPLY =>
       val spn = span
@@ -844,7 +843,8 @@ private[tasties] class TreeUnpickler private (
     case SHAREDterm =>
       val spn = span
       reader.readByte()
-      forkAt(reader.readAddr()).readPattern.withPos(spn)
+      val shared = forkAt(reader.readAddr()).readPattern
+      if spn.isUnknown then shared else shared.withPos(spn)
     case _ =>
       val expr = readTerm
       ExprPattern(expr)(expr.pos)
@@ -939,11 +939,12 @@ private[tasties] class TreeUnpickler private (
       val cls = readTypeTree
       New(cls)(spn)
     case TYPED =>
+      val spn = span
       reader.readByte()
       reader.readEnd()
       val expr = readTerm
       val tpt = readTypeTree
-      Typed(expr, tpt)(expr.pos.union(tpt.pos))
+      Typed(expr, tpt)(spn)
     case THROW =>
       val spn = span
       reader.readByte()
@@ -958,13 +959,12 @@ private[tasties] class TreeUnpickler private (
       val finalizer = reader.ifBeforeOpt(end)(readTerm)
       Try(expr, catchCases, finalizer)(spn)
     case ASSIGN =>
+      val spn = span
       reader.readByte()
       reader.readEnd()
-      val lhsSpan = span
       val lhs = readTerm
-      val rhsSpan = span
       val rhs = readTerm
-      Assign(lhs, rhs)(lhsSpan.union(rhsSpan))
+      Assign(lhs, rhs)(spn)
     case BLOCK =>
       val spn = span
       reader.readByte()
@@ -1046,7 +1046,8 @@ private[tasties] class TreeUnpickler private (
     case SHAREDterm =>
       val spn = span
       reader.readByte()
-      forkAt(reader.readAddr()).readTerm.withPos(spn)
+      val shared = forkAt(reader.readAddr()).readTerm
+      if spn.isUnknown then shared else shared.withPos(spn)
 
     // paths
     case THIS =>
@@ -1091,7 +1092,8 @@ private[tasties] class TreeUnpickler private (
     case SHAREDtype =>
       val spn = span
       reader.readByte()
-      forkAt(reader.readAddr()).readTerm.withPos(spn)
+      val shared = forkAt(reader.readAddr()).readTerm
+      if spn.isUnknown then shared else shared.withPos(spn)
     case tag if isConstantTag(tag) =>
       val spn = span
       Literal(readConstant)(spn)
@@ -1585,7 +1587,8 @@ private[tasties] class TreeUnpickler private (
     case SHAREDterm =>
       val spn = span
       reader.readByte()
-      forkAt(reader.readAddr()).readTypeTree.withPos(spn)
+      val shared = forkAt(reader.readAddr()).readTypeTree
+      if spn.isUnknown then shared else shared.withPos(spn)
     case tag if isTypeTreeTag(tag) =>
       throw TastyFormatException(s"Unexpected type tree tag ${astTagToString(tag)} $posErrorMsg")
     case _ =>

tasty-query/shared/src/test/scala/tastyquery/PositionSuite.scala

@@ -78,7 +78,7 @@ class PositionSuite extends RestrictedUnpicklingSuite {
   }
 
   testUnpickleWithCode("apply", "simple_trees.EtaExpansion") { (tree, code) =>
-    assertEquals(collectCode[Apply](tree, code), List("f(0)", "takesFunction(intMethod)"))
+    assertEquals(collectCode[Apply](tree, code), List("f(0)", "takesFunction(intMethod)", "intMethod"))
   }
 
   /** Control structures */
@@ -109,18 +109,23 @@ class PositionSuite extends RestrictedUnpicklingSuite {
   testUnpickleWithCode("match", "simple_trees.Match") { (tree, code) =>
     assertEquals(
       collectCode[Match](tree, code),
-      List("""x match {
-             |    case 0 => 0
-             |    case 1 | -1 | 2 => x + 1
-             |    case 7 if x == 7 => x - 1
-             |    case 3 | 4 | 5 if x < 5 => 0
-             |    case _ if (x % 2 == 0) => x / 2
-             |    case _ => -x
-             |  }""".stripMargin)
+      List(
+        """x match {
+          |    case 0 => 0
+          |    case 1 | -1 | 2 => x + 1
+          |    case 7 if x == 7 => x - 1
+          |    case 3 | 4 | 5 if x < 5 => 0
+          |    case _ if (x % 2 == 0) => x / 2
+          |    case _ => -x
+          |  }""".stripMargin,
+        """xs match
+          |    case List(elems: _*) => 0
+          |    case _               => 1""".stripMargin
+      )
     )
 
-    val matchPos = findTree(tree) {
-      case matchTree: Match if !matchTree.pos.isUnknown => matchTree.pos
+    val matchPos = findTree(tree) { case matchTree @ Match(Ident(SimpleName("x")), _) =>
+      matchTree.pos
     }
     assert(clue(matchPos.startLine) == 3)
     assert(clue(matchPos.startColumn) == 23)
@@ -318,7 +323,7 @@ class PositionSuite extends RestrictedUnpicklingSuite {
   }
 
   testUnpickleWithCode("type-apply", "simple_trees.TypeApply") { (tree, code) =>
-    assertEquals(collectCode[TypeApply](tree, code), List("Seq[Int]", "A[Int, Seq[String]]"))
+    assertEquals(collectCode[TypeApply](tree, code), List("Seq", "Seq[Int]", "A[Int, Seq[String]]"))
   }
 
   testUnpickleWithCode("type-ident", "simple_trees.Typed") { (tree, code) =>
@@ -356,6 +361,8 @@ class PositionSuite extends RestrictedUnpicklingSuite {
         """X match {
           |    case Int => String
           |  }""".stripMargin,
+        """Product = X match {
+          |    case Int => Some[Int]""".stripMargin,
         """X match {
           |    case _ => Int
           |  }""".stripMargin,
@@ -375,11 +382,33 @@ class PositionSuite extends RestrictedUnpicklingSuite {
   }
 
   testUnpickleWithCode("type-definition-tree-1", "simple_trees.TypeMember") { (tree, code) =>
-    assertEquals(collectCode[TypeDefinitionTree](tree, code), List("Int", ">: Null <: Product", "Int"))
+    assertEquals(
+      collectCode[TypeDefinitionTree](tree, code),
+      List(
+        "Int",
+        /* The following makes no sense; it is the position of the type def
+         * of `type AbstractType`. It's probably dotc's auto-assigning of
+         * positions that goes wild.
+         */
+        """type AbstractType
+          |  type AbstractWithBounds >: Null <: Product""".stripMargin,
+        ">: Null <: Product",
+        "Int",
+        "Null <: Product = Null",
+        "Null <: Product"
+      )
+    )
   }
 
   testUnpickleWithCode("type-definition-tree-2", "simple_trees.TypeLambda") { (tree, code) =>
-    assertEquals(collectCode[TypeDefinitionTree](tree, code), List("[X] =>> List[X]", "List[X]"))
+    assertEquals(
+      collectCode[TypeDefinitionTree](tree, code),
+      List(
+        "[X] =>> List[X]",
+        "X] =>> List[X]", // TODO Improve this
+        "List[X]"
+      )
+    )
   }
 
   /** Inlined */