KSP2: reimplement Resolver.getDeclarationsFromPackage

(cherry picked from commit bca0c74)

Author: ting-yuan

kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt

@@ -279,25 +279,16 @@ class ResolverAAImpl(
     @KspExperimental
     override fun getDeclarationsFromPackage(packageName: String): Sequence<KSDeclaration> {
         return analyze {
-            val packageNames = FqName(packageName).pathSegments().map { it.asString() }
-            var packages = listOf(useSiteSession.rootPackageSymbol)
-            for (curName in packageNames) {
-                packages = packages
-                    .flatMap { it.packageScope.getPackageSymbols { it.asString() == curName } }
-                    .distinct()
-            }
-            packages.flatMap {
-                it.packageScope.declarations.distinct().mapNotNull { symbol ->
-                    when (symbol) {
-                        is KaNamedClassSymbol -> KSClassDeclarationImpl.getCached(symbol)
-                        is KaFunctionSymbol -> KSFunctionDeclarationImpl.getCached(symbol)
-                        is KaPropertySymbol -> KSPropertyDeclarationImpl.getCached(symbol)
-                        is KaTypeAliasSymbol -> KSTypeAliasImpl.getCached(symbol)
-                        is KaJavaFieldSymbol -> KSPropertyDeclarationJavaImpl.getCached(symbol)
-                        else -> null
-                    }
+            findPackage(FqName(packageName))?.packageScope?.declarations?.distinct()?.mapNotNull { symbol ->
+                when (symbol) {
+                    is KaNamedClassSymbol -> KSClassDeclarationImpl.getCached(symbol)
+                    is KaFunctionSymbol -> KSFunctionDeclarationImpl.getCached(symbol)
+                    is KaPropertySymbol -> KSPropertyDeclarationImpl.getCached(symbol)
+                    is KaTypeAliasSymbol -> KSTypeAliasImpl.getCached(symbol)
+                    is KaJavaFieldSymbol -> KSPropertyDeclarationJavaImpl.getCached(symbol)
+                    else -> null
                 }
-            }.asSequence()
+            }?.asSequence() ?: emptySequence()
         }
     }
 

kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/test/KSPAATest.kt

@@ -714,4 +714,10 @@ class KSPAATest : AbstractKSPAATest() {
     fun testExitCode() {
         runTest("../test-utils/testData/api/exitCode.kt")
     }
+
+    @TestMetadata("packageProviderForGenerated.kt")
+    @Test
+    fun testPackageProviderForGenerated() {
+        runTest("../test-utils/testData/api/packageProviderForGenerated.kt")
+    }
 }

test-utils/src/main/kotlin/com/google/devtools/ksp/processor/PackageProviderForGeneratedProcessor.kt

@@ -0,0 +1,47 @@
+package com.google.devtools.ksp.processor
+
+import com.google.devtools.ksp.KspExperimental
+import com.google.devtools.ksp.processing.Dependencies
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.processing.SymbolProcessor
+import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
+import com.google.devtools.ksp.symbol.KSAnnotated
+
+class PackageProviderForGeneratedProcessor : AbstractTestProcessor() {
+    val result = mutableListOf<String>()
+    override fun toResult(): List<String> {
+        return result
+    }
+
+    var round = 0
+
+    @OptIn(KspExperimental::class)
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        if (round == 0) {
+            createJavaFile("foo.bar", "MyGeneratedJavaClass")
+        }
+        result.add("Processing round $round")
+        val getByName = resolver.getClassDeclarationByName(
+            resolver.getKSNameFromString("foo.bar.MyGeneratedJavaClass")
+        )?.qualifiedName?.asString()
+        result.add("resolver.getClassDeclarationByName: $getByName")
+        val getByPackage = resolver.getDeclarationsFromPackage("foo.bar").toList().map { it.qualifiedName?.asString() }
+        result.add("resolver.getDeclarationsFromPackage: $getByPackage")
+        round++
+        return emptyList()
+    }
+
+    lateinit var env: SymbolProcessorEnvironment
+
+    override fun create(environment: SymbolProcessorEnvironment): SymbolProcessor {
+        env = environment
+        return this
+    }
+
+    private fun createJavaFile(packageName: String, className: String) {
+        val dependencies = Dependencies(aggregating = false, sources = arrayOf())
+        env.codeGenerator.createNewFile(dependencies, packageName, className, "java").use {
+            it.write("package foo.bar;\nclass MyGeneratedJavaClass{}".toByteArray())
+        }
+    }
+}

test-utils/testData/api/packageProviderForGenerated.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2025 Google LLC
+ * Copyright 2010-2025 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// TEST PROCESSOR: PackageProviderForGeneratedProcessor
+// EXPECTED:
+// Processing round 0
+// resolver.getClassDeclarationByName: null
+// resolver.getDeclarationsFromPackage: []
+// Processing round 1
+// resolver.getClassDeclarationByName: foo.bar.MyGeneratedJavaClass
+// resolver.getDeclarationsFromPackage: [foo.bar.MyGeneratedJavaClass]
+// END
+
+// FILE: Dummy.kt
+class Dummy