From 9d7c3726d2359e09760ec45ca5e2138ec1b52cc2 Mon Sep 17 00:00:00 2001
From: Vadim Mishenev <vad-mishenev@yandex.ru>
Date: Tue, 8 Aug 2023 17:15:51 +0300
Subject: [PATCH] Fix and refactor Sample Transformer (#3102)

(cherry picked from commit 2fd8e9096706545f8b77e1e66bcc876d7e29f82c)
---
 plugins/base/src/main/kotlin/DokkaBase.kt     |  16 +--
 .../DefaultTemplateModelFactory.kt            |  17 +--
 .../pages/DefaultSamplesTransformer.kt        | 113 ++++++++++++++++++
 .../content/samples/ContentForSamplesTest.kt  |  12 +-
 .../api/analysis-kotlin-api.api               |  15 +++
 .../internal/InternalKotlinAnalysisPlugin.kt  |   2 +
 .../kotlin/internal/SampleProvider.kt         |  30 +++++
 .../compiler/api/compiler.api                 |  15 +++
 .../CompilerDescriptorAnalysisPlugin.kt       |   9 +-
 .../compiler/impl/KotlinSampleProvider.kt     | 109 +++++++++++++++++
 10 files changed, 319 insertions(+), 19 deletions(-)
 create mode 100644 plugins/base/src/main/kotlin/transformers/pages/DefaultSamplesTransformer.kt
 create mode 100644 subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/SampleProvider.kt
 create mode 100644 subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/impl/KotlinSampleProvider.kt

diff --git a/plugins/base/src/main/kotlin/DokkaBase.kt b/plugins/base/src/main/kotlin/DokkaBase.kt
index 6da71b4bd8..9907d9b735 100644
--- a/plugins/base/src/main/kotlin/DokkaBase.kt
+++ b/plugins/base/src/main/kotlin/DokkaBase.kt
@@ -1,5 +1,3 @@
-@file:Suppress("unused")
-
 package org.jetbrains.dokka.base
 
 import org.jetbrains.dokka.CoreExtensions
@@ -19,6 +17,7 @@ import org.jetbrains.dokka.base.signatures.KotlinSignatureProvider
 import org.jetbrains.dokka.base.signatures.SignatureProvider
 import org.jetbrains.dokka.base.templating.ImmediateHtmlCommandConsumer
 import org.jetbrains.dokka.base.transformers.documentables.*
+import org.jetbrains.dokka.base.transformers.pages.DefaultSamplesTransformer
 import org.jetbrains.dokka.base.transformers.pages.annotations.SinceKotlinTransformer
 import org.jetbrains.dokka.base.transformers.pages.comments.CommentsToContentConverter
 import org.jetbrains.dokka.base.transformers.pages.comments.DocTagToContentConverter
@@ -33,6 +32,7 @@ import org.jetbrains.dokka.plugability.PluginApiPreviewAcknowledgement
 import org.jetbrains.dokka.transformers.documentation.PreMergeDocumentableTransformer
 import org.jetbrains.dokka.transformers.pages.PageTransformer
 
+@Suppress("unused")
 class DokkaBase : DokkaPlugin() {
 
     val preMergeDocumentableTransformer by extensionPoint<PreMergeDocumentableTransformer>()
@@ -149,7 +149,6 @@ class DokkaBase : DokkaPlugin() {
 
     val pageMerger by extending {
         CoreExtensions.pageTransformer providing ::PageMerger order {
-            // TODO [beresnev] make last() or at least after samples transformer
         }
     }
 
@@ -191,6 +190,12 @@ class DokkaBase : DokkaPlugin() {
         htmlPreprocessors with RootCreator applyIf { !delayTemplateSubstitution }
     }
 
+    val defaultSamplesTransformer by extending {
+        CoreExtensions.pageTransformer providing ::DefaultSamplesTransformer order {
+            before(pageMerger)
+        }
+    }
+
     val sourceLinksTransformer by extending {
         htmlPreprocessors providing ::SourceLinksTransformer order { after(rootCreator) }
     }
@@ -270,11 +275,6 @@ class DokkaBase : DokkaPlugin() {
     val defaultKotlinAnalysis: org.jetbrains.dokka.plugability.Extension<org.jetbrains.dokka.analysis.KotlinAnalysis, *, *>
         get() = throw org.jetbrains.dokka.base.deprecated.AnalysisApiDeprecatedError()
 
-    @Suppress("DeprecatedCallableAddReplaceWith")
-    @Deprecated(message = org.jetbrains.dokka.base.deprecated.ANALYSIS_API_DEPRECATION_MESSAGE, level = DeprecationLevel.ERROR)
-    val defaultSamplesTransformer: org.jetbrains.dokka.plugability.Extension<PageTransformer, *, *>
-        get() = throw org.jetbrains.dokka.base.deprecated.AnalysisApiDeprecatedError()
-
     @Suppress("DEPRECATION_ERROR", "DeprecatedCallableAddReplaceWith")
     @Deprecated(message = org.jetbrains.dokka.base.deprecated.ANALYSIS_API_DEPRECATION_MESSAGE, level = DeprecationLevel.ERROR)
     val defaultExternalDocumentablesProvider: org.jetbrains.dokka.plugability.Extension<org.jetbrains.dokka.base.translators.descriptors.ExternalDocumentablesProvider, *, *>
diff --git a/plugins/base/src/main/kotlin/renderers/html/innerTemplating/DefaultTemplateModelFactory.kt b/plugins/base/src/main/kotlin/renderers/html/innerTemplating/DefaultTemplateModelFactory.kt
index aae2f65d24..4f4cdd7cca 100644
--- a/plugins/base/src/main/kotlin/renderers/html/innerTemplating/DefaultTemplateModelFactory.kt
+++ b/plugins/base/src/main/kotlin/renderers/html/innerTemplating/DefaultTemplateModelFactory.kt
@@ -91,29 +91,30 @@ class DefaultTemplateModelFactory(val context: DokkaContext) : TemplateModelFact
         get() = URI(this).isAbsolute
 
     private fun Appendable.resourcesForPage(pathToRoot: String, resources: List<String>): Unit =
-        resources.forEach {
+        resources.forEach { resource ->
+
             val resourceHtml = with(createHTML()) {
                 when {
-                    it.URIExtension == "css" ->
+
+                    resource.URIExtension == "css" ->
                         link(
                             rel = LinkRel.stylesheet,
-                            href = if (it.isAbsolute) it else "$pathToRoot$it"
+                            href = if (resource.isAbsolute) resource else "$pathToRoot$resource"
                         )
-                    it.URIExtension == "js" ->
+                    resource.URIExtension == "js" ->
                         script(
                             type = ScriptType.textJavaScript,
-                            src = if (it.isAbsolute) it else "$pathToRoot$it"
+                            src = if (resource.isAbsolute) resource else "$pathToRoot$resource"
                         ) {
-                            if (it == "scripts/main.js" || it.endsWith("_deferred.js"))
+                            if (resource == "scripts/main.js" || resource.endsWith("_deferred.js"))
                                 defer = true
                             else
                                 async = true
                         }
-                    it.isImage() -> link(href = if (it.isAbsolute) it else "$pathToRoot$it")
+                    resource.isImage() -> link(href = if (resource.isAbsolute) resource else "$pathToRoot$resource")
                     else -> null
                 }
             }
-
             if (resourceHtml != null) {
                 append(resourceHtml)
             }
diff --git a/plugins/base/src/main/kotlin/transformers/pages/DefaultSamplesTransformer.kt b/plugins/base/src/main/kotlin/transformers/pages/DefaultSamplesTransformer.kt
new file mode 100644
index 0000000000..0602bde20e
--- /dev/null
+++ b/plugins/base/src/main/kotlin/transformers/pages/DefaultSamplesTransformer.kt
@@ -0,0 +1,113 @@
+package org.jetbrains.dokka.base.transformers.pages
+
+import org.jetbrains.dokka.links.DRI
+import org.jetbrains.dokka.model.DisplaySourceSet
+import org.jetbrains.dokka.model.doc.Sample
+import org.jetbrains.dokka.model.properties.PropertyContainer
+import org.jetbrains.dokka.pages.*
+import org.jetbrains.dokka.plugability.DokkaContext
+import org.jetbrains.dokka.plugability.plugin
+import org.jetbrains.dokka.plugability.querySingle
+import org.jetbrains.dokka.transformers.pages.PageTransformer
+import org.jetbrains.dokka.analysis.kotlin.internal.InternalKotlinAnalysisPlugin
+import org.jetbrains.dokka.analysis.kotlin.internal.SampleProvider
+import org.jetbrains.dokka.analysis.kotlin.internal.SampleProviderFactory
+
+internal const val KOTLIN_PLAYGROUND_SCRIPT = "https://unpkg.com/kotlin-playground@1/dist/playground.min.js"
+
+internal class DefaultSamplesTransformer(val context: DokkaContext) : PageTransformer {
+
+    private val sampleProviderFactory: SampleProviderFactory = context.plugin<InternalKotlinAnalysisPlugin>().querySingle { sampleProviderFactory }
+
+    override fun invoke(input: RootPageNode): RootPageNode {
+        return sampleProviderFactory.build().use { sampleProvider ->
+            input.transformContentPagesTree { page ->
+                val samples = (page as? WithDocumentables)?.documentables?.flatMap {
+                    it.documentation.entries.flatMap { entry ->
+                        entry.value.children.filterIsInstance<Sample>().map { entry.key to it }
+                    }
+                } ?: return@transformContentPagesTree page
+
+                val newContent = samples.fold(page.content) { acc, (sampleSourceSet, sample) ->
+                    sampleProvider.getSample(sampleSourceSet, sample.name)
+                        ?.let {
+                            acc.addSample(page, sample.name, it)
+                        } ?: acc
+                }
+
+                page.modified(
+                    content = newContent,
+                    embeddedResources = page.embeddedResources + KOTLIN_PLAYGROUND_SCRIPT
+                )
+            }
+        }
+    }
+
+
+    private fun ContentNode.addSample(
+        contentPage: ContentPage,
+        fqLink: String,
+        sample: SampleProvider.SampleSnippet,
+    ): ContentNode {
+        val node = contentCode(contentPage.content.sourceSets, contentPage.dri, createSampleBody(sample.imports, sample.body), "kotlin")
+        return dfs(fqLink, node)
+    }
+
+    fun createSampleBody(imports: String, body: String) =
+        """ |$imports
+            |fun main() { 
+            |   //sampleStart 
+            |   $body 
+            |   //sampleEnd
+            |}""".trimMargin()
+
+    private fun ContentNode.dfs(fqName: String, node: ContentCodeBlock): ContentNode {
+        return when (this) {
+            is ContentHeader -> copy(children.map { it.dfs(fqName, node) })
+            is ContentDivergentGroup -> @Suppress("UNCHECKED_CAST") copy(children.map {
+                it.dfs(fqName, node)
+            } as List<ContentDivergentInstance>)
+            is ContentDivergentInstance -> copy(
+                before.let { it?.dfs(fqName, node) },
+                divergent.dfs(fqName, node),
+                after.let { it?.dfs(fqName, node) })
+            is ContentCodeBlock -> copy(children.map { it.dfs(fqName, node) })
+            is ContentCodeInline -> copy(children.map { it.dfs(fqName, node) })
+            is ContentDRILink -> copy(children.map { it.dfs(fqName, node) })
+            is ContentResolvedLink -> copy(children.map { it.dfs(fqName, node) })
+            is ContentEmbeddedResource -> copy(children.map { it.dfs(fqName, node) })
+            is ContentTable -> copy(children = children.map { it.dfs(fqName, node) as ContentGroup })
+            is ContentList -> copy(children.map { it.dfs(fqName, node) })
+            is ContentGroup -> copy(children.map { it.dfs(fqName, node) })
+            is PlatformHintedContent -> copy(inner.dfs(fqName, node))
+            is ContentText -> if (text == fqName) node else this
+            is ContentBreakLine -> this
+            else -> this.also { context.logger.error("Could not recognize $this ContentNode in SamplesTransformer") }
+        }
+    }
+
+    private fun contentCode(
+        sourceSets: Set<DisplaySourceSet>,
+        dri: Set<DRI>,
+        content: String,
+        language: String,
+        styles: Set<Style> = emptySet(),
+        extra: PropertyContainer<ContentNode> = PropertyContainer.empty()
+    ) =
+        ContentCodeBlock(
+            children = listOf(
+                ContentText(
+                    text = content,
+                    dci = DCI(dri, ContentKind.Sample),
+                    sourceSets = sourceSets,
+                    style = emptySet(),
+                    extra = PropertyContainer.empty()
+                )
+            ),
+            language = language,
+            dci = DCI(dri, ContentKind.Sample),
+            sourceSets = sourceSets,
+            style = styles + ContentStyle.RunnableSample + TextStyle.Monospace,
+            extra = extra
+        )
+}
diff --git a/plugins/base/src/test/kotlin/content/samples/ContentForSamplesTest.kt b/plugins/base/src/test/kotlin/content/samples/ContentForSamplesTest.kt
index 37009e4622..9b5ab7ad3c 100644
--- a/plugins/base/src/test/kotlin/content/samples/ContentForSamplesTest.kt
+++ b/plugins/base/src/test/kotlin/content/samples/ContentForSamplesTest.kt
@@ -2,12 +2,15 @@ package content.samples
 
 import matchers.content.*
 import org.jetbrains.dokka.base.testApi.testRunner.BaseAbstractTest
+import org.jetbrains.dokka.base.transformers.pages.KOTLIN_PLAYGROUND_SCRIPT
 import org.jetbrains.dokka.model.DisplaySourceSet
 import org.junit.jupiter.api.Test
+import utils.TestOutputWriterPlugin
 import utils.classSignature
 import utils.findTestType
 import java.nio.file.Paths
 import kotlin.test.assertEquals
+import kotlin.test.assertNotEquals
 
 class ContentForSamplesTest : BaseAbstractTest() {
     private val testDataDir = getTestDataDir("content/samples").toAbsolutePath()
@@ -61,6 +64,7 @@ class ContentForSamplesTest : BaseAbstractTest() {
 
     @Test
     fun `samples block is rendered in the description`() {
+        val writerPlugin = TestOutputWriterPlugin()
         testInline(
             """
             |/src/main/kotlin/test/source.kt
@@ -70,10 +74,12 @@ class ContentForSamplesTest : BaseAbstractTest() {
             | * @sample [test.sampleForClassDescription]
             | */
             |class Foo
-        """.trimIndent(), testConfiguration
+        """.trimIndent(), testConfiguration,
+            pluginOverrides = listOf(writerPlugin)
         ) {
             pagesTransformationStage = { module ->
                 val page = module.findTestType("test", "Foo")
+                assert(KOTLIN_PLAYGROUND_SCRIPT in page.embeddedResources)
                 page.content.assertNode {
                     group {
                         header(1) { +"Foo" }
@@ -101,6 +107,9 @@ class ContentForSamplesTest : BaseAbstractTest() {
                     skipAllNotMatching()
                 }
             }
+            renderingStage = { _, _ ->
+                assertNotEquals(-1, writerPlugin.writer.contents["root/test/-foo/index.html"]?.indexOf(KOTLIN_PLAYGROUND_SCRIPT))
+            }
         }
     }
 
@@ -134,6 +143,7 @@ class ContentForSamplesTest : BaseAbstractTest() {
         ) {
             pagesTransformationStage = { module ->
                 val page = module.findTestType("pageMerger", "Parent")
+                assert(KOTLIN_PLAYGROUND_SCRIPT in page.embeddedResources)
                 page.content.assertNode {
                     group {
                         header(1) { +"Parent" }
diff --git a/subprojects/analysis-kotlin-api/api/analysis-kotlin-api.api b/subprojects/analysis-kotlin-api/api/analysis-kotlin-api.api
index a057470ea6..c65dfe5af8 100644
--- a/subprojects/analysis-kotlin-api/api/analysis-kotlin-api.api
+++ b/subprojects/analysis-kotlin-api/api/analysis-kotlin-api.api
@@ -51,6 +51,7 @@ public final class org/jetbrains/dokka/analysis/kotlin/internal/InternalKotlinAn
 	public final fun getInheritanceBuilder ()Lorg/jetbrains/dokka/plugability/ExtensionPoint;
 	public final fun getKotlinToJavaService ()Lorg/jetbrains/dokka/plugability/ExtensionPoint;
 	public final fun getModuleAndPackageDocumentationReader ()Lorg/jetbrains/dokka/plugability/ExtensionPoint;
+	public final fun getSampleProviderFactory ()Lorg/jetbrains/dokka/plugability/ExtensionPoint;
 	public final fun getSyntheticDocumentableDetector ()Lorg/jetbrains/dokka/plugability/ExtensionPoint;
 }
 
@@ -64,6 +65,20 @@ public abstract interface class org/jetbrains/dokka/analysis/kotlin/internal/Mod
 	public abstract fun read (Lorg/jetbrains/dokka/model/DPackage;)Ljava/util/Map;
 }
 
+public abstract interface class org/jetbrains/dokka/analysis/kotlin/internal/SampleProvider : java/lang/AutoCloseable {
+	public abstract fun getSample (Lorg/jetbrains/dokka/DokkaConfiguration$DokkaSourceSet;Ljava/lang/String;)Lorg/jetbrains/dokka/analysis/kotlin/internal/SampleProvider$SampleSnippet;
+}
+
+public final class org/jetbrains/dokka/analysis/kotlin/internal/SampleProvider$SampleSnippet {
+	public fun <init> (Ljava/lang/String;Ljava/lang/String;)V
+	public final fun getBody ()Ljava/lang/String;
+	public final fun getImports ()Ljava/lang/String;
+}
+
+public abstract interface class org/jetbrains/dokka/analysis/kotlin/internal/SampleProviderFactory {
+	public abstract fun build ()Lorg/jetbrains/dokka/analysis/kotlin/internal/SampleProvider;
+}
+
 public abstract interface class org/jetbrains/dokka/analysis/kotlin/internal/SyntheticDocumentableDetector {
 	public abstract fun isSynthetic (Lorg/jetbrains/dokka/model/Documentable;Lorg/jetbrains/dokka/DokkaConfiguration$DokkaSourceSet;)Z
 }
diff --git a/subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/InternalKotlinAnalysisPlugin.kt b/subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/InternalKotlinAnalysisPlugin.kt
index 330522f361..30c2c2f8a4 100644
--- a/subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/InternalKotlinAnalysisPlugin.kt
+++ b/subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/InternalKotlinAnalysisPlugin.kt
@@ -26,6 +26,8 @@ class InternalKotlinAnalysisPlugin : DokkaPlugin() {
 
     val documentableSourceLanguageParser by extensionPoint<DocumentableSourceLanguageParser>()
 
+    val sampleProviderFactory by extensionPoint<SampleProviderFactory>()
+
     @OptIn(DokkaPluginApiPreview::class)
     override fun pluginApiPreviewAcknowledgement(): PluginApiPreviewAcknowledgement = PluginApiPreviewAcknowledgement
 }
diff --git a/subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/SampleProvider.kt b/subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/SampleProvider.kt
new file mode 100644
index 0000000000..195f3c3561
--- /dev/null
+++ b/subprojects/analysis-kotlin-api/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/internal/SampleProvider.kt
@@ -0,0 +1,30 @@
+package org.jetbrains.dokka.analysis.kotlin.internal
+
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.InternalDokkaApi
+
+@InternalDokkaApi
+interface SampleProviderFactory {
+    /**
+     * [SampleProvider] is a short-lived closeable instance.
+     * It assumes that [SampleProvider] scope of use is not big.
+     * Otherwise, it can lead to high memory consumption / leaks during Dokka running.
+     */
+    fun build(): SampleProvider
+}
+
+/**
+ * It is closeable.
+ * Otherwise, there is a chance of high memory consumption / leak.
+ * In general case, it creates a separate project to analysis samples directories.
+ */
+@InternalDokkaApi
+interface SampleProvider: AutoCloseable {
+    class SampleSnippet(val imports: String, val body:String)
+
+
+    /**
+     * @return [SampleSnippet] or null if it has not found by [fqLink]
+     */
+    fun getSample(sourceSet: DokkaConfiguration.DokkaSourceSet, fqLink: String): SampleSnippet?
+}
\ No newline at end of file
diff --git a/subprojects/analysis-kotlin-descriptors/compiler/api/compiler.api b/subprojects/analysis-kotlin-descriptors/compiler/api/compiler.api
index 7d53418ff4..470def08ab 100644
--- a/subprojects/analysis-kotlin-descriptors/compiler/api/compiler.api
+++ b/subprojects/analysis-kotlin-descriptors/compiler/api/compiler.api
@@ -66,3 +66,18 @@ public abstract class org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/c
 	public final fun get (Lorg/jetbrains/dokka/DokkaConfiguration$DokkaSourceSet;)Lorg/jetbrains/dokka/analysis/kotlin/descriptors/compiler/configuration/AnalysisContext;
 }
 
+public class org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/impl/KotlinSampleProvider : org/jetbrains/dokka/analysis/kotlin/internal/SampleProvider {
+	public fun <init> (Lorg/jetbrains/dokka/plugability/DokkaContext;)V
+	public fun close ()V
+	public final fun getContext ()Lorg/jetbrains/dokka/plugability/DokkaContext;
+	public fun getSample (Lorg/jetbrains/dokka/DokkaConfiguration$DokkaSourceSet;Ljava/lang/String;)Lorg/jetbrains/dokka/analysis/kotlin/internal/SampleProvider$SampleSnippet;
+	protected fun processBody (Lcom/intellij/psi/PsiElement;)Ljava/lang/String;
+	protected fun processImports (Lcom/intellij/psi/PsiElement;)Ljava/lang/String;
+}
+
+public final class org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/impl/KotlinSampleProviderFactory : org/jetbrains/dokka/analysis/kotlin/internal/SampleProviderFactory {
+	public fun <init> (Lorg/jetbrains/dokka/plugability/DokkaContext;)V
+	public fun build ()Lorg/jetbrains/dokka/analysis/kotlin/internal/SampleProvider;
+	public final fun getContext ()Lorg/jetbrains/dokka/plugability/DokkaContext;
+}
+
diff --git a/subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/CompilerDescriptorAnalysisPlugin.kt b/subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/CompilerDescriptorAnalysisPlugin.kt
index 6d63fe14ba..2d79affa92 100644
--- a/subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/CompilerDescriptorAnalysisPlugin.kt
+++ b/subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/CompilerDescriptorAnalysisPlugin.kt
@@ -21,6 +21,7 @@ import org.jetbrains.dokka.renderers.PostAction
 import org.jetbrains.dokka.analysis.kotlin.internal.InternalKotlinAnalysisPlugin
 import org.jetbrains.kotlin.asJava.elements.KtLightAbstractAnnotation
 
+@Suppress("unused")
 @InternalDokkaApi
 class CompilerDescriptorAnalysisPlugin : DokkaPlugin() {
 
@@ -42,7 +43,7 @@ class CompilerDescriptorAnalysisPlugin : DokkaPlugin() {
         plugin<InternalKotlinAnalysisPlugin>().documentableSourceLanguageParser providing { CompilerDocumentableSourceLanguageParser() }
     }
 
-    internal  val defaultKotlinAnalysis by extending {
+    internal val defaultKotlinAnalysis by extending {
         kotlinAnalysis providing { ctx ->
             ProjectKotlinAnalysis(
                 sourceSets = ctx.configuration.sourceSets,
@@ -51,7 +52,7 @@ class CompilerDescriptorAnalysisPlugin : DokkaPlugin() {
         }
     }
 
-    internal  val descriptorToDocumentableTranslator by extending {
+    internal val descriptorToDocumentableTranslator by extending {
         CoreExtensions.sourceToDocumentableTranslator providing ::DefaultDescriptorToDocumentableTranslator
     }
 
@@ -63,6 +64,10 @@ class CompilerDescriptorAnalysisPlugin : DokkaPlugin() {
         plugin<InternalKotlinAnalysisPlugin>().fullClassHierarchyBuilder providing { DescriptorFullClassHierarchyBuilder() }
     }
 
+    internal val kotlinSampleProviderFactory by extending {
+        plugin<InternalKotlinAnalysisPlugin>().sampleProviderFactory providing ::KotlinSampleProviderFactory
+    }
+
     internal val descriptorSyntheticDocumentableDetector by extending {
         plugin<InternalKotlinAnalysisPlugin>().syntheticDocumentableDetector providing { DescriptorSyntheticDocumentableDetector() }
     }
diff --git a/subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/impl/KotlinSampleProvider.kt b/subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/impl/KotlinSampleProvider.kt
new file mode 100644
index 0000000000..73f70b9cc5
--- /dev/null
+++ b/subprojects/analysis-kotlin-descriptors/compiler/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/descriptors/compiler/impl/KotlinSampleProvider.kt
@@ -0,0 +1,109 @@
+package org.jetbrains.dokka.analysis.kotlin.descriptors.compiler.impl
+
+import com.intellij.psi.PsiElement
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runBlocking
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.InternalDokkaApi
+import org.jetbrains.dokka.analysis.kotlin.descriptors.compiler.CompilerDescriptorAnalysisPlugin
+import org.jetbrains.dokka.analysis.kotlin.descriptors.compiler.KDocFinder
+import org.jetbrains.dokka.analysis.kotlin.descriptors.compiler.configuration.SamplesKotlinAnalysis
+import org.jetbrains.dokka.plugability.DokkaContext
+import org.jetbrains.dokka.plugability.plugin
+import org.jetbrains.dokka.plugability.querySingle
+import org.jetbrains.dokka.analysis.kotlin.internal.SampleProvider
+import org.jetbrains.dokka.analysis.kotlin.internal.SampleProviderFactory
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtDeclarationWithBody
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.lazy.ResolveSession
+
+class KotlinSampleProviderFactory(val context: DokkaContext): SampleProviderFactory {
+    override fun build(): SampleProvider {
+        return KotlinSampleProvider(context)
+    }
+
+}
+/**
+ * It's declared as open since StdLib has its own sample transformer
+ * with [processBody] and [processImports]
+ */
+@InternalDokkaApi
+open class KotlinSampleProvider(val context: DokkaContext): SampleProvider {
+    private val kDocFinder: KDocFinder = context.plugin<CompilerDescriptorAnalysisPlugin>().querySingle { kdocFinder }
+    private val analysis = lazy {
+        /**
+         * Run from the thread of [Dispatchers.Default]. It can help to avoid a memory leaks in `ThreadLocal`s (that keep `URLCLassLoader`)
+         * since we shut down Dispatchers.Default at the end of each task (see [org.jetbrains.dokka.DokkaConfiguration.finalizeCoroutines]).
+         * Currently, all `ThreadLocal`s are in a compiler/IDE codebase.
+         */
+        runBlocking(Dispatchers.Default) {
+            SamplesKotlinAnalysis(
+                sourceSets = context.configuration.sourceSets,
+                context = context,
+                projectKotlinAnalysis = context.plugin<CompilerDescriptorAnalysisPlugin>()
+                    .querySingle { kotlinAnalysis }
+            )
+        }
+    }
+    protected open fun processBody(psiElement: PsiElement): String {
+        val text = processSampleBody(psiElement).trim { it == '\n' || it == '\r' }.trimEnd()
+        val lines = text.split("\n")
+        val indent = lines.filter(String::isNotBlank).map { it.takeWhile(Char::isWhitespace).count() }.minOrNull() ?: 0
+        return lines.joinToString("\n") { it.drop(indent) }
+    }
+
+    private fun processSampleBody(psiElement: PsiElement): String = when (psiElement) {
+        is KtDeclarationWithBody -> {
+            when (val bodyExpression = psiElement.bodyExpression) {
+                is KtBlockExpression -> bodyExpression.text.removeSurrounding("{", "}")
+                else -> bodyExpression!!.text
+            }
+        }
+        else -> psiElement.text
+    }
+
+    protected open fun processImports(psiElement: PsiElement): String {
+        val psiFile = psiElement.containingFile
+        return when(val text = (psiFile as? KtFile)?.importList?.text) {
+            is String -> text
+            else -> ""
+        }
+    }
+
+    /**
+     * @return [SampleProvider.SampleSnippet] or null if it has not found by [fqLink]
+     */
+    override fun getSample(sourceSet: DokkaConfiguration.DokkaSourceSet, fqLink: String): SampleProvider.SampleSnippet? {
+        return runBlocking(Dispatchers.Default) {
+                val resolveSession = analysis.value[sourceSet].resolveSession
+                val psiElement = fqNameToPsiElement(resolveSession, fqLink, sourceSet)
+                    ?: return@runBlocking null.also { context.logger.warn("Cannot find PsiElement corresponding to $fqLink") }
+                val imports =
+                    processImports(psiElement)
+                val body = processBody(psiElement)
+                return@runBlocking SampleProvider.SampleSnippet(imports, body)
+        }
+    }
+    override fun close() {
+        if(analysis.isInitialized())
+            analysis.value.close()
+    }
+
+    private fun fqNameToPsiElement(resolveSession: ResolveSession, functionName: String, dokkaSourceSet: DokkaConfiguration.DokkaSourceSet): PsiElement? {
+        val packageName = functionName.takeWhile { it != '.' }
+        val descriptor = resolveSession.getPackageFragment(FqName(packageName))
+            ?: return null.also { context.logger.warn("Cannot find descriptor for package $packageName") }
+
+        with (kDocFinder) {
+            val symbol = resolveKDocLink(
+                descriptor,
+                functionName,
+                dokkaSourceSet,
+                emptyBindingContext = true
+            ).firstOrNull() ?: return null.also { context.logger.warn("Unresolved function $functionName in @sample") }
+            return org.jetbrains.kotlin.resolve.DescriptorToSourceUtils.descriptorToDeclaration(symbol)
+        }
+    }
+}
\ No newline at end of file