From ce4a5d8e435e0864b84b3bde69e918850b63bb64 Mon Sep 17 00:00:00 2001
From: vmishenev <vad-mishenev@yandex.ru>
Date: Wed, 22 Nov 2023 20:55:09 +0200
Subject: [PATCH] Apply review suggestions

---
 .../SymbolSampleAnalysisEnvironment.kt        | 32 +++++++------------
 1 file changed, 11 insertions(+), 21 deletions(-)

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
index 86518ccc45..e784aabcee 100644
--- 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
@@ -79,18 +79,6 @@ private class SymbolSampleAnalysisEnvironment(
         return SampleSnippet(imports, body)
     }
 
-    // copy-pasted from StdLib 1.5
-    // TODO use it from Stdlib
-     private inline fun <T, R : Any> Iterable<T>.firstNotNullOfOrNull(transform: (T) -> R?): R? {
-        for (element in this) {
-            val result = transform(element)
-            if (result != null) {
-                return result
-            }
-        }
-        return null
-    }
-
     // TODO: remove after KT-53669 and use [org.jetbrains.kotlin.analysis.api.symbols.sourcePsiSafe] from Analysis API
     private inline fun <reified PSI : PsiElement> KtSymbol.kotlinAndJavaSourcePsiSafe(): PSI? {
         // TODO: support Java sources after KT-53669
@@ -114,15 +102,17 @@ private class SymbolSampleAnalysisEnvironment(
     }
 
     private fun findPsiElement(sourceSet: DokkaSourceSet, fqLink: String): PsiElement? {
-        // fallback to default roots of the source set even if sample roots are assigned
-        val ktSourceModules =
-            listOfNotNull(samplesKotlinAnalysis.getModuleOrNull(sourceSet), projectKotlinAnalysis.getModule(sourceSet))
-
-        return ktSourceModules.firstNotNullOfOrNull { ktSourceModule ->
-            analyze(ktSourceModule) {
-                resolveKDocTextLinkSymbol(fqLink)
-                    ?.kotlinAndJavaSourcePsiSafe()
-            }
+        // fallback to default roots of the source set even if sample roots are assigned,
+        // because `@sample` tag can contain links to functions from project sources
+        return samplesKotlinAnalysis.findPsiElement(sourceSet, fqLink)
+            ?: projectKotlinAnalysis.findPsiElement(sourceSet, fqLink)
+    }
+
+    private fun KotlinAnalysis.findPsiElement(sourceSet: DokkaSourceSet, fqLink: String): PsiElement? {
+        val ktSourceModule = this.getModuleOrNull(sourceSet) ?: return null
+        return analyze(ktSourceModule) {
+            resolveKDocTextLinkSymbol(fqLink)
+                ?.kotlinAndJavaSourcePsiSafe()
         }
     }