1 files changed, 32 insertions, 4 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 b4bbc2c6..c6b70d10 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
@@ -20,7 +20,8 @@ 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.analysis.api.symbols.KtSymbol
+import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtDeclarationWithBody
@@ -78,13 +79,40 @@ private class SymbolSampleAnalysisEnvironment(
         return SampleSnippet(imports, body)
     }
 
+    // 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
+        val sourcePsi = when (origin) {
+            KtSymbolOrigin.SOURCE -> this.psi
+            KtSymbolOrigin.JAVA -> this.psi
+
+            KtSymbolOrigin.SOURCE_MEMBER_GENERATED -> null
+            KtSymbolOrigin.LIBRARY -> null
+            KtSymbolOrigin.SAM_CONSTRUCTOR -> null
+            KtSymbolOrigin.INTERSECTION_OVERRIDE -> null
+            KtSymbolOrigin.SUBSTITUTION_OVERRIDE -> null
+            KtSymbolOrigin.DELEGATED -> null
+            KtSymbolOrigin.JAVA_SYNTHETIC_PROPERTY -> null
+            KtSymbolOrigin.PROPERTY_BACKING_FIELD -> null
+            KtSymbolOrigin.PLUGIN -> null
+            KtSymbolOrigin.JS_DYNAMIC -> null
+        }
+
+        return sourcePsi as? PSI
+    }
+
     private fun findPsiElement(sourceSet: DokkaSourceSet, fqLink: String): PsiElement? {
-        val ktSourceModule = samplesKotlinAnalysis.getModuleOrNull(sourceSet)
-            ?: projectKotlinAnalysis.getModule(sourceSet)
+        // 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) {
             resolveKDocTextLinkToSymbol(fqLink)
-                ?.sourcePsiSafe()
+                ?.kotlinAndJavaSourcePsiSafe()
         }
     }