Implement Resolver.getModuleName API

Resolves #1621
google · Apr 18, 2024 · c262051 · c262051
1 parent cf3b401
commit c262051
Show file tree

Hide file tree

Showing 8 changed files with 40 additions and 8 deletions.
diff --git a/api/api.base b/api/api.base
@@ -170,6 +170,7 @@ package com.google.devtools.ksp.processing {
     method @Nullable @com.google.devtools.ksp.KspExperimental public String getJvmName(@NonNull com.google.devtools.ksp.symbol.KSFunctionDeclaration declaration);
     method @Nullable @com.google.devtools.ksp.KspExperimental public String getJvmName(@NonNull com.google.devtools.ksp.symbol.KSPropertyAccessor accessor);
     method @NonNull public com.google.devtools.ksp.symbol.KSName getKSNameFromString(@NonNull String name);
+    method @NonNull @com.google.devtools.ksp.KspExperimental public com.google.devtools.ksp.symbol.KSName getModuleName();
     method @NonNull public kotlin.sequences.Sequence<com.google.devtools.ksp.symbol.KSFile> getNewFiles();
     method @Nullable @com.google.devtools.ksp.KspExperimental public String getOwnerJvmClassName(@NonNull com.google.devtools.ksp.symbol.KSFunctionDeclaration declaration);
     method @Nullable @com.google.devtools.ksp.KspExperimental public String getOwnerJvmClassName(@NonNull com.google.devtools.ksp.symbol.KSPropertyDeclaration declaration);

diff --git a/api/src/main/kotlin/com/google/devtools/ksp/processing/Resolver.kt b/api/src/main/kotlin/com/google/devtools/ksp/processing/Resolver.kt
@@ -310,12 +310,18 @@ interface Resolver {
     @KspExperimental
     fun getPackageAnnotations(packageName: String): Sequence<KSAnnotation>
 
-    @KspExperimental
     /**
      * Returns name of packages with given annotation.
      *
      * @param annotationName name of the annotation to be queried.
      * @return a sequence of package names with corresponding annotation name.
      */
+    @KspExperimental
     fun getPackagesWithAnnotation(annotationName: String): Sequence<String>
+
+    /**
+     * @return the name of the kotlin module this resolver is running on.
+     */
+    @KspExperimental
+    fun getModuleName(): KSName
 }
diff --git a/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/processing/impl/ResolverImpl.kt b/compiler-plugin/src/main/kotlin/com/google/devtools/ksp/processing/impl/ResolverImpl.kt
@@ -140,9 +140,10 @@ class ResolverImpl(
     private val functionAsMemberOfCache: MutableMap<Pair<KSFunctionDeclaration, KSType>, KSFunction>
     private val propertyAsMemberOfCache: MutableMap<Pair<KSPropertyDeclaration, KSType>, KSType>
 
+    private val moduleIdentifier = module.name.getNonSpecialIdentifier()
     private val typeMapper = KotlinTypeMapper(
         BindingContext.EMPTY, ClassBuilderMode.LIGHT_CLASSES,
-        module.name.getNonSpecialIdentifier(),
+        moduleIdentifier,
         KotlinTypeMapper.LANGUAGE_VERSION_SETTINGS_DEFAULT, // TODO use proper LanguageVersionSettings
         true
     )
@@ -1407,6 +1408,9 @@ class ResolverImpl(
         }.map { it.packageName.asString() }
     }
 
+    @KspExperimental
+    override fun getModuleName(): KSName = KSNameImpl.getCached(moduleIdentifier)
+
     private val psiJavaFiles = allKSFiles.filterIsInstance<KSFileJavaImpl>().map {
         Pair(it.psi.virtualFile.path, it.psi)
     }.toMap()

diff --git a/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt b/gradle-plugin/src/main/kotlin/com/google/devtools/ksp/gradle/KspAATask.kt
@@ -163,7 +163,7 @@ abstract class KspAATask @Inject constructor(
                 kspAATask.kspClasspath.from(kspAADepCfg)
                 kspAATask.kspConfig.let { cfg ->
                     cfg.processorClasspath.from(processorClasspath)
-                    cfg.moduleName.value(kotlinCompilation.defaultSourceSet.name)
+                    cfg.moduleName.value(project.name)
                     val kotlinOutputDir = KspGradleSubplugin.getKspKotlinOutputDir(project, sourceSetName, target)
                     val javaOutputDir = KspGradleSubplugin.getKspJavaOutputDir(project, sourceSetName, target)
                     val filteredTasks =
@@ -256,7 +256,7 @@ abstract class KspAATask @Inject constructor(
                         cfg.jvmDefaultMode.value(jvmDefaultMode)
 
                         val jvmTarget = project.provider {
-                            (compilerOptions as KotlinJvmCompilerOptions).jvmTarget.get().target
+                            compilerOptions.jvmTarget.get().target
                         }
                         cfg.jvmTarget.value(jvmTarget)
                     }

diff --git a/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/PlaygroundIT.kt b/integration-tests/src/test/kotlin/com/google/devtools/ksp/test/PlaygroundIT.kt
@@ -381,6 +381,15 @@ class PlaygroundIT(val useKSP2: Boolean) {
         project.restore(buildFile.path)
     }
 
+    @Test
+    fun testModuleName() {
+        File(project.root, "workload/build.gradle.kts").createNewFile()
+        val gradleRunner = GradleRunner.create().withProjectDir(project.root)
+        gradleRunner.withArguments("build").build().let { result ->
+            Assert.assertTrue(result.output.contains("Module name is workload"))
+        }
+    }
+
     companion object {
         @JvmStatic
         @Parameterized.Parameters(name = "KSP2={0}")

diff --git a/...ation-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt b/...ation-tests/src/test/resources/playground/test-processor/src/main/kotlin/TestProcessor.kt
@@ -1,3 +1,4 @@
+import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.containingFile
 import com.google.devtools.ksp.processing.*
 import com.google.devtools.ksp.symbol.*
@@ -6,6 +7,7 @@ import java.io.OutputStream
 class TestProcessor : SymbolProcessor {
     lateinit var codeGenerator: CodeGenerator
     lateinit var file: OutputStream
+    lateinit var logger: KSPLogger
     var invoked = false
 
     fun emit(s: String, indent: String) {
@@ -18,6 +20,7 @@ class TestProcessor : SymbolProcessor {
         codeGenerator: CodeGenerator,
         logger: KSPLogger
     ) {
+        this.logger = logger
         logger.warn("This is a harmless warning.")
         this.codeGenerator = codeGenerator
         file = codeGenerator.createNewFile(Dependencies(false), "", "TestProcessor", "log")
@@ -27,7 +30,9 @@ class TestProcessor : SymbolProcessor {
         javaFile.appendText("class Generated {}")
     }
 
+    @OptIn(KspExperimental::class)
     override fun process(resolver: Resolver): List<KSAnnotated> {
+        logger.warn("Module name is ${resolver.getModuleName().asString()}")
         if (invoked) {
             return emptyList()
         }

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
@@ -74,6 +74,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirLibrarySymbolProviderFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLSealedInheritorsProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.builder.KtModuleBuilder
 import org.jetbrains.kotlin.analysis.project.structure.builder.KtModuleProviderBuilder
 import org.jetbrains.kotlin.analysis.project.structure.builder.buildKtSdkModule
@@ -448,7 +449,8 @@ class KotlinSymbolProcessing(
 
         val psiManager = PsiManager.getInstance(project)
         val providers: List<SymbolProcessorProvider> = symbolProcessorProviders
-        ResolverAAImpl.ktModule = modules.single()
+        // KspModuleBuilder ensures this is always a KtSourceModule
+        ResolverAAImpl.ktModule = modules.single() as KtSourceModule
 
         // Initializing environments
         val allKSFiles = prepareAllKSFiles(kotlinCoreProjectEnvironment, modules, compilerConfiguration)

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/ResolverAAImpl.kt
@@ -62,7 +62,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtTypeAliasSymbol
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.decompiler.stub.file.ClsKotlinBinaryClassCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCliJavaFileManagerImpl
 import org.jetbrains.kotlin.fir.types.isRaw
@@ -88,13 +88,13 @@ class ResolverAAImpl(
 ) : Resolver {
     companion object {
         val instance_prop: ThreadLocal<ResolverAAImpl> = ThreadLocal()
-        private val ktModule_prop: ThreadLocal<KtModule> = ThreadLocal()
+        private val ktModule_prop: ThreadLocal<KtSourceModule> = ThreadLocal()
         var instance
             get() = instance_prop.get()
             set(value) {
                 instance_prop.set(value)
             }
-        var ktModule: KtModule
+        var ktModule: KtSourceModule
             get() = ktModule_prop.get()
             set(value) {
                 ktModule_prop.set(value)
@@ -621,6 +621,11 @@ class ResolverAAImpl(
         }.map { it.packageName.asString() }
     }
 
+    @KspExperimental
+    override fun getModuleName(): KSName {
+        return KSNameImpl.getCached((ktModule.stableModuleName ?: ktModule.moduleName).removeSurrounding("<", ">"))
+    }
+
     @KspExperimental
     override fun mapJavaNameToKotlin(javaName: KSName): KSName? {
         return JavaToKotlinClassMap.mapJavaToKotlin(FqName(javaName.asString()))?.toKSName()