[KAPT] Adhere to the call conventions when instantiating and using KtAnalysisSession

Merge-request: KT-MR-13098
Merged-by: Pavel Mikhailovskii <Pavel.Mikhailovskii@jetbrains.com>

build.gradle.kts

@@ -493,6 +493,7 @@ val projectsWithEnabledContextReceivers by extra {
         ":kotlinx-serialization-compiler-plugin.k2",
         ":plugins:parcelize:parcelize-compiler:parcelize.k2",
         ":plugins:fir-plugin-prototype",
+        ":plugins:kapt4",
     )
 }
 

plugins/kapt4/build.gradle.kts

@@ -74,3 +74,7 @@ publish()
 runtimeJar()
 sourcesJar()
 javadocJar()
+
+tasks.withType<org.jetbrains.kotlin.gradle.tasks.KotlinCompile>().configureEach {
+    kotlinOptions.freeCompilerArgs += "-Xcontext-receivers"
+}

plugins/kapt4/src/org/jetbrains/kotlin/kapt4/Kapt4AnalysisHandlerExtension.kt

@@ -5,13 +5,13 @@
 
 package org.jetbrains.kotlin.kapt4
 
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiJavaFile
 import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
-import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenProvider
-import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.api.standalone.KtAlwaysAccessibleLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.api.standalone.buildStandaloneAnalysisAPISession
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
@@ -30,7 +30,6 @@ import org.jetbrains.kotlin.kapt3.base.util.info
 import org.jetbrains.kotlin.kapt3.measureTimeMillis
 import org.jetbrains.kotlin.kapt3.util.MessageCollectorBackedKaptLogger
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
-import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 private class Kapt4AnalysisHandlerExtension : FirAnalysisHandlerExtension() {
@@ -73,7 +72,7 @@ private class Kapt4AnalysisHandlerExtension : FirAnalysisHandlerExtension() {
                 registerProjectService(KtLifetimeTokenProvider::class.java, KtAlwaysAccessibleLifetimeTokenProvider())
             }
 
-        val (module, psiFiles) = standaloneAnalysisAPISession.modulesWithFiles.entries.single()
+        val (module, files) = standaloneAnalysisAPISession.modulesWithFiles.entries.single()
 
         optionsBuilder.apply {
             projectBaseDir = projectBaseDir ?: module.project.basePath?.let(::File)
@@ -91,16 +90,15 @@ private class Kapt4AnalysisHandlerExtension : FirAnalysisHandlerExtension() {
 
         return try {
             if (options.mode.generateStubs) {
-                KtAnalysisSessionProvider.getInstance(module.project).analyze(module) {
-                    generateAndSaveStubs(
-                        psiFiles.filterIsInstance<KtFile>(),
-                        options,
-                        logger,
-                        configuration.getBoolean(CommonConfigurationKeys.REPORT_OUTPUT_FILES),
-                        this,
-                        configuration[CommonConfigurationKeys.METADATA_VERSION]
-                    )
-                }
+                generateAndSaveStubs(
+                    module,
+                    files,
+                    options,
+                    logger,
+                    configuration.getBoolean(CommonConfigurationKeys.REPORT_OUTPUT_FILES),
+                    configuration[CommonConfigurationKeys.METADATA_VERSION]
+                )
+
             }
             if (options.mode.runAnnotationProcessing) {
                 KaptContext(
@@ -120,11 +118,11 @@ private class Kapt4AnalysisHandlerExtension : FirAnalysisHandlerExtension() {
     }
 
     private fun generateAndSaveStubs(
-        files: List<KtFile>,
+        module: KtSourceModule,
+        files: List<PsiFile>,
         options: KaptOptions,
         logger: MessageCollectorBackedKaptLogger,
         reportOutputFiles: Boolean,
-        analysisSession: KtAnalysisSession,
         overriddenMetadataVersion: BinaryVersion?
     ) {
         fun onError(message: String) {
@@ -135,7 +133,7 @@ private class Kapt4AnalysisHandlerExtension : FirAnalysisHandlerExtension() {
             }
         }
         val (stubGenerationTime, classesToStubs) = measureTimeMillis {
-            generateStubs(files, options, ::onError, analysisSession, overriddenMetadataVersion)
+            generateStubs(module, files, options, ::onError, overriddenMetadataVersion)
         }
 
         logger.info { "Java stub generation took $stubGenerationTime ms" }

plugins/kapt4/src/org/jetbrains/kotlin/kapt4/StubGenerator.kt

@@ -13,10 +13,12 @@ import com.intellij.lang.ASTNode
 import com.intellij.psi.*
 import com.intellij.psi.impl.source.tree.LeafPsiElement
 import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
 import org.jetbrains.kotlin.asJava.elements.*
@@ -47,14 +49,16 @@ import org.jetbrains.kotlin.utils.Printer
 import java.io.File
 
 internal fun generateStubs(
-    files: List<KtFile>,
+    module: KtSourceModule,
+    files: List<PsiFile>,
     options: KaptOptions,
     onError: (messages: String) -> Unit,
-    analysisSession: KtAnalysisSession,
     overriddenMetadataVersion: BinaryVersion? = null,
     metadataRenderer: (Printer.(Metadata) -> Unit)? = null
 ): Map<KtLightClass, KaptStub?> =
-    StubGenerator(files, options, onError, analysisSession, metadataRenderer, overriddenMetadataVersion).generateStubs()
+    analyze(module) {
+        StubGenerator(files.filterIsInstance<KtFile>(), options, onError, metadataRenderer, overriddenMetadataVersion).generateStubs()
+    }
 
 class KaptStub(val source: String, val kaptMetadata: ByteArray) {
     fun writeMetadata(forSource: File) {
@@ -63,11 +67,11 @@ class KaptStub(val source: String, val kaptMetadata: ByteArray) {
     }
 }
 
+context(KtAnalysisSession)
 private class StubGenerator(
     private val files: List<KtFile>,
     options: KaptOptions,
     private val onError: (String) -> Unit,
-    private val analysisSession: KtAnalysisSession,
     private val metadataRenderer: (Printer.(Metadata) -> Unit)? = null,
     private val overriddenMetadataVersion: BinaryVersion? = null,
 ) {
@@ -160,14 +164,12 @@ private class StubGenerator(
 
                 if (!acceptableByName) continue
 
-                val importedSymbols = with(analysisSession) {
-                    val importedReference = importDirective.importedReference
-                        ?.getCalleeExpressionIfAny()
-                        ?.references
-                        ?.firstOrNull() as? KtReference
-                    importedReference?.resolveToSymbols().orEmpty()
-                }
 
+                val importedReference = importDirective.importedReference
+                    ?.getCalleeExpressionIfAny()
+                    ?.references
+                    ?.firstOrNull() as? KtReference
+                val importedSymbols = importedReference?.resolveToSymbols().orEmpty()
                 val isAllUnderClassifierImport = importDirective.isAllUnder && importedSymbols.any { it is KtClassOrObjectSymbol }
                 val isCallableImport = !importDirective.isAllUnder && importedSymbols.any { it is KtCallableSymbol }
                 val isEnumEntryImport = !importDirective.isAllUnder && importedSymbols.any { it is KtEnumEntrySymbol }

plugins/kapt4/test/org/jetbrains/kotlin/kapt4/Kapt4Facade.kt

@@ -11,7 +11,6 @@ import kotlinx.metadata.jvm.KotlinClassMetadata
 import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.api.lifetime.KtReadActionConfinementLifetimeTokenProvider
-import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.api.standalone.buildStandaloneAnalysisAPISession
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoots
@@ -25,7 +24,6 @@ import org.jetbrains.kotlin.kapt3.test.KaptMessageCollectorProvider
 import org.jetbrains.kotlin.kapt3.test.kaptOptionsProvider
 import org.jetbrains.kotlin.kotlinp.Kotlinp
 import org.jetbrains.kotlin.kotlinp.KotlinpSettings
-import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.model.*
 import org.jetbrains.kotlin.test.services.*
 import org.jetbrains.kotlin.utils.Printer
@@ -80,23 +78,20 @@ private fun run(
         @Suppress("DEPRECATION")
         buildKtModuleProviderByCompilerConfiguration(configuration)
     }
-    val (module, psiFiles) = standaloneAnalysisAPISession.modulesWithFiles.entries.single()
-
-    return KtAnalysisSessionProvider.getInstance(module.project).analyze(module) {
-        val context = KaptContext(
-            options,
-            withJdk = false,
-            WriterBackedKaptLogger(isVerbose = false),
-        )
-        val onError = { message: String ->
-            if (context.options[KaptFlag.STRICT]) {
-                context.reportKaptError(*message.split("\n").toTypedArray())
-            } else {
-                context.logger.warn(message)
-            }
+    val (module, files) = standaloneAnalysisAPISession.modulesWithFiles.entries.single()
+    val context = KaptContext(
+        options,
+        withJdk = false,
+        WriterBackedKaptLogger(isVerbose = false),
+    )
+    val onError = { message: String ->
+        if (context.options[KaptFlag.STRICT]) {
+            context.reportKaptError(*message.split("\n").toTypedArray())
+        } else {
+            context.logger.warn(message)
         }
-        context to generateStubs(psiFiles.filterIsInstance<KtFile>(), options, onError,this@analyze, metadataRenderer = { renderMetadata(it) })
     }
+    return context to generateStubs(module, files, options, onError, metadataRenderer = { renderMetadata(it) })
 }
 
 internal data class Kapt4ContextBinaryArtifact(