From 4db79285c146c8e48d2e9627782cfac54dfdcbab Mon Sep 17 00:00:00 2001
From: Ting-Yuan Huang <laszio@google.com>
Date: Mon, 15 Mar 2021 20:58:23 -0700
Subject: [PATCH] Add AnalysisHandlerExtension extension point for
 K2MetadataCompiler

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.
---
 .../kotlin/cli/metadata/CommonAnalysis.kt     | 12 +++++++++-
 .../common/CommonResolverForModuleFactory.kt  | 24 ++++++++++++++++---
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/compiler/cli/src/org/jetbrains/kotlin/cli/metadata/CommonAnalysis.kt b/compiler/cli/src/org/jetbrains/kotlin/cli/metadata/CommonAnalysis.kt
index 6dc7a889a94..a122a53995e 100644
--- a/compiler/cli/src/org/jetbrains/kotlin/cli/metadata/CommonAnalysis.kt
+++ b/compiler/cli/src/org/jetbrains/kotlin/cli/metadata/CommonAnalysis.kt
@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.cli.metadata
 
+import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.analyzer.common.CommonDependenciesContainer
 import org.jetbrains.kotlin.analyzer.common.CommonResolverForModuleFactory
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
@@ -33,7 +34,16 @@ internal fun runCommonAnalysisForSerialization(
         return null
     }
 
-    return runCommonAnalysis(environment, dependOnBuiltins, dependencyContainer)
+    var analyzer: AnalyzerWithCompilerReport
+    do {
+        analyzer = runCommonAnalysis(environment, dependOnBuiltins, dependencyContainer)
+        val result = analyzer.analysisResult
+        if (result is AnalysisResult.RetryWithAdditionalRoots) {
+            environment.addKotlinSourceRoots(result.additionalKotlinRoots)
+        }
+    } while (result is AnalysisResult.RetryWithAdditionalRoots)
+
+    return if (analyzer.analysisResult.shouldGenerateCode) analyzer else null
 }
 
 private fun runCommonAnalysis(
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/analyzer/common/CommonResolverForModuleFactory.kt b/compiler/frontend/src/org/jetbrains/kotlin/analyzer/common/CommonResolverForModuleFactory.kt
index cf98ea40da0..086a69d1d42 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/analyzer/common/CommonResolverForModuleFactory.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/analyzer/common/CommonResolverForModuleFactory.kt
@@ -40,11 +40,13 @@ import org.jetbrains.kotlin.platform.TargetPlatform
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.checkers.ExpectedActualDeclarationChecker
+import org.jetbrains.kotlin.resolve.extensions.AnalysisHandlerExtension
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession
 import org.jetbrains.kotlin.resolve.lazy.declarations.DeclarationProviderFactory
 import org.jetbrains.kotlin.resolve.lazy.declarations.DeclarationProviderFactoryService
 import org.jetbrains.kotlin.serialization.deserialization.MetadataPackageFragmentProvider
 import org.jetbrains.kotlin.serialization.deserialization.MetadataPartProvider
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 
 class CommonAnalysisParameters(
     val metadataPartProviderFactory: (ModuleContent<*>) -> MetadataPartProvider
@@ -148,10 +150,11 @@ class CommonResolverForModuleFactory(
                 dependenciesContainer
             )
 
+            val projectContext = ProjectContext(project, "metadata serializer")
             @Suppress("NAME_SHADOWING")
             val resolver = ResolverForSingleModuleProject<ModuleInfo>(
                 "sources for metadata serializer",
-                ProjectContext(project, "metadata serializer"),
+                projectContext,
                 moduleInfo,
                 resolverForModuleFactory,
                 GlobalSearchScope.allScope(project),
@@ -167,9 +170,24 @@ class CommonResolverForModuleFactory(
 
             val container = resolver.resolverForModule(moduleInfo).componentProvider
 
-            container.get<LazyTopDownAnalyzer>().analyzeDeclarations(TopDownAnalysisMode.TopLevelDeclarations, files)
+            val analysisHandlerExtensions = AnalysisHandlerExtension.getInstances(project)
+            val trace = container.get<BindingTrace>()
 
-            return AnalysisResult.success(container.get<BindingTrace>().bindingContext, moduleDescriptor)
+            // 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, moduleDescriptor, projectContext, files, trace, container)
+            } ?: run {
+                container.get<LazyTopDownAnalyzer>().analyzeDeclarations(TopDownAnalysisMode.TopLevelDeclarations, files)
+                AnalysisResult.success(trace.bindingContext, moduleDescriptor)
+            }
+
+            result = analysisHandlerExtensions.firstNotNullResult { extension ->
+                extension.analysisCompleted(project, moduleDescriptor, trace, files)
+            } ?: result
+
+            return result
         }
     }
 }