1 files changed, 64 insertions, 24 deletions

diff --git a/plugins/templating/src/main/kotlin/templates/TemplateProcessor.kt b/plugins/templating/src/main/kotlin/templates/TemplateProcessor.kt
index 8fbd76b6..c67654ec 100644
--- a/plugins/templating/src/main/kotlin/templates/TemplateProcessor.kt
+++ b/plugins/templating/src/main/kotlin/templates/TemplateProcessor.kt
@@ -1,15 +1,26 @@
 package org.jetbrains.dokka.templates
 
 import kotlinx.coroutines.*
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.base.DokkaBase
 import org.jetbrains.dokka.base.templating.Command
+import org.jetbrains.dokka.model.withDescendants
+import org.jetbrains.dokka.pages.RootPageNode
 import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.dokka.plugability.plugin
 import org.jetbrains.dokka.plugability.query
+import org.jetbrains.dokka.plugability.querySingle
 import org.jsoup.nodes.Element
 import java.io.File
 
-interface TemplateProcessor {
-    fun process()
+interface TemplateProcessor
+
+interface SubmoduleTemplateProcessor : TemplateProcessor {
+    fun process(modules: List<DokkaConfiguration.DokkaModuleDescription>): TemplatingResult
+}
+
+interface MultiModuleTemplateProcessor : TemplateProcessor {
+    fun process(generatedPagesTree: RootPageNode)
 }
 
 interface TemplateProcessingStrategy {
@@ -17,40 +28,69 @@ interface TemplateProcessingStrategy {
     fun finish(output: File) {}
 }
 
-class DefaultTemplateProcessor(
+class DefaultSubmoduleTemplateProcessor(
     private val context: DokkaContext,
-): TemplateProcessor {
+) : SubmoduleTemplateProcessor {
 
-    private val strategies: List<TemplateProcessingStrategy> = context.plugin<TemplatingPlugin>().query { templateProcessingStrategy }
+    private val strategies: List<TemplateProcessingStrategy> =
+        context.plugin<TemplatingPlugin>().query { templateProcessingStrategy }
 
-    override fun process() = runBlocking(Dispatchers.Default) {
-        coroutineScope {
-            context.configuration.modules.forEach {
-                launch {
-                    it.sourceOutputDirectory.visit(context.configuration.outputDir.resolve(it.relativePathToOutputDirectory))
+    private val configuredModulesPaths =
+        context.configuration.modules.map { it.sourceOutputDirectory.absolutePath to it.name }.toMap()
+
+    override fun process(modules: List<DokkaConfiguration.DokkaModuleDescription>) =
+        runBlocking(Dispatchers.Default) {
+            coroutineScope {
+                modules.fold(TemplatingResult()) { acc, module ->
+                    acc + module.sourceOutputDirectory.visit(context.configuration.outputDir.resolve(module.relativePathToOutputDirectory))
                 }
             }
         }
-        strategies.map { it.finish(context.configuration.outputDir) }
-        Unit
-    }
 
-    private suspend fun File.visit(target: File): Unit = coroutineScope {
-        val source = this@visit
-        if (source.isDirectory) {
-           target.mkdir()
-           source.list()?.forEach {
-               launch { source.resolve(it).visit(target.resolve(it)) }
-           }
-        } else {
-            strategies.first { it.process(source, target) }
+    private suspend fun File.visit(target: File, acc: TemplatingResult = TemplatingResult()): TemplatingResult =
+        coroutineScope {
+            val source = this@visit
+            if (source.isDirectory) {
+                target.mkdir()
+                val files = source.list().orEmpty()
+                val accWithSelf = configuredModulesPaths[source.absolutePath]
+                    ?.takeIf { files.firstOrNull { !it.startsWith(".") } != null }
+                    ?.let { acc.copy(modules = acc.modules + it) }
+                    ?: acc
+
+                files.fold(accWithSelf) { acc, path ->
+                    source.resolve(path).visit(target.resolve(path), acc)
+                }
+            } else {
+                strategies.first { it.process(source, target) }
+                acc
+            }
         }
+}
+
+class DefaultMultiModuleTemplateProcessor(
+    context: DokkaContext,
+) : MultiModuleTemplateProcessor {
+    private val strategies: List<TemplateProcessingStrategy> =
+        context.plugin<TemplatingPlugin>().query { templateProcessingStrategy }
+
+    private val locationProviderFactory = context.plugin<DokkaBase>().querySingle { locationProviderFactory }
+
+    override fun process(generatedPagesTree: RootPageNode) {
+        val locationProvider = locationProviderFactory.getLocationProvider(generatedPagesTree)
+        generatedPagesTree.withDescendants().mapNotNull { locationProvider.resolve(it)?.let { File(it) } }
+            .forEach { location -> strategies.first { it.process(location, location) } }
     }
+
 }
 
-data class TemplatingContext<out T: Command>(
+data class TemplatingContext<out T : Command>(
     val input: File,
     val output: File,
     val element: Element,
     val command: T,
-)
\ No newline at end of file
+)
+
+data class TemplatingResult(val modules: List<String> = emptyList()) {
+    operator fun plus(rhs: TemplatingResult): TemplatingResult = TemplatingResult((modules + rhs.modules).distinct())
+}
\ No newline at end of file