* Do not walk directories for resolving content roots, use KotlinCoreEnvironment for getting sources.

* Move SymbolProcessorProvider loading into KSPCommandLineProcessor.
* Move compilerConfiguration processing into KSPCommandLineProcessor.
* Move KSPCoreEnvironment initialization out of KotlinSymbolProcessing constructor.

4 files changed, 67 insertions, 89 deletions

diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KSPCommandLineProcessor.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KSPCommandLineProcessor.kt
index 6b8676fd..fd9e61a0 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KSPCommandLineProcessor.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KSPCommandLineProcessor.kt
@@ -18,40 +18,36 @@
 package com.google.devtools.ksp.impl
 
 import com.google.devtools.ksp.KspOptions
+import com.google.devtools.ksp.processing.SymbolProcessorProvider
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoots
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.cli.jvm.config.addJavaSourceRoots
+import org.jetbrains.kotlin.cli.jvm.plugins.ServiceLoaderLite
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import java.io.File
-import java.nio.file.Files
+import java.net.URLClassLoader
 
-class KSPCommandLineProcessor(val args: Array<String>) {
-    val compilerConfiguration = CompilerConfiguration()
-    // TODO: support KSP options
-    val sources = args.toList()
-    lateinit var kspOptions: KspOptions
+class KSPCommandLineProcessor(val compilerConfiguration: CompilerConfiguration) {
+    private val kspOptionsBuilder = KspOptions.Builder()
 
-    val ktFiles = sources
-        .map { File(it) }
-        .sortedBy { Files.isSymbolicLink(it.toPath()) } // Get non-symbolic paths first
-        .flatMap { root -> root.walk().filter { it.isFile && it.extension == "kt" }.toList() }
-        .sortedBy { Files.isSymbolicLink(it.toPath()) }
-        .distinctBy { it.canonicalPath }
+    val kspOptions: KspOptions
+        get() = kspOptionsBuilder.build()
 
-    val javaFiles = sources
-        .map { File(it) }
-        .sortedBy { Files.isSymbolicLink(it.toPath()) } // Get non-symbolic paths first
-        .flatMap { root -> root.walk().filter { it.isFile && it.extension == "java" }.toList() }
-        .sortedBy { Files.isSymbolicLink(it.toPath()) }
-        .distinctBy { it.canonicalPath }
+    lateinit var providers: List<SymbolProcessorProvider>
 
-    init {
-        compilerConfiguration.addKotlinSourceRoots(ktFiles.map { it.absolutePath })
-        compilerConfiguration.addJavaSourceRoots(javaFiles)
+    fun processArgs(args: Array<String>) {
+        // TODO: support KSP options
+        val sources = args.toList()
+        compilerConfiguration.addKotlinSourceRoots(sources)
+        compilerConfiguration.addJavaSourceRoots(sources.map { File(it) })
         compilerConfiguration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, MessageCollector.NONE)
+        val processingClasspath = kspOptionsBuilder.processingClasspath
+        val classLoader = URLClassLoader(
+            processingClasspath.map { it.toURI().toURL() }.toTypedArray(),
+            javaClass.classLoader
+        )
 
-        val kspOptionsBuilder = KspOptions.Builder()
-        kspOptions = kspOptionsBuilder.build()
+        providers = ServiceLoaderLite.loadImplementations(SymbolProcessorProvider::class.java, classLoader)
     }
 }
diff --git a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
index c77ad1da..6db08650 100644
--- a/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
+++ b/kotlin-analysis-api/src/main/kotlin/com/google/devtools/ksp/impl/KotlinSymbolProcessing.kt
@@ -27,50 +27,31 @@ import com.google.devtools.ksp.symbol.KSAnnotated
 import com.intellij.mock.MockApplication
 import com.intellij.mock.MockProject
 import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
-import com.intellij.openapi.vfs.StandardFileSystems
 import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem
-import com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.analysis.api.standalone.configureApplicationEnvironment
 import org.jetbrains.kotlin.analysis.api.standalone.configureProjectEnvironment
-import org.jetbrains.kotlin.cli.common.config.kotlinSourceRoots
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
-import org.jetbrains.kotlin.cli.jvm.plugins.ServiceLoaderLite
 import org.jetbrains.kotlin.config.CompilerConfiguration
-import org.jetbrains.kotlin.psi.KtFile
-import java.net.URLClassLoader
 
 class KotlinSymbolProcessing(
     val compilerConfiguration: CompilerConfiguration,
     val options: KspOptions,
     val logger: KSPLogger,
-    val testProcessor: SymbolProcessorProvider? = null
-) {
-
+    val env: KotlinCoreEnvironment,
     val providers: List<SymbolProcessorProvider>
-    val env: KotlinCoreEnvironment = KotlinCoreEnvironment.createForProduction(
-        Disposer.newDisposable(), compilerConfiguration,
-        EnvironmentConfigFiles.JVM_CONFIG_FILES
-    )
+) {
     val project = env.project as MockProject
     val kspCoreEnvironment = KSPCoreEnvironment(project)
 
     var finished = false
     val deferredSymbols = mutableMapOf<SymbolProcessor, List<KSAnnotated>>()
-    val ktFiles = convertFilesToKtFiles(project, compilerConfiguration.kotlinSourceRoots.map { it.path })
-    val codeGenerator: CodeGeneratorImpl
-    val processors: List<SymbolProcessor>
-
-    init {
-        configureProjectEnvironment(
-            project,
-            compilerConfiguration,
-            env::createPackagePartProvider,
-            env.projectEnvironment.environment.jarFileSystem as CoreJarFileSystem
-        )
+    val ktFiles = env.getSourceFiles()
+    lateinit var codeGenerator: CodeGeneratorImpl
+    lateinit var processors: List<SymbolProcessor>
 
+    fun prepare() {
         val ksFiles = ktFiles.map { KSFileImpl(it) }
         val anyChangesWildcard = AnyChanges(options.projectBaseDir)
         codeGenerator = CodeGeneratorImpl(
@@ -83,18 +64,6 @@ class KotlinSymbolProcessing(
             ksFiles,
             options.incremental
         )
-        val application = ApplicationManager.getApplication() as MockApplication
-        configureApplicationEnvironment(application)
-        providers = if (testProcessor != null) {
-            listOf(testProcessor)
-        } else {
-            val processingClasspath = options.processingClasspath
-            val classLoader =
-                URLClassLoader(processingClasspath.map { it.toURI().toURL() }.toTypedArray(), javaClass.classLoader)
-
-            ServiceLoaderLite.loadImplementations(SymbolProcessorProvider::class.java, classLoader)
-        }
-
         processors = providers.mapNotNull { provider ->
             var processor: SymbolProcessor? = null
             processor = provider.create(
@@ -117,27 +86,37 @@ class KotlinSymbolProcessing(
         val resolver = ResolverAAImpl(ktFiles)
         processors.forEach { it.process(resolver) }
     }
-
-    private fun convertFilesToKtFiles(project: Project, filePaths: List<String>): List<KtFile> {
-        val fs = StandardFileSystems.local()
-        val psiManager = PsiManager.getInstance(project)
-        val ktFiles = mutableListOf<KtFile>()
-        for (path in filePaths) {
-            val vFile = fs.findFileByPath(path) ?: continue
-            val ktFile = psiManager.findFile(vFile) as? KtFile ?: continue
-            ktFiles.add(ktFile)
-        }
-        return ktFiles
-    }
 }
 
 fun main(args: Array<String>) {
-    val commandLineProcessor = KSPCommandLineProcessor(args)
+    val compilerConfiguration = CompilerConfiguration()
+    val commandLineProcessor = KSPCommandLineProcessor(compilerConfiguration)
     val logger = CommandLineKSPLogger()
+
+    val application = ApplicationManager.getApplication() as MockApplication
+    configureApplicationEnvironment(application)
+
+    commandLineProcessor.processArgs(args)
+
+    val kotlinCoreEnvironment = KotlinCoreEnvironment.createForProduction(
+        Disposer.newDisposable(), commandLineProcessor.compilerConfiguration,
+        EnvironmentConfigFiles.JVM_CONFIG_FILES
+    )
+
+    configureProjectEnvironment(
+        kotlinCoreEnvironment.project as MockProject,
+        compilerConfiguration,
+        kotlinCoreEnvironment::createPackagePartProvider,
+        kotlinCoreEnvironment.projectEnvironment.environment.jarFileSystem as CoreJarFileSystem
+    )
+
     val kotlinSymbolProcessing = KotlinSymbolProcessing(
         commandLineProcessor.compilerConfiguration,
         commandLineProcessor.kspOptions,
-        logger
+        logger,
+        kotlinCoreEnvironment,
+        commandLineProcessor.providers
     )
+    kotlinSymbolProcessing.prepare()
     kotlinSymbolProcessing.execute()
 }
diff --git a/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/AbstractKSPAATest.kt b/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/AbstractKSPAATest.kt
index 3e3a8875..56dbcfac 100644
--- a/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/AbstractKSPAATest.kt
+++ b/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/impl/test/AbstractKSPAATest.kt
@@ -77,13 +77,6 @@ abstract class AbstractKSPAATest : AbstractKSPTest(FrontendKinds.FIR) {
         val application = ApplicationManager.getApplication() as MockApplication
         configureApplicationEnvironment(application)
 
-        // TODO: other platforms
-        val kotlinCoreEnvironment = KotlinCoreEnvironment.createForTests(
-            disposable,
-            compilerConfiguration,
-            EnvironmentConfigFiles.JVM_CONFIG_FILES
-        )
-
         // Some underlying service needs files backed by local fs.
         // Therefore, this doesn't work:
         //  val ktFiles = mainModule.loadKtFiles(kotlinCoreEnvironment.project)
@@ -91,12 +84,16 @@ abstract class AbstractKSPAATest : AbstractKSPTest(FrontendKinds.FIR) {
         val kotlinSourceFiles = mainModule.files.filter { it.isKtFile }.map {
             File(mainModule.kotlinSrc, it.relativePath)
         }
-        val ktFiles = kotlinSourceFiles
+        val ktSourceRoots = kotlinSourceFiles
             .sortedBy { Files.isSymbolicLink(it.toPath()) } // Get non-symbolic paths first
-            .flatMap { root -> root.walk().filter { it.isFile && it.extension == "kt" }.toList() }
-            .sortedBy { Files.isSymbolicLink(it.toPath()) }
             .distinctBy { it.canonicalPath }
-        compilerConfiguration.addKotlinSourceRoots(ktFiles.map { it.absolutePath })
+        compilerConfiguration.addKotlinSourceRoots(ktSourceRoots.map { it.absolutePath })
+        // TODO: other platforms
+        val kotlinCoreEnvironment = KotlinCoreEnvironment.createForTests(
+            disposable,
+            compilerConfiguration,
+            EnvironmentConfigFiles.JVM_CONFIG_FILES
+        )
 
         configureProjectEnvironment(
             kotlinCoreEnvironment.project as MockProject,
@@ -119,9 +116,15 @@ abstract class AbstractKSPAATest : AbstractKSPTest(FrontendKinds.FIR) {
             cachesDir = File(testRoot, "kspTest/kspCaches")
             kspOutputDir = File(testRoot, "kspTest")
         }.build()
-        val ksp = KotlinSymbolProcessing(compilerConfiguration, kspOptions, CommandLineKSPLogger(), testProcessor)
+        val ksp = KotlinSymbolProcessing(
+            compilerConfiguration,
+            kspOptions,
+            CommandLineKSPLogger(),
+            kotlinCoreEnvironment,
+            listOf(testProcessor)
+        )
+        ksp.prepare()
         ksp.execute()
-
         return testProcessor.toResult()
     }
 }
diff --git a/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/processor/HelloWorldProcessor.kt b/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/processor/HelloWorldProcessor.kt
index c37ab7df..b3378adc 100644
--- a/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/processor/HelloWorldProcessor.kt
+++ b/kotlin-analysis-api/src/test/kotlin/com/google/devtools/ksp/processor/HelloWorldProcessor.kt
@@ -9,9 +9,9 @@ class HelloWorldProcessor : AbstractTestProcessor() {
     override fun process(resolver: Resolver): List<KSAnnotated> {
         val fooKt = resolver.getAllFiles().single()
         results.add(fooKt.fileName)
-        val bar = resolver.getAllFiles().single { it.fileName == "Foo.kt" }
+        val foo = resolver.getAllFiles().single { it.fileName == "Foo.kt" }
             .declarations.single { it.simpleName.asString() == "Foo" }
-        results.add(bar.simpleName.asString())
+        results.add(foo.simpleName.asString())
         return emptyList()
     }