[K2] Migrate to new standalone API (#3201)

* Migrate K2-Based dokka to use the new standalone mode API

see https://youtrack.jetbrains.com/issue/KT-60884

* Do not use copy-n-pasted code from Analysis API in K2-Based Dokka

see https://youtrack.jetbrains.com/issue/KT-60884

* Remove copy-n-pasted API from Analysis API in K2-Based Dokka

as it's now unused

see https://youtrack.jetbrains.com/issue/KT-60884

* Update version Analysis API to 1.9.30-dev-3330

---------

Co-authored-by: Ilya Kirillov <ilya.kirillov@jetbrains.com>

5 files changed, 34 insertions, 173 deletions

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 6acd2a57..83bfe621 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -11,7 +11,7 @@ kotlinx-bcv = "0.12.1"
 
 ## Analysis
 kotlin-compiler = "1.9.10"
-kotlin-compiler-k2 = "1.9.0-release-358"
+kotlin-compiler-k2 = "1.9.30-dev-3330"
 
 # MUST match the version of the intellij platform used in the kotlin compiler,
 # otherwise this will lead to different versions of psi API and implementations
diff --git a/subprojects/analysis-kotlin-symbols/build.gradle.kts b/subprojects/analysis-kotlin-symbols/build.gradle.kts
index c1705e9f..6733b778 100644
--- a/subprojects/analysis-kotlin-symbols/build.gradle.kts
+++ b/subprojects/analysis-kotlin-symbols/build.gradle.kts
@@ -54,7 +54,6 @@ dependencies {
     listOf(
         libs.kotlin.high.level.api.api,
         libs.kotlin.analysis.api.standalone,
-        libs.kotlin.high.level.api.impl // for Standalone prototype
     ).forEach {
         implementation(it) {
             isTransitive = false // see KTIJ-19820
@@ -67,7 +66,7 @@ dependencies {
         libs.kotlin.low.level.api.fir,
         libs.kotlin.analysis.project.structure,
         libs.kotlin.analysis.api.providers,
-        libs.kotlin.symbol.light.classes
+        libs.kotlin.symbol.light.classes,
     ).forEach {
         runtimeOnly(it) {
             isTransitive = false // see KTIJ-19820
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt
index 9ccd52b2..cf57e815 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt
@@ -125,11 +125,12 @@ internal open class EnvironmentKotlinAnalysis(
 internal interface AnalysisContext: Closeable {
     val project: Project
     val mainModule: KtSourceModule
+    val analysisSession: StandaloneAnalysisAPISession
 }
 
 private class AnalysisContextImpl(
     override val mainModule: KtSourceModule,
-    private val analysisSession: StandaloneAnalysisAPISession,
+    override val analysisSession: StandaloneAnalysisAPISession,
     private val applicationDisposable: Disposable,
     private val projectDisposable: Disposable
 ) : AnalysisContext {
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt
index 9419ec65..a6155fb0 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt
@@ -4,19 +4,11 @@
 
 package org.jetbrains.dokka.analysis.kotlin.symbols.plugin
 
-import com.intellij.core.CoreApplicationEnvironment
 import com.intellij.openapi.Disposable
-import com.intellij.openapi.project.Project
-import com.intellij.openapi.vfs.StandardFileSystems
-import com.intellij.openapi.vfs.VirtualFileManager
-import com.intellij.psi.PsiFileSystemItem
-import com.intellij.psi.PsiManager
-import com.intellij.psi.search.GlobalSearchScope
-import com.intellij.psi.search.ProjectScope
-import com.intellij.util.io.URLUtil
 import org.jetbrains.dokka.Platform
-import org.jetbrains.kotlin.analysis.api.impl.base.util.LibraryUtils
-import org.jetbrains.kotlin.analysis.api.resolve.extensions.KtResolveExtensionProvider
+import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
+import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenProvider
+import org.jetbrains.kotlin.analysis.api.standalone.KtAlwaysAccessibleLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.api.standalone.StandaloneAnalysisAPISession
 import org.jetbrains.kotlin.analysis.api.standalone.buildStandaloneAnalysisAPISession
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
@@ -24,18 +16,12 @@ import org.jetbrains.kotlin.analysis.project.structure.builder.KtModuleBuilder
 import org.jetbrains.kotlin.analysis.project.structure.builder.buildKtLibraryModule
 import org.jetbrains.kotlin.analysis.project.structure.builder.buildKtSdkModule
 import org.jetbrains.kotlin.analysis.project.structure.builder.buildKtSourceModule
-import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.config.*
-import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.platform.CommonPlatforms
 import org.jetbrains.kotlin.platform.js.JsPlatforms
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.platform.konan.NativePlatforms
-import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
-import java.io.IOException
-import java.nio.file.*
-import java.nio.file.attribute.BasicFileAttributes
 
 internal fun Platform.toTargetPlatform() = when (this) {
     Platform.js, Platform.wasm -> JsPlatforms.defaultJsPlatform
@@ -71,6 +57,7 @@ internal fun getLanguageVersionSettings(
 }
 
 // it should be changed after https://github.com/Kotlin/dokka/issues/3114
+@OptIn(KtAnalysisApiInternals::class)
 internal fun createAnalysisSession(
     classpath: List<File>,
     sourceRoots: Set<File>,
@@ -87,163 +74,40 @@ internal fun createAnalysisSession(
         projectDisposable = projectDisposable,
         withPsiDeclarationFromBinaryModuleProvider = false
     ) {
-        val project = project
+        registerProjectService(KtLifetimeTokenProvider::class.java, KtAlwaysAccessibleLifetimeTokenProvider())
         val targetPlatform = analysisPlatform.toTargetPlatform()
-        fun KtModuleBuilder.addModuleDependencies(moduleName: String) {
+
+        buildKtModuleProvider {
             val libraryRoots = classpath
-            addRegularDependency(
-                buildKtLibraryModule {
-                    contentScope = ProjectScope.getLibrariesScope(project)
-                    this.platform = targetPlatform
-                    this.project = project
-                    binaryRoots = libraryRoots.map { it.toPath() }
-                    libraryName = "Library for $moduleName"
-                }
-            )
-            getJdkHomeFromSystemProperty()?.let { jdkHome ->
-                val vfm = VirtualFileManager.getInstance()
-                val jdkHomePath = jdkHome.toPath()
-                val jdkHomeVirtualFile = vfm.findFileByNioPath(jdkHome.toPath())//vfm.findFileByPath(jdkHomePath)
-                val binaryRoots = LibraryUtils.findClassesFromJdkHome(jdkHomePath).map {
-                    Paths.get(URLUtil.extractPath(it))
-                }
+            fun KtModuleBuilder.addModuleDependencies(moduleName: String) {
                 addRegularDependency(
-                    buildKtSdkModule {
-                        contentScope = GlobalSearchScope.fileScope(project, jdkHomeVirtualFile)
+                    buildKtLibraryModule {
                         this.platform = targetPlatform
-                        this.project = project
-                        this.binaryRoots = binaryRoots
-                        sdkName = "JDK for $moduleName"
+                        addBinaryRoots(libraryRoots.map { it.toPath() })
+                        libraryName = "Library for $moduleName"
                     }
                 )
+                getJdkHomeFromSystemProperty()?.let { jdkHome ->
+                    addRegularDependency(
+                        buildKtSdkModule {
+                            this.platform = targetPlatform
+                            addBinaryRootsFromJdkHome(jdkHome.toPath(), isJre = true)
+                            sdkName = "JDK for $moduleName"
+                        }
+                    )
+                }
+            }
+            sourceModule = buildKtSourceModule {
+                languageVersionSettings = getLanguageVersionSettings(languageVersion, apiVersion)
+                platform = targetPlatform
+                moduleName = "<module>"
+                // TODO: We should handle (virtual) file changes announced via LSP with the VFS
+                addSourceRoots(sourceRoots.map { it.toPath() })
+                addModuleDependencies(moduleName)
             }
-        }
-        sourceModule = buildKtSourceModule {
-            this.languageVersionSettings = getLanguageVersionSettings(languageVersion, apiVersion)
-
-            //val fs = StandardFileSystems.local()
-            //val psiManager = PsiManager.getInstance(project)
-            // TODO: We should handle (virtual) file changes announced via LSP with the VFS
-            /*val ktFiles = sources
-                .flatMap { Files.walk(it).toList() }
-                .mapNotNull { fs.findFileByPath(it.toString()) }
-                .mapNotNull { psiManager.findFile(it) }
-                .map { it as KtFile }*/
-            val sourcePaths = sourceRoots.map { it.absolutePath }
-            val (ktFilePath, javaFilePath) = getSourceFilePaths(sourcePaths).partition { it.endsWith(KotlinFileType.EXTENSION) }
-            val javaFiles: List<PsiFileSystemItem> = getPsiFilesFromPaths(project, javaFilePath)
-            val ktFiles: List<KtFile> = getPsiFilesFromPaths(project, getSourceFilePaths(ktFilePath))
-            addSourceRoots(ktFiles + javaFiles)
-            contentScope = TopDownAnalyzerFacadeForJVM.newModuleSearchScope(project, ktFiles)
-            platform = targetPlatform
-            moduleName = "<module>"
-            this.project = project
-            addModuleDependencies(moduleName)
-        }
-
-        buildKtModuleProvider {
             platform = targetPlatform
-            this.project = project
             addModule(sourceModule!!)
         }
     }
-    // TODO remove further
-    CoreApplicationEnvironment.registerExtensionPoint(
-        analysisSession.project.extensionArea,
-        KtResolveExtensionProvider.EP_NAME.name,
-        KtResolveExtensionProvider::class.java
-    )
     return Pair(analysisSession, sourceModule ?: throw IllegalStateException())
-}
-
-// ----------- copy-paste from Analysis API ----------------------------------------------------------------------------
-/**
- * Collect source file path from the given [root] store them in [result].
- *
- * E.g., for `project/app/src` as a [root], this will walk the file tree and
- * collect all `.kt` and `.java` files under that folder.
- *
- * Note that this util gracefully skips [IOException] during file tree traversal.
- */
-internal fun collectSourceFilePaths(
-    root: Path,
-    result: MutableSet<String>
-) {
-    // NB: [Files#walk] throws an exception if there is an issue during IO.
-    // With [Files#walkFileTree] with a custom visitor, we can take control of exception handling.
-    Files.walkFileTree(
-        root,
-        object : SimpleFileVisitor<Path>() {
-            override fun preVisitDirectory(dir: Path, attrs: BasicFileAttributes): FileVisitResult {
-                return if (Files.isReadable(dir))
-                    FileVisitResult.CONTINUE
-                else
-                    FileVisitResult.SKIP_SUBTREE
-            }
-
-            override fun visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult {
-                if (!Files.isRegularFile(file) || !Files.isReadable(file))
-                    return FileVisitResult.CONTINUE
-                val ext = file.toFile().extension
-                if (ext == KotlinFileType.EXTENSION || ext == "java"/*JavaFileType.DEFAULT_EXTENSION*/) {
-                    result.add(file.toString())
-                }
-                return FileVisitResult.CONTINUE
-            }
-
-            override fun visitFileFailed(file: Path, exc: IOException?): FileVisitResult {
-                // TODO: report or log [IOException]?
-                // NB: this intentionally swallows the exception, hence fail-safe.
-                // Skipping subtree doesn't make any sense, since this is not a directory.
-                // Skipping sibling may drop valid file paths afterward, so we just continue.
-                return FileVisitResult.CONTINUE
-            }
-        }
-    )
-}
-
-/**
- * Collect source file path as [String] from the given source roots in [sourceRoot].
- *
- * this util collects all `.kt` and `.java` files under source roots.
- */
-internal fun getSourceFilePaths(
-    sourceRoot: Collection<String>,
-    includeDirectoryRoot: Boolean = false,
-): Set<String> {
-    val result = mutableSetOf<String>()
-    sourceRoot.forEach { srcRoot ->
-        val path = Paths.get(srcRoot)
-        if (Files.isDirectory(path)) {
-            // E.g., project/app/src
-            collectSourceFilePaths(path, result)
-            if (includeDirectoryRoot) {
-                result.add(srcRoot)
-            }
-        } else {
-            // E.g., project/app/src/some/pkg/main.kt
-            result.add(srcRoot)
-        }
-    }
-
-    return result
-}
-
-internal inline fun <reified T : PsiFileSystemItem> getPsiFilesFromPaths(
-    project: Project,
-    paths: Collection<String>,
-): List<T> {
-    val fs = StandardFileSystems.local()
-    val psiManager = PsiManager.getInstance(project)
-    val result = mutableListOf<T>()
-    for (path in paths) {
-        val vFile = fs.findFileByPath(path) ?: continue
-        val psiFileSystemItem =
-            if (vFile.isDirectory)
-                psiManager.findDirectory(vFile) as? T
-            else
-                psiManager.findFile(vFile) as? T
-        psiFileSystemItem?.let { result.add(it) }
-    }
-    return result
-}
+}
\ No newline at end of file
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt
index f217c88f..74d31269 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt
@@ -118,10 +118,7 @@ internal class DokkaSymbolVisitor(
     }
 
     fun visitModule(): DModule {
-        val ktFiles: List<KtFile> = getPsiFilesFromPaths(
-            analysisContext.project,
-            getSourceFilePaths(sourceSet.sourceRoots.map { it.canonicalPath })
-        )
+        val ktFiles = analysisContext.analysisSession.modulesWithFiles.entries.single().value.filterIsInstance<KtFile>()
         val processedPackages: MutableSet<FqName> = mutableSetOf()
         return analyze(analysisContext.mainModule) {
             val packageSymbols: List<DPackage> = ktFiles