Add AnalysisHandlerExtension extension point for K2JsCompiler

AnalysisHandlerExtension allows compiler plugins to:
1. Intercept and override the default analysis.
2. Utilize the compiler infrastructure to do custom analysis.

A well know plugin on the JVM platform is KAPT.

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JSCompiler.java

@@ -252,17 +252,28 @@ public class K2JSCompiler extends CLICompiler<K2JSCompilerArguments> {
             return COMPILATION_ERROR;
         }
 
-        AnalyzerWithCompilerReport analyzerWithCompilerReport = new AnalyzerWithCompilerReport(
+        AnalysisResult analysisResult;
-                messageCollector, CommonConfigurationKeysKt.getLanguageVersionSettings(configuration)
+        do {
-        );
+            AnalyzerWithCompilerReport analyzerWithCompilerReport = new AnalyzerWithCompilerReport(
-        analyzerWithCompilerReport.analyzeAndReport(sourcesFiles, () -> TopDownAnalyzerFacadeForJS.analyzeFiles(sourcesFiles, config));
+                    messageCollector, CommonConfigurationKeysKt.getLanguageVersionSettings(configuration)
-        if (analyzerWithCompilerReport.hasErrors()) {
+            );
-            return COMPILATION_ERROR;
+            List<KtFile> sources = environmentForJS.getSourceFiles();
-        }
+            analyzerWithCompilerReport.analyzeAndReport(sourcesFiles, () -> TopDownAnalyzerFacadeForJS.analyzeFiles(sources, config));
+            if (analyzerWithCompilerReport.hasErrors()) {
+                return COMPILATION_ERROR;
+            }
 
-        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled();
+            ProgressIndicatorAndCompilationCanceledStatus.checkCanceled();
+            analysisResult = analyzerWithCompilerReport.getAnalysisResult();
+
+            if (analysisResult instanceof JsAnalysisResult.RetryWithAdditionalRoots) {
+                environmentForJS.addKotlinSourceRoots(((JsAnalysisResult.RetryWithAdditionalRoots) analysisResult).getAdditionalKotlinRoots());
+            }
+        } while(analysisResult instanceof JsAnalysisResult.RetryWithAdditionalRoots);
+
+        if (!analysisResult.getShouldGenerateCode())
+            return OK;
 
-        AnalysisResult analysisResult = analyzerWithCompilerReport.getAnalysisResult();
         assert analysisResult instanceof JsAnalysisResult : "analysisResult should be instance of JsAnalysisResult, but " + analysisResult;
         JsAnalysisResult jsAnalysisResult = (JsAnalysisResult) analysisResult;
 

compiler/tests-common/tests/org/jetbrains/kotlin/checkers/AbstractDiagnosticsTestWithJsStdLib.kt

@@ -60,7 +60,7 @@ abstract class AbstractDiagnosticsTestWithJsStdLib : AbstractDiagnosticsTest() {
         moduleTrace.record<ModuleDescriptor, ModuleKind>(MODULE_KIND, moduleContext.module, getModuleKind(files))
         config.configuration.languageVersionSettings = languageVersionSettings
         return TopDownAnalyzerFacadeForJS.analyzeFilesWithGivenTrace(
-            files, moduleTrace, moduleContext, config.configuration, CompilerEnvironment,
+            files, moduleTrace, moduleContext, config.configuration, CompilerEnvironment, config.project
         )
     }
 

compiler/tests/org/jetbrains/kotlin/serialization/js/JsVersionRequirementTest.kt

@@ -48,7 +48,7 @@ class JsVersionRequirementTest : AbstractVersionRequirementTest() {
         }
         val trace = BindingTraceContext()
         val analysisResult = TopDownAnalyzerFacadeForJS.analyzeFilesWithGivenTrace(
-            ktFiles, trace, createModule(environment), environment.configuration, CompilerEnvironment,
+            ktFiles, trace, createModule(environment), environment.configuration, CompilerEnvironment, environment.project
         )
 
         // There are INVISIBLE_REFERENCE errors on RequireKotlin and K2JSTranslator refuses to translate the code otherwise
@@ -63,7 +63,7 @@ class JsVersionRequirementTest : AbstractVersionRequirementTest() {
     override fun loadModule(directory: File): ModuleDescriptor {
         val environment = createEnvironment(extraDependencies = listOf(File(directory, "lib.meta.js")))
         return TopDownAnalyzerFacadeForJS.analyzeFilesWithGivenTrace(
-            emptyList(), BindingTraceContext(), createModule(environment), environment.configuration, CompilerEnvironment,
+            emptyList(), BindingTraceContext(), createModule(environment), environment.configuration, CompilerEnvironment, environment.project
         ).moduleDescriptor
     }
 

js/js.frontend/src/org/jetbrains/kotlin/frontend/js/di/injection.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.frontend.js.di
 
 import org.jetbrains.kotlin.config.LanguageVersionSettings
+import org.jetbrains.kotlin.container.StorageComponentContainer
 import org.jetbrains.kotlin.container.get
 import org.jetbrains.kotlin.container.useInstance
 import org.jetbrains.kotlin.context.ModuleContext
@@ -37,7 +38,7 @@ import org.jetbrains.kotlin.resolve.createContainer
 import org.jetbrains.kotlin.resolve.lazy.KotlinCodeAnalyzer
 import org.jetbrains.kotlin.resolve.lazy.declarations.DeclarationProviderFactory
 
-fun createTopDownAnalyzerForJs(
+fun createContainerForJS(
     moduleContext: ModuleContext,
     bindingTrace: BindingTrace,
     declarationProviderFactory: DeclarationProviderFactory,
@@ -46,7 +47,7 @@ fun createTopDownAnalyzerForJs(
     expectActualTracker: ExpectActualTracker,
     additionalPackages: List<PackageFragmentProvider>,
     targetEnvironment: TargetEnvironment,
-): LazyTopDownAnalyzer {
+): StorageComponentContainer {
     val storageComponentContainer = createContainer("TopDownAnalyzerForJs", JsPlatformAnalyzerServices) {
         configureModule(
             moduleContext,
@@ -67,5 +68,21 @@ fun createTopDownAnalyzerForJs(
         packagePartProviders += additionalPackages
         moduleDescriptor.initialize(CompositePackageFragmentProvider(packagePartProviders))
     }
-    return storageComponentContainer.get<LazyTopDownAnalyzer>()
+    return storageComponentContainer
+}
+
+fun createTopDownAnalyzerForJs(
+    moduleContext: ModuleContext,
+    bindingTrace: BindingTrace,
+    declarationProviderFactory: DeclarationProviderFactory,
+    languageVersionSettings: LanguageVersionSettings,
+    lookupTracker: LookupTracker,
+    expectActualTracker: ExpectActualTracker,
+    additionalPackages: List<PackageFragmentProvider>,
+    targetEnvironment: TargetEnvironment,
+): LazyTopDownAnalyzer {
+    return createContainerForJS(
+        moduleContext, bindingTrace, declarationProviderFactory, languageVersionSettings,
+        lookupTracker, expectActualTracker, additionalPackages, targetEnvironment
+    ).get<LazyTopDownAnalyzer>()
 }

js/js.frontend/src/org/jetbrains/kotlin/js/analyze/TopDownAnalyzerFacadeForJS.kt

@@ -6,18 +6,20 @@
 package org.jetbrains.kotlin.js.analyze
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.builtins.DefaultBuiltIns
 import org.jetbrains.kotlin.builtins.functions.functionInterfacePackageFragmentProvider
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.container.get
 import org.jetbrains.kotlin.context.ContextForNewModule
 import org.jetbrains.kotlin.context.ModuleContext
 import org.jetbrains.kotlin.context.ProjectContext
 import org.jetbrains.kotlin.descriptors.PackageFragmentProvider
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
-import org.jetbrains.kotlin.frontend.js.di.createTopDownAnalyzerForJs
+import org.jetbrains.kotlin.frontend.js.di.createContainerForJS
 import org.jetbrains.kotlin.incremental.components.ExpectActualTracker
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.incremental.js.IncrementalDataProvider
@@ -31,12 +33,13 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.js.JsPlatforms
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.*
+import org.jetbrains.kotlin.resolve.extensions.AnalysisHandlerExtension
 import org.jetbrains.kotlin.resolve.lazy.declarations.FileBasedDeclarationProviderFactory
 import org.jetbrains.kotlin.serialization.js.KotlinJavascriptSerializationUtil
 import org.jetbrains.kotlin.serialization.js.ModuleKind
 import org.jetbrains.kotlin.serialization.js.PackagesWithHeaderMetadata
 import org.jetbrains.kotlin.utils.JsMetadataVersion
-import java.lang.Exception
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 abstract class AbstractTopDownAnalyzerFacadeForJS {
 
@@ -82,7 +85,7 @@ abstract class AbstractTopDownAnalyzerFacadeForJS {
 
         val trace = BindingTraceContext()
         trace.record(MODULE_KIND, context.module, moduleKind)
-        return analyzeFilesWithGivenTrace(files, trace, context, configuration, targetEnvironment, additionalPackages)
+        return analyzeFilesWithGivenTrace(files, trace, context, configuration, targetEnvironment, project, additionalPackages)
     }
 
     protected abstract fun loadIncrementalCacheMetadata(
@@ -98,6 +101,7 @@ abstract class AbstractTopDownAnalyzerFacadeForJS {
         moduleContext: ModuleContext,
         configuration: CompilerConfiguration,
         targetEnvironment: TargetEnvironment,
+        project: Project,
         additionalPackages: List<PackageFragmentProvider> = emptyList()
     ): JsAnalysisResult {
         val lookupTracker = configuration.get(CommonConfigurationKeys.LOOKUP_TRACKER) ?: LookupTracker.DO_NOTHING
@@ -106,7 +110,8 @@ abstract class AbstractTopDownAnalyzerFacadeForJS {
         val packageFragment = configuration[JSConfigurationKeys.INCREMENTAL_DATA_PROVIDER]?.let {
             loadIncrementalCacheMetadata(it, moduleContext, lookupTracker, languageVersionSettings)
         }
-        val analyzerForJs = createTopDownAnalyzerForJs(
+
+        val container = createContainerForJS(
             moduleContext, trace,
             FileBasedDeclarationProviderFactory(moduleContext.storageManager, files),
             languageVersionSettings,
@@ -115,8 +120,38 @@ abstract class AbstractTopDownAnalyzerFacadeForJS {
             additionalPackages + listOfNotNull(packageFragment),
             targetEnvironment,
         )
-        analyzerForJs.analyzeDeclarations(TopDownAnalysisMode.TopLevelDeclarations, files)
+
-        return JsAnalysisResult.success(trace, moduleContext.module)
+        val analysisHandlerExtensions = AnalysisHandlerExtension.getInstances(project)
+
+        // Mimic the behavior in the jvm frontend. The extensions have 2 chances to override the normal analysis:
+        // * If any of the extensions returns a non-null result, it. Otherwise do the normal analysis.
+        // * `analysisCompleted` can be used to override the result, too.
+        var result = analysisHandlerExtensions.firstNotNullResult { extension ->
+            extension.doAnalysis(project, moduleContext.module, moduleContext, files, trace, container)
+        } ?: run {
+            container.get<LazyTopDownAnalyzer>().analyzeDeclarations(TopDownAnalysisMode.TopLevelDeclarations, files)
+            AnalysisResult.success(trace.bindingContext, moduleContext.module)
+        }
+
+        result = analysisHandlerExtensions.firstNotNullResult { extension ->
+            extension.analysisCompleted(project, moduleContext.module, trace, files)
+        } ?: result
+
+        return when (result) {
+            is JsAnalysisResult -> result
+            else -> {
+                // AnalysisHandlerExtension returns a BindingContext, not BindingTrace. Therefore, synthesize one here.
+                val bindingTrace = DelegatingBindingTrace(result.bindingContext, "DelegatingBindingTrace by AnalysisHandlerExtension")
+                when (result) {
+                    is AnalysisResult.RetryWithAdditionalRoots -> JsAnalysisResult.RetryWithAdditionalRoots(
+                        bindingTrace,
+                        result.moduleDescriptor,
+                        result.additionalKotlinRoots
+                    )
+                    else -> JsAnalysisResult.success(bindingTrace, result.moduleDescriptor, result.shouldGenerateCode)
+                }
+            }
+        }
     }
 
     fun checkForErrors(allFiles: Collection<KtFile>, bindingContext: BindingContext, errorPolicy: ErrorTolerancePolicy): Boolean {

js/js.frontend/src/org/jetbrains/kotlin/js/analyzer/JsAnalysisResult.kt

@@ -21,15 +21,27 @@ import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.types.ErrorUtils
+import java.io.File
 
-class JsAnalysisResult(
+open class JsAnalysisResult(
         val bindingTrace: BindingTrace,
-        moduleDescriptor: ModuleDescriptor
+        moduleDescriptor: ModuleDescriptor,
-) : AnalysisResult(bindingTrace.bindingContext, moduleDescriptor) {
+        shouldGenerateCode: Boolean
+) : AnalysisResult(bindingTrace.bindingContext, moduleDescriptor, shouldGenerateCode) {
 
     companion object {
         @JvmStatic fun success(trace: BindingTrace, module: ModuleDescriptor): JsAnalysisResult {
-            return JsAnalysisResult(trace, module)
+            return JsAnalysisResult(trace, module, true)
+        }
+
+        @JvmStatic fun success(trace: BindingTrace, module: ModuleDescriptor, shouldGenerateCode: Boolean): JsAnalysisResult {
+            return JsAnalysisResult(trace, module, shouldGenerateCode)
         }
     }
+
+    class RetryWithAdditionalRoots(
+        bindingTrace: BindingTrace,
+        moduleDescriptor: ModuleDescriptor,
+        val additionalKotlinRoots: List<File>,
+    ) : JsAnalysisResult(bindingTrace, moduleDescriptor,false)
 }