Implement Resolver.getSubpackagesOf()

api/api.base

@@ -176,6 +176,7 @@ package com.google.devtools.ksp.processing {
     method @NonNull @com.google.devtools.ksp.KspExperimental public kotlin.sequences.Sequence<com.google.devtools.ksp.symbol.KSAnnotation> getPackageAnnotations(@NonNull String packageName);
     method @NonNull @com.google.devtools.ksp.KspExperimental public kotlin.sequences.Sequence<java.lang.String> getPackagesWithAnnotation(@NonNull String annotationName);
     method @Nullable public com.google.devtools.ksp.symbol.KSPropertyDeclaration getPropertyDeclarationByName(@NonNull com.google.devtools.ksp.symbol.KSName name, boolean includeTopLevel = false);
+    method @NonNull @com.google.devtools.ksp.KspExperimental public kotlin.sequences.Sequence<java.lang.String> getSubpackagesOf(@NonNull String packageName);
     method @NonNull public kotlin.sequences.Sequence<com.google.devtools.ksp.symbol.KSAnnotated> getSymbolsWithAnnotation(@NonNull String annotationName, boolean inDepth = false);
     method @NonNull public com.google.devtools.ksp.symbol.KSTypeArgument getTypeArgument(@NonNull com.google.devtools.ksp.symbol.KSTypeReference typeRef, @NonNull com.google.devtools.ksp.symbol.Variance variance);
     method @com.google.devtools.ksp.KspExperimental public boolean isJavaRawType(@NonNull com.google.devtools.ksp.symbol.KSType type);

api/src/main/kotlin/com/google/devtools/ksp/processing/Resolver.kt

@@ -238,6 +238,15 @@ interface Resolver {
     @KspExperimental
     fun getDeclarationsFromPackage(packageName: String): Sequence<KSDeclaration>
 
+    /**
+     * Returns a sequence of subpackage names of the given [packageName].
+     * @param packageName the root package name to search from.
+     * @return A sequence of subpackage names, if any. Does not include the input package name.
+     * This will return declarations from both dependencies and source.
+     */
+    @KspExperimental
+    fun getSubpackagesOf(packageName: String): Sequence<String>
+
     /**
      * Returns the corresponding Kotlin class with the given Java class.
      *

compiler-plugin/src/main/kotlin/com/google/devtools/ksp/processing/impl/ResolverImpl.kt

@@ -95,6 +95,7 @@ import org.jetbrains.kotlin.resolve.constants.ConstantValue
 import org.jetbrains.kotlin.resolve.constants.KClassValue
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.descriptorUtil.classId
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameOrNull
 import org.jetbrains.kotlin.resolve.descriptorUtil.getAllSuperClassifiers
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.descriptorUtil.propertyIfAccessor
@@ -1013,6 +1014,25 @@ class ResolverImpl(
             .plus(javaPackageToClassMap.getOrDefault(packageName, emptyList()).asSequence())
     }
 
+    @KspExperimental
+    override fun getSubpackagesOf(packageName: String): Sequence<String> {
+        fun PackageViewDescriptor.subPackages(): Sequence<PackageViewDescriptor> = memberScope
+            .getContributedDescriptors(DescriptorKindFilter.PACKAGES)
+            .asSequence()
+            .filterIsInstance<PackageViewDescriptor>()
+            .filterNot { it == this@subPackages }
+
+        return generateSequence(listOf(module.getPackage(FqName(packageName)))) { subPackages ->
+            subPackages
+                .flatMap { it.subPackages() }
+                .ifEmpty { null }
+        }
+            .flatMap { it.asSequence() }
+            // Drop the input package
+            .drop(1)
+            .mapNotNull { it.fqNameOrNull()?.asString() }
+    }
+
     override fun getTypeArgument(typeRef: KSTypeReference, variance: Variance): KSTypeArgument {
         return KSTypeArgumentLiteImpl.getCached(typeRef, variance)
     }

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

@@ -55,6 +55,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtFunctionLikeSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtJavaFieldSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtNamedClassOrObjectSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtPackageSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPropertyAccessorSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtPropertySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
@@ -346,6 +347,26 @@ class ResolverAAImpl(
         }
     }
 
+    @KspExperimental
+    override fun getSubpackagesOf(packageName: String): Sequence<String> {
+        return analyze {
+            fun KtPackageSymbol.subPackages(): Sequence<KtPackageSymbol> = getPackageScope()
+                .getPackageSymbols()
+                .distinct()
+
+            val packages = listOfNotNull(analysisSession.getPackageSymbolIfPackageExists(FqName(packageName)))
+            generateSequence(packages) { subPackages ->
+                subPackages
+                    .flatMap { it.subPackages() }
+                    .ifEmpty { null }
+            }
+                .flatMap { it.asSequence() }
+                // Drop the input package
+                .drop(1)
+                .map { it.fqName.asString() }
+        }
+    }
+
     override fun getDeclarationsInSourceOrder(container: KSDeclarationContainer): Sequence<KSDeclaration> {
         if (container.origin != Origin.KOTLIN_LIB) {
             return container.declarations

kotlin-analysis-api/testData/getSubPackages.kt

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2021 Google LLC
+ * Copyright 2010-2021 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: GetSubPackagesProcessor
+// EXPECTED:
+// subpackages of lib1
+// subpackages of lib2
+// subpackages of main
+// main.test
+// main.test.nested
+// main.test.main
+// main.test.main.test
+// main.test.main.test.nested
+// subpackages of non.exist
+// END
+
+// MODULE: lib1
+// FILE: foo.kt
+package lib1
+
+class Foo
+
+fun funcFoo(): Int {
+    return 1
+}
+
+// FILE: Bar.java
+package lib1;
+
+class Bar {}
+
+// MODULE: lib2
+// FILE: foo.kt
+package lib2
+
+class Foo
+
+typealias FooTypeAlias = Foo
+
+val a = 0
+
+// FILE: Bar.java
+
+class Bar {}
+
+// MODULE: main(lib1, lib2)
+// FILE: a.kt
+package lib1
+class FooInSource
+
+val propInSource = 1
+// FILE: main.kt
+package main.test
+
+class KotlinMain
+
+// FILE: main/test/C.java
+package main.test;
+
+public class C {
+
+}
+
+class D {
+
+}
+
+// FILE: wrongDir/K.java
+package main;
+
+public class K {
+
+}
+
+class KK {}
+
+
+// FILE: main/test/main/test/L.java
+package main.test;
+
+public class L {
+
+}
+
+
+// FILE: main/test/main/test/nested/M.java
+package main.test.nested;
+
+public class M {
+
+}

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

@@ -0,0 +1,28 @@
+package com.google.devtools.ksp.processor
+
+import com.google.devtools.ksp.KspExperimental
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.symbol.KSAnnotated
+
+class GetSubPackagesProcessor : AbstractTestProcessor() {
+    val results = mutableListOf<String>()
+
+    override fun toResult(): List<String> {
+        return results
+    }
+
+    @OptIn(KspExperimental::class)
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        addPackage("lib1", resolver)
+        addPackage("lib2", resolver)
+        addPackage("main", resolver)
+        addPackage("non.exist", resolver)
+        return emptyList()
+    }
+
+    @KspExperimental
+    private fun addPackage(name: String, resolver: Resolver) {
+        results.add("subpackages of $name")
+        results += resolver.getSubpackagesOf(name).toList()
+    }
+}

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

@@ -265,6 +265,12 @@ class KSPAATest : AbstractKSPAATest() {
         runTest("../kotlin-analysis-api/testData/getPackage.kt")
     }
 
+    @TestMetadata("getSubPackages.kt")
+    @Test
+    fun testGetSubPackages() {
+        runTest("../kotlin-analysis-api/testData/getSubPackages.kt")
+    }
+
     @TestMetadata("getByName.kt")
     @Test
     fun testGetByName() {

test-utils/src/test/kotlin/com/google/devtools/ksp/test/KSPCompilerPluginTest.kt

@@ -255,6 +255,13 @@ class KSPCompilerPluginTest : AbstractKSPCompilerPluginTest() {
         runTest("../test-utils/testData/api/getPackage.kt")
     }
 
+    @DisabledOnOs(OS.WINDOWS)
+    @TestMetadata("getSubPackages.kt")
+    @Test
+    fun testGetSubPackages() {
+        runTest("../test-utils/testData/api/getSubPackages.kt")
+    }
+
     @TestMetadata("getByName.kt")
     @Test
     fun testGetByName() {

test-utils/testData/api/getSubPackages.kt

@@ -0,0 +1,104 @@
+/*
+ * Copyright 2021 Google LLC
+ * Copyright 2010-2021 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: GetSubPackagesProcessor
+// EXPECTED:
+// subpackages of lib1
+// subpackages of lib2
+// subpackages of main
+// main.test
+// main.test.main
+// main.test.main.test
+// main.test.main.test.nested
+// subpackages of non.exist
+// END
+
+// MODULE: lib1
+// FILE: foo.kt
+package lib1
+
+class Foo
+
+fun funcFoo(): Int {
+    return 1
+}
+
+// FILE: Bar.java
+package lib1;
+
+class Bar {}
+
+// MODULE: lib2
+// FILE: foo.kt
+package lib2
+
+class Foo
+
+typealias FooTypeAlias = Foo
+
+val a = 0
+
+// FILE: Bar.java
+
+class Bar {}
+
+// MODULE: main(lib1, lib2)
+// FILE: a.kt
+package lib1
+class FooInSource
+
+val propInSource = 1
+// FILE: main.kt
+package main.test
+
+class KotlinMain
+
+// FILE: main/test/C.java
+package main.test;
+
+public class C {
+
+}
+
+class D {
+
+}
+
+// FILE: wrongDir/K.java
+package main;
+
+public class K {
+
+}
+
+class KK {}
+
+
+// FILE: main/test/main/test/L.java
+package main.test;
+
+public class L {
+
+}
+
+
+// FILE: main/test/main/test/nested/M.java
+package main.test.nested;
+
+public class M {
+
+}