Cache DisplaySourceSet as it's stored a lot inside of ContentPages (

#4008)

Additionally rewrite sourceSet uniqueness test in a way that will not cause full Dokka pipeline execution, as if we stop at verification stage we should not proceed with execution

dokka-subprojects/core/src/main/kotlin/org/jetbrains/dokka/DokkaGenerator.kt

@@ -7,6 +7,7 @@ package org.jetbrains.dokka
 import kotlinx.coroutines.DelicateCoroutinesApi
 import kotlinx.coroutines.Dispatchers
 import org.jetbrains.dokka.generation.GracefulGenerationExit
+import org.jetbrains.dokka.model.DisplaySourceSetCaches
 import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.dokka.plugability.DokkaPlugin
 import org.jetbrains.dokka.utilities.DokkaLogger
@@ -20,23 +21,24 @@ public class DokkaGenerator(
     private val configuration: DokkaConfiguration,
     private val logger: DokkaLogger
 ) {
+    init {
+        DisplaySourceSetCaches.clear()
+    }
 
     public fun generate() {
         timed(logger) {
             report("Initializing plugins")
             val context = initializePlugins(configuration, logger)
 
-            runCatching {
+            try {
                 context.single(CoreExtensions.generation).run {
                     logger.progress("Dokka is performing: $generationName")
                     generate()
                 }
-            }.exceptionOrNull()?.let { e ->
+            } finally {
+                DisplaySourceSetCaches.clear()
                 finalizeCoroutines()
-                throw e
             }
-
-            finalizeCoroutines()
         }.dump("\n\n === TIME MEASUREMENT ===\n")
     }
 

dokka-subprojects/core/src/main/kotlin/org/jetbrains/dokka/model/DisplaySourceSet.kt

@@ -5,7 +5,9 @@
 package org.jetbrains.dokka.model
 
 import org.jetbrains.dokka.*
+import org.jetbrains.dokka.pages.*
 import org.jetbrains.dokka.DokkaConfiguration.DokkaSourceSet
+import java.util.concurrent.ConcurrentHashMap
 
 /**
  * Represents a final user-visible source set in the documentable model that is
@@ -44,13 +46,16 @@ public data class DisplaySourceSet(
  * Transforms the current [DokkaSourceSet] into [DisplaySourceSet],
  * matching the corresponding subset of its properties to [DisplaySourceSet] properties.
  */
-public fun DokkaSourceSet.toDisplaySourceSet(): DisplaySourceSet = DisplaySourceSet(this)
+public fun DokkaSourceSet.toDisplaySourceSet(): DisplaySourceSet {
+    return DisplaySourceSetCaches.displaySourceSet(this)
+}
 
 /**
  * Transforms all the given [DokkaSourceSet]s into [DisplaySourceSet]s.
  */
-public fun Iterable<DokkaSourceSet>.toDisplaySourceSets(): Set<DisplaySourceSet> =
-    map { it.toDisplaySourceSet() }.toSet()
+public fun Iterable<DokkaSourceSet>.toDisplaySourceSets(): Set<DisplaySourceSet> {
+    return DisplaySourceSetCaches.displaySourceSets(this)
+}
 
 @InternalDokkaApi
 @Deprecated("Use computeSourceSetIds() and cache its results instead", replaceWith = ReplaceWith("computeSourceSetIds()"))
@@ -59,3 +64,30 @@ public val Iterable<DisplaySourceSet>.sourceSetIDs: List<DokkaSourceSetID> get()
 @InternalDokkaApi
 public fun Iterable<DisplaySourceSet>.computeSourceSetIds(): Set<DokkaSourceSetID> =
     fold(hashSetOf()) { acc, set -> acc.addAll(set.sourceSetIDs.all); acc }
+
+internal object DisplaySourceSetCaches {
+    private val instanceCache = ConcurrentHashMap<DokkaSourceSet, DisplaySourceSet>()
+    private val setCache = ConcurrentHashMap<Iterable<DokkaSourceSet>, Set<DisplaySourceSet>>()
+
+    /**
+     * Construction of [DisplaySourceSet] from [DokkaSourceSet] requires to create a [CompositeSourceSetID], which has a big memory footprint.
+     * [toDisplaySourceSets] is called for almost each [ContentNode] a LOT.
+     * In reality, there will be only several [DokkaSourceSet] instances during Dokka execution.
+     */
+    fun displaySourceSet(sourceSet: DokkaSourceSet): DisplaySourceSet {
+        return instanceCache.computeIfAbsent(sourceSet, ::DisplaySourceSet)
+    }
+
+    /**
+     * [toDisplaySourceSets] mostly called on the same Set<DokkaSourceSet>,
+     * so it makes sense to cache them by the set itself.
+     */
+    fun displaySourceSets(sourceSets: Iterable<DokkaSourceSet>): Set<DisplaySourceSet> {
+        return setCache.computeIfAbsent(sourceSets) { it.map(::displaySourceSet).toSet() }
+    }
+
+    fun clear() {
+        instanceCache.clear()
+        setCache.clear()
+    }
+}

dokka-subprojects/plugin-base/src/test/kotlin/generation/SourceSetIdUniquenessCheckerTest.kt

@@ -4,22 +4,19 @@
 
 package generation
 
+import org.jetbrains.dokka.CoreExtensions
 import org.jetbrains.dokka.DokkaException
+import org.jetbrains.dokka.base.generation.SingleModuleGeneration
 import org.jetbrains.dokka.base.testApi.testRunner.BaseAbstractTest
+import org.jetbrains.dokka.plugability.DokkaContext
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 
 class SourceSetIdUniquenessCheckerTest : BaseAbstractTest() {
     @Test
-    fun `pre-generation check should fail if there are sourceSets with the same id`() = testInline(
-        """
-        |/src/main1/file1.kt
-        |fun someFunction2()
-        |/src/main2/file2.kt
-        |fun someFunction2()
-        """.trimMargin(),
-        dokkaConfiguration {
+    fun `pre-generation check should fail if there are sourceSets with the same id`() {
+        val configuration = dokkaConfiguration {
             sourceSets {
                 sourceSet {
                     sourceRoots = listOf("src/main1")
@@ -33,25 +30,19 @@ class SourceSetIdUniquenessCheckerTest : BaseAbstractTest() {
                 }
             }
         }
-    ) {
-        verificationStage = { verification ->
-            val exception = assertFailsWith(DokkaException::class, verification)
-            assertEquals(
-                exception.message,
-                "Pre-generation validity check failed: Source sets 'S1' and 'S2' have the same `sourceSetID=root/JVM`. Every source set should have unique sourceSetID."
-            )
-        }
+        val context = DokkaContext.create(configuration, logger, emptyList())
+        val generation = context.single(CoreExtensions.generation) as SingleModuleGeneration
+
+        val exception = assertFailsWith<DokkaException> { generation.validityCheck(context) }
+        assertEquals(
+            exception.message,
+            "Pre-generation validity check failed: Source sets 'S1' and 'S2' have the same `sourceSetID=root/JVM`. Every source set should have unique sourceSetID."
+        )
     }
 
     @Test
-    fun `pre-generation check should not fail if sourceSets have different ids`() = testInline(
-        """
-        |/src/main1/file1.kt
-        |fun someFunction2()
-        |/src/main2/file2.kt
-        |fun someFunction2()
-        """.trimMargin(),
-        dokkaConfiguration {
+    fun `pre-generation check should not fail if sourceSets have different ids`() {
+        val configuration = dokkaConfiguration {
             sourceSets {
                 sourceSet {
                     sourceRoots = listOf("src/main1")
@@ -65,10 +56,11 @@ class SourceSetIdUniquenessCheckerTest : BaseAbstractTest() {
                 }
             }
         }
-    ) {
-        verificationStage = { verification ->
-            // we check that there is no error thrown
-            assertEquals(verification(), Unit)
-        }
+        val context = DokkaContext.create(configuration, logger, emptyList())
+        val generation = context.single(CoreExtensions.generation) as SingleModuleGeneration
+
+        // check no error thrown
+        // assertEquals is needed not to have a dangling declaration, it has `assertNotFails` semantics.
+        assertEquals(generation.validityCheck(context), Unit)
     }
 }