[GH] [Room][Compiler Processing] Support getting elements by package name

## Proposed Changes
KSP's recent version added the new API "Resolver.getDeclarationsFromPackage()" which is conceptually similar to javac's "Elements.getPackageElement". The Room compiler processing library can include an abstraction in XProcessingEnv to surface both of these and return XElements contained in a given package.

This PR adds the function `getElementsFromPackage` to XProcessingEnv, and has the implementations for KSP and Javac simply use their appropriate package lookup. This was very straightforward; the main source of code changes was to handle cases where xprocessing does not support certain element types, in which case those elements are skipped. I introduced a new exception to throw in this case, which callers can catch.

I'm personally partial to a "result" type monad to represent success and failure in a safer, clearer way, but Room doesn't seem to use that pattern.

## Testing

Test: Updated XRoundEnvTest.kt to test that elements from packages in both local sources and binaries are queried correctly.

## Issues Fixed

Fixes: https://issuetracker.google.com/issues/185609823

This is an imported pull request from https://github.com/androidx/androidx/pull/163.

Resolves #163
Github-Pr-Head-Sha: c7a7c1d3f156a7816f8547cfe879645756cdd976
GitOrigin-RevId: 0560adf51554b3e61fb8356199fd2e0d60647a32
Change-Id: I13c8fcba04d356323d55ade93d5a941357254037
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XProcessingEnv.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XProcessingEnv.kt
index da29394..7a860d6 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XProcessingEnv.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/XProcessingEnv.kt
@@ -151,4 +151,16 @@
             resolver = resolver
         )
     }
+
+    /**
+     * Returns [XTypeElement]s with the given package name. Note that this call can be expensive.
+     *
+     * @param packageName the package name to look up.
+     *
+     * @return A list of [XTypeElement] with matching package name. This will return declarations
+     * from both dependencies and source.
+     */
+    fun getTypeElementsFromPackage(packageName: String): List<XTypeElement>
+
+    // TODO: Add support for getting top level members in a package
 }
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacProcessingEnv.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacProcessingEnv.kt
index 9e8f7e4..e6ca777 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacProcessingEnv.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/javac/JavacProcessingEnv.kt
@@ -73,6 +73,16 @@
         return typeElementStore[qName]
     }
 
+    override fun getTypeElementsFromPackage(packageName: String): List<XTypeElement> {
+        // Note, to support Java Modules we would need to use "getAllPackageElements",
+        // but that is only available in Java 9+.
+        val packageElement = delegate.elementUtils.getPackageElement(packageName)
+
+        return packageElement.enclosedElements
+            .filterIsInstance<TypeElement>()
+            .map { wrapTypeElement(it) }
+    }
+
     override fun findType(qName: String): XType? {
         // check for primitives first
         PRIMITIVE_TYPES[qName]?.let {
diff --git a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt
index 499634e..46889e0 100644
--- a/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt
+++ b/room/compiler-processing/src/main/java/androidx/room/compiler/processing/ksp/KspProcessingEnv.kt
@@ -22,6 +22,7 @@
 import androidx.room.compiler.processing.XType
 import androidx.room.compiler.processing.XTypeElement
 import androidx.room.compiler.processing.javac.XTypeElementStore
+import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.getClassDeclarationByName
 import com.google.devtools.ksp.processing.CodeGenerator
 import com.google.devtools.ksp.processing.KSPLogger
@@ -82,6 +83,14 @@
         return typeElementStore[qName]
     }
 
+    @OptIn(KspExperimental::class)
+    override fun getTypeElementsFromPackage(packageName: String): List<XTypeElement> {
+        return resolver.getDeclarationsFromPackage(packageName)
+            .filterIsInstance<KSClassDeclaration>()
+            .map { KspTypeElement.create(this, it) }
+            .toList()
+    }
+
     override fun findType(qName: String): XType? {
         val kotlinTypeName = KspTypeMapper.swapWithKotlinType(qName)
         return resolver.findClass(kotlinTypeName)?.let {
diff --git a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XRoundEnvTest.kt b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XRoundEnvTest.kt
index 3b79246..1645db1 100644
--- a/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XRoundEnvTest.kt
+++ b/room/compiler-processing/src/test/java/androidx/room/compiler/processing/XRoundEnvTest.kt
@@ -216,6 +216,44 @@
         }
     }
 
+    @Test
+    fun getElementsFromPackageIncludesSources() {
+        val source = Source.kotlin(
+            "foo/Baz.kt",
+            """
+            package foo
+            class Baz 
+            """.trimIndent()
+        )
+
+        runProcessorTest(listOf(source)) { testInvocation ->
+            val elements = testInvocation.processingEnv.getTypeElementsFromPackage("foo")
+            val targetElement = testInvocation.processingEnv.requireTypeElement(
+                "foo.Baz"
+            )
+            assertThat(
+                elements
+            ).contains(targetElement)
+        }
+    }
+
+    @Test
+    fun getElementsFromPackageIncludesBinaries() {
+        runProcessorTest { testInvocation ->
+            val kspElements = testInvocation.processingEnv.getTypeElementsFromPackage(
+                "com.google.devtools.ksp.processing"
+            )
+
+            val symbolProcessorType = testInvocation.processingEnv.requireTypeElement(
+                "com.google.devtools.ksp.processing.SymbolProcessor"
+            )
+
+            assertThat(
+                kspElements
+            ).contains(symbolProcessorType)
+        }
+    }
+
     annotation class TopLevelAnnotation
 
     @Suppress("unused") // used in tests