Stabilize Sample analysis API (#3195)

8 files changed, 154 insertions, 114 deletions

diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/KDocProvider.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/KDocProvider.kt
index d8a4e476..36492cdd 100644
--- a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/KDocProvider.kt
+++ b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/KDocProvider.kt
@@ -64,7 +64,7 @@ internal fun KtAnalysisSession.getKDocDocumentationFrom(symbol: KtSymbol, logger
 
     parseFromKDocTag(
         kDocTag = kDocContent.contentTag,
-        externalDri = { link -> resolveKDocLink(link).ifUnresolved { logger.logUnresolvedLink(link.getLinkText(), kdocLocation) } },
+        externalDri = { link -> resolveKDocLinkDRI(link).ifUnresolved { logger.logUnresolvedLink(link.getLinkText(), kdocLocation) } },
         kdocLocation = kdocLocation
     )
 }
diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/ResolveKDocLink.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/ResolveKDocLink.kt
index 9a0b81bd..c719c855 100644
--- a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/ResolveKDocLink.kt
+++ b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/ResolveKDocLink.kt
@@ -9,6 +9,7 @@ import org.jetbrains.dokka.analysis.kotlin.symbols.translators.getDRIFromSymbol
 import org.jetbrains.dokka.links.DRI
 import org.jetbrains.dokka.utilities.DokkaLogger
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.kdoc.psi.api.KDoc
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocLink
@@ -34,7 +35,23 @@ internal inline fun DRI?.ifUnresolved(action: () -> Unit): DRI? = this ?: run {
  *
  * @return [DRI] or null if the [link] is unresolved
  */
-internal fun KtAnalysisSession.resolveKDocTextLink(link: String, context: PsiElement? = null): DRI? {
+internal fun KtAnalysisSession.resolveKDocTextLinkDRI(link: String, context: PsiElement? = null): DRI? {
+    val kDocLink = createKDocLink(link, context)
+    return kDocLink?.let { resolveKDocLinkDRI(it) }
+}
+
+/**
+ * If the [link] is ambiguous, i.e. leads to more than one declaration,
+ * it returns deterministically any declaration.
+ *
+ * @return [KtSymbol] or null if the [link] is unresolved
+ */
+internal fun KtAnalysisSession.resolveKDocTextLinkSymbol(link: String, context: PsiElement? = null): KtSymbol? {
+    val kDocLink = createKDocLink(link, context)
+    return kDocLink?.let { resolveToSymbol(it) }
+}
+
+private fun KtAnalysisSession.createKDocLink(link: String, context: PsiElement? = null): KDocLink? {
     val psiFactory = context?.let { KtPsiFactory.contextual(it) } ?: KtPsiFactory(this.useSiteModule.project)
     val kDoc = psiFactory.createComment(
         """
@@ -43,8 +60,8 @@ internal fun KtAnalysisSession.resolveKDocTextLink(link: String, context: PsiEle
     */
  """.trimIndent()
     ) as? KDoc
-    val kDocLink = kDoc?.getDefaultSection()?.children?.filterIsInstance<KDocLink>()?.singleOrNull()
-    return kDocLink?.let { resolveKDocLink(it) }
+
+    return kDoc?.getDefaultSection()?.children?.filterIsInstance<KDocLink>()?.singleOrNull()
 }
 
 /**
@@ -53,9 +70,13 @@ internal fun KtAnalysisSession.resolveKDocTextLink(link: String, context: PsiEle
  *
  * @return [DRI] or null if the [link] is unresolved
  */
-internal fun KtAnalysisSession.resolveKDocLink(link: KDocLink): DRI? {
-    val lastNameSegment = link.children.filterIsInstance<KDocName>().lastOrNull()
-    val linkedSymbol = lastNameSegment?.mainReference?.resolveToSymbols()?.firstOrNull()
+internal fun KtAnalysisSession.resolveKDocLinkDRI(link: KDocLink): DRI? {
+    val linkedSymbol = resolveToSymbol(link)
     return if (linkedSymbol == null) null
     else getDRIFromSymbol(linkedSymbol)
 }
+
+private fun KtAnalysisSession.resolveToSymbol(kDocLink: KDocLink): KtSymbol? {
+    val lastNameSegment = kDocLink.children.filterIsInstance<KDocName>().lastOrNull()
+    return lastNameSegment?.mainReference?.resolveToSymbols()?.firstOrNull()
+}
diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/SyntheticKDocProvider.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/SyntheticKDocProvider.kt
index 4622ffd8..9f2e6873 100644
--- a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/SyntheticKDocProvider.kt
+++ b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/SyntheticKDocProvider.kt
@@ -54,7 +54,7 @@ internal fun KtAnalysisSession.getGeneratedKDocDocumentationFrom(symbol: KtSymbo
 private fun KtAnalysisSession.loadTemplate(filePath: String): DocumentationNode? {
     val kdoc = loadContent(filePath) ?: return null
     val externalDriProvider = { link: String ->
-        resolveKDocTextLink(link)
+        resolveKDocTextLinkDRI(link)
     }
 
     val parser = MarkdownParser(externalDriProvider, filePath)
diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt
index 0ee95e45..846ee61d 100644
--- a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt
+++ b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt
@@ -11,7 +11,7 @@ import org.jetbrains.dokka.analysis.java.parsers.DocCommentParser
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.*
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.logUnresolvedLink
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.parseFromKDocTag
-import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.resolveKDocLink
+import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.resolveKDocLinkDRI
 import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.SymbolsAnalysisPlugin
 import org.jetbrains.dokka.model.doc.DocumentationNode
 import org.jetbrains.dokka.plugability.DokkaContext
@@ -42,7 +42,7 @@ internal class KotlinDocCommentParser(
         return analyze(kotlinAnalysis.getModule(sourceSet)) {
             parseFromKDocTag(
                 kDocTag = element.comment,
-                externalDri = { link -> resolveKDocLink(link).ifUnresolved { context.logger.logUnresolvedLink(link.getLinkText(), elementName) } },
+                externalDri = { link -> resolveKDocLinkDRI(link).ifUnresolved { context.logger.logUnresolvedLink(link.getLinkText(), elementName) } },
                 kdocLocation = null,
                 parseWithChildren = parseWithChildren
             )
diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt
index f5cfbdb9..1fdebe24 100644
--- a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt
+++ b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt
@@ -10,7 +10,7 @@ import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.logUnresolvedLink
 import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.KotlinAnalysis
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.moduledocs.ModuleAndPackageDocumentation.Classifier.Module
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.moduledocs.ModuleAndPackageDocumentation.Classifier.Package
-import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.resolveKDocTextLink
+import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.resolveKDocTextLinkDRI
 import org.jetbrains.dokka.analysis.markdown.jb.MarkdownParser
 import org.jetbrains.dokka.model.doc.DocumentationNode
 import org.jetbrains.dokka.utilities.DokkaLogger
@@ -47,7 +47,7 @@ internal fun ModuleAndPackageDocumentationParsingContext(
         MarkdownParser(
             externalDri = { link ->
                 analyze(sourceModule) {
-                    resolveKDocTextLink(
+                    resolveKDocTextLinkDRI(
                         link,
                         contextPsi
                     ).ifUnresolved {
diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/SymbolsAnalysisPlugin.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/SymbolsAnalysisPlugin.kt
index 122e0b10..23bb0dc5 100644
--- a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/SymbolsAnalysisPlugin.kt
+++ b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/SymbolsAnalysisPlugin.kt
@@ -9,6 +9,7 @@ import com.intellij.psi.PsiAnnotation
 import org.jetbrains.dokka.CoreExtensions
 import org.jetbrains.dokka.analysis.java.BreakingAbstractionKotlinLightMethodChecker
 import org.jetbrains.dokka.analysis.java.JavaAnalysisPlugin
+import org.jetbrains.dokka.analysis.kotlin.KotlinAnalysisPlugin
 import org.jetbrains.dokka.analysis.kotlin.internal.InternalKotlinAnalysisPlugin
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.java.KotlinInheritDocTagContentProvider
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.java.DescriptorKotlinDocCommentCreator
@@ -117,8 +118,8 @@ public class SymbolsAnalysisPlugin : DokkaPlugin() {
         plugin<InternalKotlinAnalysisPlugin>().externalDocumentablesProvider providing ::SymbolExternalDocumentablesProvider
     }
 
-    internal val kotlinSampleProviderFactory by extending {
-        plugin<InternalKotlinAnalysisPlugin>().sampleProviderFactory providing ::KotlinSampleProviderFactory
+    internal val symbolSampleAnalysisEnvironmentCreator  by extending {
+        plugin<KotlinAnalysisPlugin>().sampleAnalysisEnvironmentCreator providing ::SymbolSampleAnalysisEnvironmentCreator
     }
 
     @OptIn(DokkaPluginApiPreview::class)
diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinSampleProvider.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinSampleProvider.kt
deleted file mode 100644
index e453c72d..00000000
--- a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinSampleProvider.kt
+++ /dev/null
@@ -1,100 +0,0 @@
-/*
- * Copyright 2014-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
- */
-
-package org.jetbrains.dokka.analysis.kotlin.symbols.services
-
-import com.intellij.psi.PsiElement
-import org.jetbrains.dokka.DokkaConfiguration
-import org.jetbrains.dokka.InternalDokkaApi
-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.dokka.analysis.kotlin.symbols.plugin.SamplesKotlinAnalysis
-import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.SymbolsAnalysisPlugin
-import org.jetbrains.kotlin.analysis.api.analyze
-import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtBlockExpression
-import org.jetbrains.kotlin.psi.KtDeclarationWithBody
-import org.jetbrains.kotlin.psi.KtFile
-
-public class KotlinSampleProviderFactory(
-    public 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
-public open class KotlinSampleProvider(
-    public val context: DokkaContext
-): SampleProvider {
-    private val kotlinAnalysisOfRegularSources = context.plugin<SymbolsAnalysisPlugin>().querySingle { kotlinAnalysis }
-
-    private val kotlinAnalysisOfSamples = SamplesKotlinAnalysis(
-        sourceSets = context.configuration.sourceSets, context = context
-    )
-
-    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).minOfOrNull { it.takeWhile(Char::isWhitespace).count() } ?: 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 kotlinAnalysisOfSamples.getModuleOrNull(sourceSet)?.let { getSampleFromModule(it, fqLink) }
-            ?: getSampleFromModule(
-                kotlinAnalysisOfRegularSources.getModule(sourceSet), fqLink
-            )
-    }
-    private fun getSampleFromModule(module: KtSourceModule, fqLink: String): SampleProvider.SampleSnippet? {
-        val psiElement = analyze(module) {
-            val lastDotIndex = fqLink.lastIndexOf('.')
-
-            val functionName = if (lastDotIndex == -1) fqLink else fqLink.substring(lastDotIndex + 1, fqLink.length)
-            val packageName = if (lastDotIndex == -1) "" else fqLink.substring(0, lastDotIndex)
-            getTopLevelCallableSymbols(FqName(packageName), Name.identifier(functionName)).firstOrNull()?.psi
-        }
-            ?: return null.also { context.logger.warn("Cannot find PsiElement corresponding to $fqLink") }
-        val imports =
-            processImports(psiElement)
-        val body = processBody(psiElement)
-
-        return SampleProvider.SampleSnippet(imports, body)
-    }
-
-    override fun close() {
-        kotlinAnalysisOfSamples.close()
-    }
-}
diff --git a/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolSampleAnalysisEnvironment.kt b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolSampleAnalysisEnvironment.kt
new file mode 100644
index 00000000..37095df5
--- /dev/null
+++ b/dokka-subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolSampleAnalysisEnvironment.kt
@@ -0,0 +1,118 @@
+/*
+ * Copyright 2014-2023 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package org.jetbrains.dokka.analysis.kotlin.symbols.services
+
+import com.intellij.psi.PsiElement
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runBlocking
+import org.jetbrains.dokka.DokkaConfiguration.DokkaSourceSet
+import org.jetbrains.dokka.analysis.kotlin.sample.SampleAnalysisEnvironment
+import org.jetbrains.dokka.analysis.kotlin.sample.SampleAnalysisEnvironmentCreator
+import org.jetbrains.dokka.analysis.kotlin.sample.SampleSnippet
+import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.resolveKDocTextLinkSymbol
+import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.KotlinAnalysis
+import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.SamplesKotlinAnalysis
+import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.SymbolsAnalysisPlugin
+import org.jetbrains.dokka.plugability.DokkaContext
+import org.jetbrains.dokka.plugability.plugin
+import org.jetbrains.dokka.plugability.querySingle
+import org.jetbrains.dokka.utilities.DokkaLogger
+import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.api.symbols.sourcePsiSafe
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtDeclarationWithBody
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtFunction
+
+internal class SymbolSampleAnalysisEnvironmentCreator(
+    private val context: DokkaContext,
+) : SampleAnalysisEnvironmentCreator {
+
+    private val projectKotlinAnalysis = context.plugin<SymbolsAnalysisPlugin>().querySingle { kotlinAnalysis }
+
+    override fun <T> use(block: SampleAnalysisEnvironment.() -> T): T {
+        return runBlocking(Dispatchers.Default) {
+            SamplesKotlinAnalysis(
+                sourceSets = context.configuration.sourceSets,
+                context = context
+            ).use { samplesKotlinAnalysis ->
+                val sampleAnalysisEnvironment = SymbolSampleAnalysisEnvironment(
+                    samplesKotlinAnalysis = samplesKotlinAnalysis,
+                    projectKotlinAnalysis = projectKotlinAnalysis,
+                    dokkaLogger = context.logger
+                )
+                block(sampleAnalysisEnvironment)
+            }
+        }
+    }
+}
+
+private class SymbolSampleAnalysisEnvironment(
+    private val samplesKotlinAnalysis: KotlinAnalysis,
+    private val projectKotlinAnalysis: KotlinAnalysis,
+    private val dokkaLogger: DokkaLogger,
+) : SampleAnalysisEnvironment {
+
+    override fun resolveSample(sourceSet: DokkaSourceSet, fullyQualifiedLink: String): SampleSnippet? {
+        val psiElement = findPsiElement(sourceSet, fullyQualifiedLink)
+        if (psiElement == null) {
+            dokkaLogger.warn(
+                "Unable to resolve a @sample link: \"$fullyQualifiedLink\". Is it used correctly? " +
+                        "Expecting a link to a reachable (resolvable) top-level Kotlin function."
+            )
+            return null
+        } else if (psiElement.language != KotlinLanguage.INSTANCE) {
+            dokkaLogger.warn("Unable to resolve non-Kotlin @sample links: \"$fullyQualifiedLink\"")
+            return null
+        } else if (psiElement !is KtFunction) {
+            dokkaLogger.warn("Unable to process a @sample link: \"$fullyQualifiedLink\". Only function links allowed.")
+            return null
+        }
+
+        val imports = processImports(psiElement)
+        val body = processBody(psiElement)
+
+        return SampleSnippet(imports, body)
+    }
+
+    private fun findPsiElement(sourceSet: DokkaSourceSet, fqLink: String): PsiElement? {
+        val ktSourceModule = samplesKotlinAnalysis.getModuleOrNull(sourceSet)
+            ?: projectKotlinAnalysis.getModule(sourceSet)
+
+        return analyze(ktSourceModule) {
+            resolveKDocTextLinkSymbol(fqLink)
+                ?.sourcePsiSafe()
+        }
+    }
+
+    private fun processImports(psiElement: PsiElement): List<String> {
+        val psiFile = psiElement.containingFile
+        val importsList = (psiFile as? KtFile)?.importList ?: return emptyList()
+        return importsList.imports
+            .map { it.text.removePrefix("import ") }
+            .filter { it.isNotBlank() }
+    }
+
+    private fun processBody(sampleElement: PsiElement): String {
+        return getSampleBody(sampleElement)
+            .trim { it == '\n' || it == '\r' }
+            .trimEnd()
+            .trimIndent()
+    }
+
+    private fun getSampleBody(sampleElement: PsiElement): String {
+        return when (sampleElement) {
+            is KtDeclarationWithBody -> {
+                when (val bodyExpression = sampleElement.bodyExpression) {
+                    is KtBlockExpression -> bodyExpression.text.removeSurrounding("{", "}")
+                    else -> bodyExpression!!.text
+                }
+            }
+
+            else -> sampleElement.text
+        }
+    }
+}