[K2] Fix sample resolving (#3361)

* [K2] Fix sample

Fallback to default roots of the source set even if sample roots are assigned

* [K2] Resolve java samples as well

`sourcePsiSafe` does not support Java KT-53669

Author: vmishenev

dokka-subprojects/analysis-kotlin-api/src/test/kotlin/org/jetbrains/dokka/analysis/test/sample/SampleAnalysisTest.kt

@@ -10,7 +10,6 @@ import org.jetbrains.dokka.analysis.test.api.mixedJvmTestProject
 import org.jetbrains.dokka.analysis.test.api.useServices
 import org.jetbrains.dokka.analysis.test.api.util.CollectingDokkaConsoleLogger
 import org.jetbrains.dokka.analysis.test.api.util.singleSourceSet
-import org.junit.jupiter.api.Tag
 import kotlin.test.*
 
 class SampleAnalysisTest {
@@ -75,7 +74,6 @@ class SampleAnalysisTest {
     }
 
     @Test
-    @Tag("onlyDescriptors") // TODO #3359
     fun `should resolve a valid sample if set via the additionalSourceRoots option`() {
         val testProject = kotlinJvmTestProject {
             dokkaConfiguration {
@@ -135,7 +133,6 @@ class SampleAnalysisTest {
     }
 
     @Test
-    @Tag("onlyDescriptors") // TODO #3359
     fun `should resolve a valid sample function that exists in the main source set`() {
         val testProject = kotlinJvmTestProject {
             ktFile("org/jetbrains/dokka/test/MyKotlinFile.kt") {
@@ -325,7 +322,6 @@ class SampleAnalysisTest {
     }
 
     @Test
-    @Tag("onlyDescriptors") // TODO #3359
     fun `should return null if trying to resolve a non-kotlin sample link`() {
         val testProject = mixedJvmTestProject {
             kotlinSourceDirectory {
@@ -366,7 +362,6 @@ class SampleAnalysisTest {
     }
 
     @Test
-    @Tag("onlyDescriptors") // TODO #3359
     fun `should filter out empty import statement lines`() {
         val testProject = kotlinJvmTestProject {
             ktFile("org/jetbrains/dokka/test/MyKotlinFile.kt") {
@@ -403,7 +398,6 @@ class SampleAnalysisTest {
     }
 
     @Test
-    @Tag("onlyDescriptors") // TODO #3359
     fun `should return an empty list of imports if sample file has none`() {
         val testProject = kotlinJvmTestProject {
             ktFile("org/jetbrains/dokka/test/MyKotlinFile.kt") {
@@ -430,7 +424,6 @@ class SampleAnalysisTest {
     }
 
     @Test
-    @Tag("onlyDescriptors") // TODO #3359
     fun `should filter out leading and trailing line breaks`() {
         val testProject = kotlinJvmTestProject {
             ktFile("org/jetbrains/dokka/test/MyKotlinFile.kt") {
@@ -459,7 +452,6 @@ class SampleAnalysisTest {
     }
 
     @Test
-    @Tag("onlyDescriptors") // TODO #3359
     fun `should filter out trailing whitespace`() {
         val testProject = kotlinJvmTestProject {
             ktFile("org/jetbrains/dokka/test/MyKotlinFile.kt") {

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()
         }
     }