[K2] Fix sample

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

Author: vmishenev

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

@@ -75,7 +75,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 +134,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") {
@@ -366,7 +364,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 +400,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 +426,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 +454,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

@@ -78,13 +78,27 @@ private class SymbolSampleAnalysisEnvironment(
         return SampleSnippet(imports, body)
     }
 
+    // copy-pasted from StdLib 1.5
+     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
+    }
+
     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
+        val ktSourceModules =
+            listOfNotNull(samplesKotlinAnalysis.getModuleOrNull(sourceSet), projectKotlinAnalysis.getModule(sourceSet))
 
-        return analyze(ktSourceModule) {
-            resolveKDocTextLinkSymbol(fqLink)
-                ?.sourcePsiSafe()
+        return ktSourceModules.firstNotNullOfOrNull { ktSourceModule ->
+            analyze(ktSourceModule) {
+                resolveKDocTextLinkSymbol(fqLink)
+                    ?.sourcePsiSafe()
+            }
         }
     }