Dispose `AnalysisEnvironment` (#2755)

3 files changed, 165 insertions, 72 deletions

diff --git a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisContext.kt b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisContext.kt
new file mode 100644
index 00000000..ca83d029
--- /dev/null
+++ b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisContext.kt
@@ -0,0 +1,94 @@
+package org.jetbrains.dokka.analysis
+
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.Platform
+import org.jetbrains.dokka.utilities.DokkaLogger
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSourceLocation
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import java.io.Closeable
+import java.io.File
+
+internal fun createAnalysisContext(
+    logger: DokkaLogger,
+    sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
+    sourceSet: DokkaConfiguration.DokkaSourceSet,
+    analysisConfiguration: DokkaAnalysisConfiguration
+): AnalysisContext {
+    val parentSourceSets = sourceSets.filter { it.sourceSetID in sourceSet.dependentSourceSets }
+    val classpath = sourceSet.classpath + parentSourceSets.flatMap { it.classpath }
+    val sources = sourceSet.sourceRoots + parentSourceSets.flatMap { it.sourceRoots }
+
+    return createAnalysisContext(
+        logger = logger,
+        classpath = classpath,
+        sourceRoots = sources,
+        sourceSet = sourceSet,
+        analysisConfiguration = analysisConfiguration
+    )
+}
+
+internal fun createAnalysisContext(
+    logger: DokkaLogger,
+    classpath: List<File>,
+    sourceRoots: Set<File>,
+    sourceSet: DokkaConfiguration.DokkaSourceSet,
+    analysisConfiguration: DokkaAnalysisConfiguration
+): AnalysisContext {
+    val analysisEnvironment = AnalysisEnvironment(DokkaMessageCollector(logger), sourceSet.analysisPlatform).apply {
+        if (analysisPlatform == Platform.jvm) {
+            configureJdkClasspathRoots()
+        }
+        addClasspath(classpath)
+        addSources(sourceRoots)
+
+        loadLanguageVersionSettings(sourceSet.languageVersion, sourceSet.apiVersion)
+    }
+
+    val environment = analysisEnvironment.createCoreEnvironment()
+    val (facade, _) = analysisEnvironment.createResolutionFacade(
+        environment,
+        analysisConfiguration.ignoreCommonBuiltIns
+    )
+
+    return AnalysisContext(environment, facade, analysisEnvironment)
+}
+
+class DokkaMessageCollector(private val logger: DokkaLogger) : MessageCollector {
+    override fun clear() {
+        seenErrors = false
+    }
+
+    private var seenErrors = false
+
+    override fun report(severity: CompilerMessageSeverity, message: String, location: CompilerMessageSourceLocation?) {
+        if (severity == CompilerMessageSeverity.ERROR) {
+            seenErrors = true
+        }
+        logger.info(MessageRenderer.PLAIN_FULL_PATHS.render(severity, message, location))
+    }
+
+    override fun hasErrors() = seenErrors
+}
+
+// It is not data class due to ill-defined equals
+class AnalysisContext(
+    environment: KotlinCoreEnvironment,
+    facade: DokkaResolutionFacade,
+    private val analysisEnvironment: AnalysisEnvironment
+) : Closeable {
+    private var isClosed: Boolean = false
+    val environment: KotlinCoreEnvironment = environment
+        get() = field.takeUnless { isClosed } ?: throw IllegalStateException("AnalysisEnvironment is already closed")
+    val facade: DokkaResolutionFacade = facade
+        get() = field.takeUnless { isClosed } ?: throw IllegalStateException("AnalysisEnvironment is already closed")
+
+    operator fun component1() = environment
+    operator fun component2() = facade
+    override fun close() {
+        isClosed = true
+        analysisEnvironment.dispose()
+    }
+}
diff --git a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/EnvironmentAndFacade.kt b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/EnvironmentAndFacade.kt
deleted file mode 100644
index b946c5bd..00000000
--- a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/EnvironmentAndFacade.kt
+++ /dev/null
@@ -1,57 +0,0 @@
-package org.jetbrains.dokka.analysis
-
-import org.jetbrains.dokka.DokkaConfiguration
-import org.jetbrains.dokka.Platform
-import org.jetbrains.dokka.utilities.DokkaLogger
-import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
-import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSourceLocation
-import org.jetbrains.kotlin.cli.common.messages.MessageCollector
-import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
-import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
-
-internal fun createEnvironmentAndFacade(
-    logger: DokkaLogger,
-    sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
-    sourceSet: DokkaConfiguration.DokkaSourceSet,
-    analysisConfiguration: DokkaAnalysisConfiguration
-): EnvironmentAndFacade =
-    AnalysisEnvironment(DokkaMessageCollector(logger), sourceSet.analysisPlatform).run {
-        if (analysisPlatform == Platform.jvm) {
-            configureJdkClasspathRoots()
-        }
-
-        val parentSourceSets = sourceSets.filter { it.sourceSetID in sourceSet.dependentSourceSets }
-        addClasspath(sourceSet.classpath + parentSourceSets.flatMap { it.classpath })
-
-        addSources(sourceSet.sourceRoots + parentSourceSets.flatMap { it.sourceRoots })
-
-        loadLanguageVersionSettings(sourceSet.languageVersion, sourceSet.apiVersion)
-
-        val environment = createCoreEnvironment()
-
-        val (facade, _) = createResolutionFacade(environment, analysisConfiguration.ignoreCommonBuiltIns)
-        EnvironmentAndFacade(environment, facade)
-    }
-
-class DokkaMessageCollector(private val logger: DokkaLogger) : MessageCollector {
-    override fun clear() {
-        seenErrors = false
-    }
-
-    private var seenErrors = false
-
-    override fun report(severity: CompilerMessageSeverity, message: String, location: CompilerMessageSourceLocation?) {
-        if (severity == CompilerMessageSeverity.ERROR) {
-            seenErrors = true
-        }
-        logger.info(MessageRenderer.PLAIN_FULL_PATHS.render(severity, message, location))
-    }
-
-    override fun hasErrors() = seenErrors
-}
-
-// It is not data class due to ill-defined equals
-class EnvironmentAndFacade(val environment: KotlinCoreEnvironment, val facade: DokkaResolutionFacade) {
-    operator fun component1() = environment
-    operator fun component2() = facade
-}
diff --git a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/KotlinAnalysis.kt b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/KotlinAnalysis.kt
index a188e3f9..64a583b6 100644
--- a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/KotlinAnalysis.kt
+++ b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/KotlinAnalysis.kt
@@ -7,18 +7,48 @@ import org.jetbrains.dokka.DokkaSourceSetID
 import org.jetbrains.dokka.model.SourceSetDependent
 import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.dokka.utilities.DokkaLogger
+import java.io.Closeable
 
-fun KotlinAnalysis(sourceSets: List<DokkaSourceSet>, logger: DokkaLogger, analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()): KotlinAnalysis {
+fun ProjectKotlinAnalysis(
+    sourceSets: List<DokkaSourceSet>,
+    logger: DokkaLogger,
+    analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()
+): KotlinAnalysis {
     val environments = sourceSets.associateWith { sourceSet ->
-        createEnvironmentAndFacade(
+        createAnalysisContext(
             logger = logger,
             sourceSets = sourceSets,
             sourceSet = sourceSet,
             analysisConfiguration = analysisConfiguration
         )
     }
+    return EnvironmentKotlinAnalysis(environments)
+}
+
+/**
+ * [projectKotlinAnalysis] needs to be closed separately
+ *  Usually the analysis created for samples is short-lived and can be closed right after
+ *  it's been used, there's no need to wait for [projectKotlinAnalysis] to be closed as it must be handled separately.
+ */
+fun SamplesKotlinAnalysis(
+    sourceSets: List<DokkaSourceSet>,
+    logger: DokkaLogger,
+    projectKotlinAnalysis: KotlinAnalysis,
+    analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()
+): KotlinAnalysis {
+    val environments = sourceSets
+        .filter { it.samples.isNotEmpty() }
+        .associateWith { sourceSet ->
+            createAnalysisContext(
+                logger = logger,
+                classpath = sourceSet.classpath,
+                sourceRoots = sourceSet.samples,
+                sourceSet = sourceSet,
+                analysisConfiguration = analysisConfiguration
+            )
+    }
 
-    return KotlinAnalysisImpl(environments)
+    return EnvironmentKotlinAnalysis(environments, projectKotlinAnalysis)
 }
 
 class DokkaAnalysisConfiguration(
@@ -32,22 +62,48 @@ class DokkaAnalysisConfiguration(
 @Deprecated(message = "Construct using list of DokkaSourceSets and logger",
     replaceWith = ReplaceWith("KotlinAnalysis(context.configuration.sourceSets, context.logger)")
 )
-fun KotlinAnalysis(context: DokkaContext): KotlinAnalysis = KotlinAnalysis(context.configuration.sourceSets, context.logger)
+fun KotlinAnalysis(context: DokkaContext): KotlinAnalysis =
+    ProjectKotlinAnalysis(context.configuration.sourceSets, context.logger)
+
+@Deprecated(message = "It was renamed to `ProjectKotlinAnalysis`",
+    replaceWith = ReplaceWith("ProjectKotlinAnalysis(sourceSets, logger, analysisConfiguration)")
+)
+fun KotlinAnalysis(
+    sourceSets: List<DokkaSourceSet>,
+    logger: DokkaLogger,
+    analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()
+) = ProjectKotlinAnalysis(sourceSets, logger, analysisConfiguration)
 
-interface KotlinAnalysis : SourceSetDependent<EnvironmentAndFacade> {
-    override fun get(key: DokkaSourceSet): EnvironmentAndFacade
-    operator fun get(sourceSetID: DokkaSourceSetID): EnvironmentAndFacade
-}
 
-internal class KotlinAnalysisImpl(
-    private val environments: SourceSetDependent<EnvironmentAndFacade>
-) : KotlinAnalysis, SourceSetDependent<EnvironmentAndFacade> by environments {
+/**
+ * First child delegation. It does not close [parent].
+ */
+abstract class KotlinAnalysis(
+    val parent: KotlinAnalysis? = null
+) : Closeable {
 
-    override fun get(key: DokkaSourceSet): EnvironmentAndFacade {
-        return environments[key] ?: throw IllegalStateException("Missing EnvironmentAndFacade for sourceSet $key")
+    operator fun get(key: DokkaSourceSet): AnalysisContext {
+        return get(key.sourceSetID)
     }
+    operator fun get(key: DokkaSourceSetID): AnalysisContext {
+        return find(key)
+            ?: parent?.get(key)
+            ?: throw IllegalStateException("Missing EnvironmentAndFacade for sourceSet ${key}")
+    }
+    protected abstract fun find(sourceSetID: DokkaSourceSetID): AnalysisContext?
+}
+
+internal open class EnvironmentKotlinAnalysis(
+    private val environments: SourceSetDependent<AnalysisContext>,
+    parent: KotlinAnalysis? = null,
+) : KotlinAnalysis(parent = parent) {
 
-    override fun get(sourceSetID: DokkaSourceSetID): EnvironmentAndFacade {
-        return environments.entries.first { (sourceSet, _) -> sourceSet.sourceSetID == sourceSetID }.value
+    override fun find(sourceSetID: DokkaSourceSetID): AnalysisContext? =
+        environments.entries.firstOrNull { (sourceSet, _) -> sourceSet.sourceSetID == sourceSetID }?.value
+
+    override fun close() {
+        environments.values.forEach(AnalysisContext::close)
     }
 }
+
+