Repository: scalacenter/tasty-query (scala)

Issue:
Nested packages can't be found from their parent packages
`lang` cannot be found as a member of `java`. Test case:
```scala
    println(resolve(name"java" / name"lang" / tname"String")) // OK: symbol[class String]
    val javaPackage = resolve(name"java".singleton).asInstanceOf[PackageClassSymbol]
    println(javaPackage) // OK: symbol[package java]
    println(javaPackage.getDecl(name"lang")) // Not OK: None
    println(TypeRef(TermRef(PackageRef(javaPackage), name"lang"), tname"String").symbol) // Not OK: throws
```

Pull request description:
store packages in decls by their short name.

remove the `lookup` method, migration is to cast to `DeclaringSymbol` and call `getDecl`

fixes #71

Patch:
diff --git a/shared/src/main/scala/tastyquery/Contexts.scala b/shared/src/main/scala/tastyquery/Contexts.scala
index 13342a41..5f56c841 100644
--- a/shared/src/main/scala/tastyquery/Contexts.scala
+++ b/shared/src/main/scala/tastyquery/Contexts.scala
@@ -77,15 +77,21 @@ object Contexts {
 
     def findSymbolFromRoot(path: List[Name]): Symbol =
       @tailrec
-      def rec(owner: Symbol, path: List[Name]): Symbol =
+      def rec(symbol: Symbol, path: List[Name]): Symbol =
         path match
           case Nil =>
-            owner
+            symbol
           case name :: pathRest =>
-            val sym = owner.lookup(name).getOrElse {
-              throw new IllegalArgumentException(s"cannot find member ${name.toDebugString} in $owner")
+            val owner = symbol match
+              case owner: DeclaringSymbol => owner
+              case _ =>
+                throw IllegalArgumentException(
+                  s"$symbol does not declare a scope, cannot find member ${name.toDebugString}"
+                )
+            val next = owner.getDecl(name).getOrElse {
+              throw IllegalArgumentException(s"cannot find member ${name.toDebugString} in $symbol")
             }
-            rec(sym, pathRest)
+            rec(next, pathRest)
 
       rec(defn.RootPackage, path)
     end findSymbolFromRoot
diff --git a/shared/src/main/scala/tastyquery/ast/Symbols.scala b/shared/src/main/scala/tastyquery/ast/Symbols.scala
index 25f6d1ea..604fcb11 100644
--- a/shared/src/main/scala/tastyquery/ast/Symbols.scala
+++ b/shared/src/main/scala/tastyquery/ast/Symbols.scala
@@ -159,10 +159,6 @@ object Symbols {
       case scope: DeclaringSymbol => scope.allOverloads(name)
       case _                      => Nil
 
-    def lookup(name: Name)(using Context): Option[Symbol] = this match
-      case scope: DeclaringSymbol => scope.getDecl(name)
-      case _                      => None
-
     override def toString: String = {
       val kind = this match
         case _: PackageClassSymbol => "package "
@@ -201,11 +197,14 @@ object Symbols {
       mutable.HashMap[Name, mutable.HashSet[Symbol]]()
 
     private[tastyquery] final def addDecl(decl: Symbol): Unit =
-      myDeclarations.getOrElseUpdate(decl.name, new mutable.HashSet) += decl
-
-    private[tastyquery] final def getDeclInternal(name: Name): Option[Symbol] = name match {
-      case overloaded: SignedName =>
-        None // need context to resolve overloads
+      val key = decl.name match
+        case QualifiedName(_, _, last) => last // drop package prefixes
+        case name                      => name
+      myDeclarations.getOrElseUpdate(key, new mutable.HashSet) += decl
+
+    /** direct lookup without requiring `Context`, can not resolve overloads */
+    private[tastyquery] final def getDeclInternal(name: Name): Option[Symbol] = name match
+      case overloaded: SignedName => None // need context to resolve overloads
       case name =>
         myDeclarations.get(name) match
           case Some(decls) =>
@@ -214,7 +213,6 @@ object Symbols {
               throw SymbolLookupException(name, s"unexpected overloads: ${decls.mkString(", ")}")
             else None // TODO: this should be an error
           case _ => None
-    }
 
     /** Forces the symbol to be initialized.
       *
@@ -356,24 +354,17 @@ object Symbols {
     this.withDeclaredType(PackageRef(this))
 
     def findPackageSymbol(packageName: TermName): Option[PackageClassSymbol] = packageName match {
-      case _: SimpleName => getPackageDecl(packageName)
+      case packageName: SimpleName => getPackageDecl(packageName)
       case QualifiedName(NameTags.QUALIFIED, prefix, suffix) =>
         if (prefix == name)
-          getPackageDecl(packageName)
+          getPackageDecl(suffix)
         else
           // recurse
           findPackageSymbol(prefix).flatMap(_.findPackageSymbol(packageName))
       case _ => throw IllegalArgumentException(s"Unexpected package name: $name")
     }
 
-    override def lookup(name: Name)(using Context): Option[Symbol] =
-      val sel = name match {
-        case name: SimpleName if this.name != nme.RootName => this.name.toTermName.select(name)
-        case _                                             => name
-      }
-      getDecl(sel)
-
-    private def getPackageDecl(packageName: TermName): Option[PackageClassSymbol] =
+    private def getPackageDecl(packageName: SimpleName): Option[PackageClassSymbol] =
       /* All subpackages are created eagerly when initializing contexts,
        * so we can use getDeclInternal here.
        */
diff --git a/shared/src/main/scala/tastyquery/ast/Types.scala b/shared/src/main/scala/tastyquery/ast/Types.scala
index b8300ace..c266d053 100644
--- a/shared/src/main/scala/tastyquery/ast/Types.scala
+++ b/shared/src/main/scala/tastyquery/ast/Types.scala
@@ -455,7 +455,7 @@ object Types {
         sym
       case Scala2ExternalSymRef(owner, path) =>
         val sym = path.foldLeft(owner) { (owner, name) =>
-          owner.lookup(name).getOrElse {
+          owner.asClass.getDecl(name).getOrElse {
             throw new SymbolLookupException(name, s"cannot find symbol $owner.$name")
           }
         }
@@ -725,7 +725,7 @@ object Types {
       val sym = resolveToSymbol
       sym match
         case sym: ClassSymbol =>
-          sym.lookup(name).getOrElse {
+          sym.getDecl(name).getOrElse {
             throw new AssertionError(s"Cannot find member named '$name' in $pre")
           }
         case _ =>
@@ -1183,7 +1183,7 @@ object Types {
 
   case class ClassType(cls: ClassSymbol, rawParents: Type) extends GroundType {
     def findMember(name: Name, pre: Type)(using Context): Symbol =
-      cls.lookup(name).getOrElse {
+      cls.getDecl(name).getOrElse {
         throw new SymbolLookupException(name, s"in $cls")
       }
 
diff --git a/shared/src/main/scala/tastyquery/reader/classfiles/JavaSignatures.scala b/shared/src/main/scala/tastyquery/reader/classfiles/JavaSignatures.scala
index fd2f96be..66f7078a 100644
--- a/shared/src/main/scala/tastyquery/reader/classfiles/JavaSignatures.scala
+++ b/shared/src/main/scala/tastyquery/reader/classfiles/JavaSignatures.scala
@@ -141,14 +141,14 @@ object JavaSignatures:
         def packageSpecifiers(acc: PackageClassSymbol): ClassSymbol =
           val next = identifier
           if consume('/') then // must have '/', identifier, and terminal char.
-            acc.lookup(next) match
+            acc.getDecl(next) match
               case Some(pkg: PackageClassSymbol) =>
                 packageSpecifiers(pkg)
               case res =>
                 sys.error(s"found $res in operation $acc lookup $next")
                 abort
           else
-            acc.lookup(next.toTypeName) match // TODO: encoded names?
+            acc.getDecl(next.toTypeName) match // TODO: encoded names?
               case Some(cls: ClassSymbol) =>
                 cls
               case res =>
@@ -162,7 +162,7 @@ object JavaSignatures:
           packageSpecifiers(init.resolveToSymbol)
         else
           val emptyPkg = PackageRef(nme.EmptyPackageName).resolveToSymbol
-          emptyPkg.lookup(firstIdent.toTypeName) match
+          emptyPkg.getDecl(firstIdent.toTypeName) match
             case Some(cls: ClassSymbol) => cls // TODO: encoded names?
             case _                      => abort
       end findRawTopLevelClass
diff --git a/shared/src/main/scala/tastyquery/reader/pickles/PickleReader.scala b/shared/src/main/scala/tastyquery/reader/pickles/PickleReader.scala
index 2beb7611..bd03bb35 100644
--- a/shared/src/main/scala/tastyquery/reader/pickles/PickleReader.scala
+++ b/shared/src/main/scala/tastyquery/reader/pickles/PickleReader.scala
@@ -129,7 +129,7 @@ class PickleReader {
         case _ =>
           owner match
             case owner: PackageClassSymbol =>
-              owner.lookup(name).getOrElse {
+              owner.getDecl(name).getOrElse {
                 //errorBadSignature(s"cannot find symbol $owner.$name")
                 ExternalSymbolRef(owner, name)
               }
@@ -304,7 +304,7 @@ class PickleReader {
             // and also for the inner Transform class in all views. We fix it by
             // replacing the this with the appropriate super.
             if (sym.owner != thispre.cls) {
-              val overriding = thispre.cls.info.decls.lookup(sym.name)
+              val overriding = thispre.cls.info.decls.asClass.getDecl(sym.name)
               if (overriding.exists && overriding != sym) {
                 val base = pre.baseType(sym.owner)
                 assert(base.exists)


Test patch:
diff --git a/jvm/src/test/scala/tastyquery/LazyLoadingSuite.scala b/jvm/src/test/scala/tastyquery/LazyLoadingSuite.scala
index c6f3307a..060659c0 100644
--- a/jvm/src/test/scala/tastyquery/LazyLoadingSuite.scala
+++ b/jvm/src/test/scala/tastyquery/LazyLoadingSuite.scala
@@ -21,7 +21,7 @@ class LazyLoadingSuite extends RestrictedUnpicklingSuite {
     * They can only do that because they access `private[tastyquery]` methods,
     * in particular `Symbol.getDeclInternal`.
     */
-  object DangerousOps {
+  object DangerousOps:
     def followPathImpl(root: DeclaringSymbol, path: DeclarationPath): Either[String, Symbol] = {
       def follow(selected: Symbol)(remainder: DeclarationPath): Either[String, Symbol] = selected match {
         case nextDecl: DeclaringSymbol =>
@@ -36,18 +36,10 @@ class LazyLoadingSuite extends RestrictedUnpicklingSuite {
       yield sym
     }
 
-    private def select(root: DeclaringSymbol, next: Name): Either[String, Symbol] = {
-      val sel = (root, next) match {
-        case (p: PackageClassSymbol, s: SimpleName) if p.name != nme.RootName =>
-          p.name.toTermName.select(s)
-        case _ => next
-      }
-      root.getDeclInternal(sel) match {
+    private def select(root: DeclaringSymbol, next: Name): Either[String, Symbol] =
+      root.getDeclInternal(next) match
         case Some(someSym) => Right(someSym)
-        case _ => Left(s"No declaration for ${next.toDebugString} [${sel.toDebugString}] in ${root.toDebugString}")
-      }
-    }
-  }
+        case _             => Left(s"No declaration for ${next.toDebugString} in ${root.toDebugString}")
 
   def assertSymbolExistsAndIsLoaded(path: DeclarationPath)(implicit ctx: Context): Unit =
     DangerousOps.followPathImpl(defn.RootPackage, path).fold(fail(_), _ => ())
diff --git a/jvm/src/test/scala/tastyquery/SymbolSuite.scala b/jvm/src/test/scala/tastyquery/SymbolSuite.scala
index 7c080435..4f464f31 100644
--- a/jvm/src/test/scala/tastyquery/SymbolSuite.scala
+++ b/jvm/src/test/scala/tastyquery/SymbolSuite.scala
@@ -132,4 +132,23 @@ class SymbolSuite extends RestrictedUnpicklingSuite {
       Set.empty
     )
   }
+
+  test("nested-package-lookup") {
+    import tastyquery.ast.Types.*
+
+    val InNestedPackage = `simple_trees.nested` / tname"InNestedPackage"
+    given Context = getUnpicklingContext(InNestedPackage)
+
+    assert(resolve(InNestedPackage).fullName.toString == "simple_trees.nested.InNestedPackage")
+
+    val (simpleTreesPkg @ _: PackageClassSymbol) = resolve(simple_trees): @unchecked
+
+    assert(simpleTreesPkg.fullName.toString == "simple_trees")
+
+    val (simpleTreesNestedPkg @ _: PackageClassSymbol) = simpleTreesPkg.getDecl(name"nested").get: @unchecked
+
+    assert(simpleTreesNestedPkg.fullName.toString == "simple_trees.nested")
+
+    assert(TermRef(PackageRef(simpleTreesPkg), name"nested").symbol == simpleTreesNestedPkg)
+  }
 }
diff --git a/jvm/src/test/scala/tastyquery/TypeSuite.scala b/jvm/src/test/scala/tastyquery/TypeSuite.scala
index b49f87e3..67e13506 100644
--- a/jvm/src/test/scala/tastyquery/TypeSuite.scala
+++ b/jvm/src/test/scala/tastyquery/TypeSuite.scala
@@ -285,8 +285,7 @@ class TypeSuite extends UnrestrictedUnpicklingSuite {
   testWithContext("basic-scala2-types") {
     val ScalaRange = name"scala" / name"collection" / name"immutable" / tname"Range"
 
-    val rangeSym = resolve(ScalaRange)
-    assert(rangeSym.isClass, rangeSym)
+    val RangeClass = resolve(ScalaRange).asClass
 
     val BooleanClass = resolve(name"scala" / tname"Boolean")
     val IntClass = resolve(name"scala" / tname"Int")
@@ -294,33 +293,33 @@ class TypeSuite extends UnrestrictedUnpicklingSuite {
     val IndexedSeqClass = resolve(name"scala" / name"collection" / name"immutable" / tname"IndexedSeq")
 
     // val start: Int
-    val startSym = rangeSym.lookup(name"start").get
+    val startSym = RangeClass.getDecl(name"start").get
     assert(startSym.declaredType.isOfClass(IntClass), clue(startSym.declaredType))
     assert(startSym.declaredType.isInstanceOf[ExprType], clue(startSym.declaredType)) // should this be TypeRef?
 
     // val isEmpty: Boolean
-    val isEmptySym = rangeSym.lookup(name"isEmpty").get
+    val isEmptySym = RangeClass.getDecl(name"isEmpty").get
     assert(isEmptySym.declaredType.isOfClass(BooleanClass), clue(isEmptySym.declaredType))
     assert(isEmptySym.declaredType.isInstanceOf[ExprType], clue(isEmptySym.declaredType)) // should this be TypeRef?
 
     // def isInclusive: Boolean
-    val isInclusiveSym = rangeSym.lookup(name"isInclusive").get
+    val isInclusiveSym = RangeClass.getDecl(name"isInclusive").get
     assert(isInclusiveSym.declaredType.isOfClass(BooleanClass), clue(isInclusiveSym.declaredType))
     assert(isInclusiveSym.declaredType.isInstanceOf[ExprType], clue(isInclusiveSym.declaredType))
 
     // def by(step: Int): Range
     locally {
-      val bySym = rangeSym.lookup(name"by").get
+      val bySym = RangeClass.getDecl(name"by").get
       val mt = bySym.declaredType.asInstanceOf[MethodType]
       assertEquals(List[TermName](name"step"), mt.paramNames, clue(mt.paramNames))
       assert(mt.paramTypes.sizeIs == 1)
       assert(mt.paramTypes.head.isOfClass(IntClass), clue(mt.paramTypes.head))
-      assert(mt.resultType.isOfClass(rangeSym), clue(mt.resultType))
+      assert(mt.resultType.isOfClass(RangeClass), clue(mt.resultType))
     }
 
     // def map[B](f: Int => B): IndexedSeq[B]
     locally {
-      val mapSym = rangeSym.lookup(name"map").get
+      val mapSym = RangeClass.getDecl(name"map").get
       val pt = mapSym.declaredType.asInstanceOf[PolyType]
       val mt = pt.resultType.asInstanceOf[MethodType]
       assertEquals(List[TypeName](name"B".toTypeName), pt.paramNames, clue(pt.paramNames))
@@ -596,7 +595,7 @@ class TypeSuite extends UnrestrictedUnpicklingSuite {
     val PolySelect = resolve(name"simple_trees" / tname"PolySelect").asClass
     val IntClass = resolve(name"scala" / tname"Int")
 
-    val Some(DefDef(_, _, _, body, _)) = PolySelect.lookup(name"testField").get.tree: @unchecked
+    val Some(DefDef(_, _, _, body, _)) = PolySelect.getDecl(name"testField").get.tree: @unchecked
 
     val Select(qual, fieldName) = body: @unchecked
 
@@ -610,7 +609,7 @@ class TypeSuite extends UnrestrictedUnpicklingSuite {
     val PolySelect = resolve(name"simple_trees" / tname"PolySelect").asClass
     val IntClass = resolve(name"scala" / tname"Int")
 
-    val Some(DefDef(_, _, _, body, _)) = PolySelect.lookup(name"testGetter").get.tree: @unchecked
+    val Some(DefDef(_, _, _, body, _)) = PolySelect.getDecl(name"testGetter").get.tree: @unchecked
 
     val Select(qual, getterName) = body: @unchecked
 
@@ -624,7 +623,7 @@ class TypeSuite extends UnrestrictedUnpicklingSuite {
     val PolySelect = resolve(name"simple_trees" / tname"PolySelect").asClass
     val IntClass = resolve(name"scala" / tname"Int")
 
-    val Some(DefDef(_, _, _, body, _)) = PolySelect.lookup(name"testMethod").get.tree: @unchecked
+    val Some(DefDef(_, _, _, body, _)) = PolySelect.getDecl(name"testMethod").get.tree: @unchecked
 
     val Apply(fun @ Select(qual, methodName), List(arg)) = body: @unchecked
 
@@ -646,7 +645,7 @@ class TypeSuite extends UnrestrictedUnpicklingSuite {
     val PolySelect = resolve(name"simple_trees" / tname"PolySelect").asClass
     val IntClass = resolve(name"scala" / tname"Int")
 
-    val Some(DefDef(_, _, _, body, _)) = PolySelect.lookup(name"testGenericMethod").get.tree: @unchecked
+    val Some(DefDef(_, _, _, body, _)) = PolySelect.getDecl(name"testGenericMethod").get.tree: @unchecked
 
     val Apply(tapp @ TypeApply(fun @ Select(qual, methodName), List(targ)), List(arg)) = body: @unchecked