[Analysis API] Add test infrastructure for code fragments

Now, Analysis API tests can define modules with code fragments.
The preceding module will be used as a context module.

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/registerBaseTestServices.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.Disposable
 import org.jetbrains.kotlin.analysis.test.framework.AnalysisApiTestDirectives
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.AnalysisApiKtModuleProvider
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.AnalysisApiKtModuleProviderImpl
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.AnalysisApiTestCodeFragmentDirectives
 import org.jetbrains.kotlin.analysis.test.framework.services.*
 import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfigurator
 import org.jetbrains.kotlin.test.ExecutionListenerBasedDisposableProvider
@@ -33,4 +34,5 @@ fun TestConfigurationBuilder.registerAnalysisApiBaseTestServices(
     useCustomCompilerConfigurationProvider(::AnalysisApiTestCompilerConfiguratorProvider)
     usePreAnalysisHandlers(::ProjectStructureInitialisationPreAnalysisHandler.bind(configurator))
     useDirectives(AnalysisApiTestDirectives)
+    useDirectives(AnalysisApiTestCodeFragmentDirectives)
 }

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/AnalysisApiKtModuleProviderImpl.kt

@@ -67,6 +67,7 @@ fun List<KtModuleWithFiles>.associateByName(): Map<String, KtModuleWithFiles> {
             is KtBuiltinsModule -> "Builtins for ${ktModule.platform}"
             is KtNotUnderContentRootModule -> ktModule.name
             is KtScriptModule -> ktModule.file.name
+            is KtDanglingFileModule -> ktModule.file.name
             else -> error("Unsupported module type: " + ktModule.javaClass.name)
         }
     }

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/KtCodeFragmentModuleFactory.kt

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.analysis.test.framework.project.structure
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.KtModuleWithFiles
+import org.jetbrains.kotlin.analysis.project.structure.DanglingFileResolutionMode
+import org.jetbrains.kotlin.analysis.project.structure.impl.KtDanglingFileModuleImpl
+import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
+import org.jetbrains.kotlin.psi.KtBlockCodeFragment
+import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtExpressionCodeFragment
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtTypeCodeFragment
+import org.jetbrains.kotlin.test.directives.model.DirectiveApplicability
+import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
+import org.jetbrains.kotlin.test.directives.model.singleOrZeroValue
+import org.jetbrains.kotlin.test.model.TestModule
+import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.sourceFileProvider
+
+object KtCodeFragmentModuleFactory : KtModuleFactory {
+    override fun createModule(
+        testModule: TestModule,
+        contextModule: KtModuleWithFiles?,
+        testServices: TestServices,
+        project: Project,
+    ): KtModuleWithFiles {
+        requireNotNull(contextModule) { "Context module should precede the code fragment module" }
+
+        val testFile = testModule.files.singleOrNull() ?: error("A single file is expected for a code fragment module")
+
+        val fileName = testFile.name
+        val fileText = testServices.sourceFileProvider.getContentOfSourceFile(testFile)
+
+        val codeFragmentKind = testFile.directives
+            .singleOrZeroValue(AnalysisApiTestCodeFragmentDirectives.CODE_FRAGMENT_KIND)
+            ?: CodeFragmentKind.BLOCK
+
+        val codeFragmentImports = testFile
+            .directives[AnalysisApiTestCodeFragmentDirectives.CODE_FRAGMENT_IMPORT]
+            .joinToString(KtCodeFragment.IMPORT_SEPARATOR)
+            .takeIf { it.isNotEmpty() }
+
+        if (codeFragmentKind == CodeFragmentKind.TYPE && codeFragmentImports != null) {
+            error("Imports cannot be configured for type code fragments")
+        }
+
+        val expressionMarkerProvider = testServices.expressionMarkerProvider
+        val contextElement = contextModule.files
+            .filterIsInstance<KtFile>()
+            .firstNotNullOfOrNull { expressionMarkerProvider.getElementOfTypeAtCaret<KtElement>(it, "context") }
+
+        val codeFragment = when (codeFragmentKind) {
+            CodeFragmentKind.EXPRESSION -> KtExpressionCodeFragment(project, fileName, fileText, codeFragmentImports, contextElement)
+            CodeFragmentKind.BLOCK -> KtBlockCodeFragment(project, fileName, fileText, codeFragmentImports, contextElement)
+            CodeFragmentKind.TYPE -> KtTypeCodeFragment(project, fileName, fileText, contextElement)
+        }
+
+        val module = KtDanglingFileModuleImpl(
+            codeFragment,
+            contextModule.ktModule,
+            DanglingFileResolutionMode.PREFER_SELF
+        )
+
+        return KtModuleWithFiles(module, listOf(codeFragment))
+    }
+}
+
+object AnalysisApiTestCodeFragmentDirectives : SimpleDirectivesContainer() {
+    val CODE_FRAGMENT_KIND by enumDirective<CodeFragmentKind>(
+        description = "Code fragment kind",
+        applicability = DirectiveApplicability.File
+    )
+
+    val CODE_FRAGMENT_IMPORT by stringDirective(
+        description = "Import local to the code fragment content",
+        applicability = DirectiveApplicability.File
+    )
+}
+
+enum class CodeFragmentKind {
+    EXPRESSION, BLOCK, TYPE
+}

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/KtLibraryBinaryModuleFactory.kt

@@ -18,7 +18,12 @@ import org.jetbrains.kotlin.test.services.TestServices
  * @see org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind.LibraryBinary
  */
 object KtLibraryBinaryModuleFactory : KtModuleFactory {
-    override fun createModule(testModule: TestModule, testServices: TestServices, project: Project): KtModuleWithFiles {
+    override fun createModule(
+        testModule: TestModule,
+        contextModule: KtModuleWithFiles?,
+        testServices: TestServices,
+        project: Project,
+    ): KtModuleWithFiles {
         val library = testServices.compiledLibraryProvider.compileToLibrary(testModule).artifact
         val decompiledFiles = testServices.testModuleDecompiler.getAllPsiFilesFromLibrary(library, project)
 

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/KtLibrarySourceModuleFactory.kt

@@ -22,7 +22,12 @@ import java.nio.file.Path
  * @see org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind.LibrarySource
  */
 object KtLibrarySourceModuleFactory : KtModuleFactory {
-    override fun createModule(testModule: TestModule, testServices: TestServices, project: Project): KtModuleWithFiles {
+    override fun createModule(
+        testModule: TestModule,
+        contextModule: KtModuleWithFiles?,
+        testServices: TestServices,
+        project: Project,
+    ): KtModuleWithFiles {
         Assume.assumeFalse("Compilation of multi-platform libraries is not supported", testModule.targetPlatform.isMultiPlatform())
 
         val (libraryJar, librarySourcesJar) = testServices.compiledLibraryProvider.compileToLibrary(testModule)

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/KtModuleFactory.kt

@@ -14,7 +14,19 @@ import org.jetbrains.kotlin.test.services.TestService
 import org.jetbrains.kotlin.test.services.TestServices
 
 fun interface KtModuleFactory : TestService {
-    fun createModule(testModule: TestModule, testServices: TestServices, project: Project): KtModuleWithFiles
+    /**
+     * Creates a [KtModule](org.jetbrains.kotlin.analysis.project.structure.KtModule) for the given [testModule].
+     *
+     * @param contextModule a module to use as a context module. Some kinds of modules (such as dangling file modules) require a
+     * context module. Modules representing code fragments also require a context element – that is why the [KtModuleWithFiles] is passed,
+     * instead of a plain (KtModule)[org.jetbrains.kotlin.analysis.project.structure.KtModule].
+     */
+    fun createModule(
+        testModule: TestModule,
+        contextModule: KtModuleWithFiles?,
+        testServices: TestServices,
+        project: Project,
+    ): KtModuleWithFiles
 }
 
 private val TestServices.ktModuleFactory: KtModuleFactory by TestServices.testServiceAccessor()
@@ -27,7 +39,7 @@ private val TestServices.ktModuleFactory: KtModuleFactory by TestServices.testSe
  * [MODULE_KIND][org.jetbrains.kotlin.analysis.test.framework.AnalysisApiTestDirectives.MODULE_KIND] directive for a specific test module.
  *
  * [DependencyKindModuleStructureTransformer][org.jetbrains.kotlin.analysis.test.framework.services.DependencyKindModuleStructureTransformer]
- * should be used to properly set up [DependencyKind][org.jetbrains.kotlin.test.model.DependencyKind] for module dependencies
+ * should be used to properly set up the [DependencyKind][org.jetbrains.kotlin.test.model.DependencyKind] for module dependencies.
  *
  * @see org.jetbrains.kotlin.analysis.test.framework.services.DependencyKindModuleStructureTransformer
  */
@@ -36,5 +48,6 @@ fun TestServices.getKtModuleFactoryForTestModule(testModule: TestModule): KtModu
     TestModuleKind.LibraryBinary -> KtLibraryBinaryModuleFactory
     TestModuleKind.LibrarySource -> KtLibrarySourceModuleFactory
     TestModuleKind.ScriptSource -> KtScriptModuleFactory
+    TestModuleKind.CodeFragment -> KtCodeFragmentModuleFactory
     else -> ktModuleFactory
 }

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/KtScriptModuleFactory.kt

@@ -15,7 +15,12 @@ import org.jetbrains.kotlin.test.services.TestServices
  * @see org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind.ScriptSource
  */
 object KtScriptModuleFactory : KtModuleFactory {
-    override fun createModule(testModule: TestModule, testServices: TestServices, project: Project): KtModuleWithFiles {
+    override fun createModule(
+        testModule: TestModule,
+        contextModule: KtModuleWithFiles?,
+        testServices: TestServices,
+        project: Project,
+    ): KtModuleWithFiles {
         val ktFile = TestModuleStructureFactory.createSourcePsiFiles(testModule, testServices, project).single() as KtFile
         val module = KtScriptModuleImpl(
             ktFile,

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/KtSourceModuleFactory.kt

@@ -15,7 +15,12 @@ import org.jetbrains.kotlin.test.services.TestServices
  * @see org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind.Source
  */
 object KtSourceModuleFactory : KtModuleFactory {
-    override fun createModule(testModule: TestModule, testServices: TestServices, project: Project): KtModuleWithFiles {
+    override fun createModule(
+        testModule: TestModule,
+        contextModule: KtModuleWithFiles?,
+        testServices: TestServices,
+        project: Project,
+    ): KtModuleWithFiles {
         val psiFiles = TestModuleStructureFactory.createSourcePsiFiles(testModule, testServices, project)
 
         val module = KtSourceModuleImpl(

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/project/structure/TestModuleStructureFactory.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.KtMod
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.KtModuleWithFiles
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.StandaloneProjectFactory
 import org.jetbrains.kotlin.analysis.project.structure.KtBinaryModule
+import org.jetbrains.kotlin.analysis.project.structure.KtDanglingFileModule
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtNotUnderContentRootModule
 import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
@@ -45,11 +46,11 @@ object TestModuleStructureFactory {
         testServices: TestServices,
         project: Project
     ): KtModuleProjectStructure {
-        val moduleEntries = moduleStructure.modules.map { testModule ->
-            testServices.getKtModuleFactoryForTestModule(testModule).createModule(testModule, testServices, project)
+        val modules = moduleStructure.modules.mapWithPrevious<TestModule, KtModuleWithFiles> { contextModule, testModule ->
+            testServices.getKtModuleFactoryForTestModule(testModule).createModule(testModule, contextModule, testServices, project)
         }
 
-        val modulesByName = moduleEntries.associateByName()
+        val modulesByName = modules.associateByName()
 
         val libraryCache: LibraryCache = mutableMapOf()
 
@@ -60,7 +61,18 @@ object TestModuleStructureFactory {
             ktModule.addDependencies(testModule, testServices, modulesByName, libraryCache)
         }
 
-        return KtModuleProjectStructure(moduleEntries, libraryCache.values)
+        return KtModuleProjectStructure(modules, libraryCache.values)
+    }
+
+    private fun <T : Any, R> List<T>.mapWithPrevious(transform: (R?, T) -> R): List<R> {
+        val result = ArrayList<R>(this.size)
+        var previousResult: R? = null
+        for (item in this) {
+            val newResult = transform(previousResult, item)
+            result.add(newResult)
+            previousResult = newResult
+        }
+        return result
     }
 
     private fun ModulesByName.getByTestModule(testModule: TestModule): KtModuleWithFiles =
@@ -87,6 +99,9 @@ object TestModuleStructureFactory {
         is KtNotUnderContentRootModule -> {
             // Not-under-content-root modules have no external dependencies on purpose
         }
+        is KtDanglingFileModule -> {
+            // Dangling file modules get dependencies from their context
+        }
         is KtModuleWithModifiableDependencies -> {
             addModuleDependencies(testModule, modulesByName, this)
             addLibraryDependencies(testModule, testServices, project, this, libraryCache::getOrPut)

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/DependencyKindModuleStructureTransformer.kt

@@ -47,10 +47,21 @@ object DependencyKindModuleStructureTransformer : ModuleStructureTransformer() {
     ): DependencyDescription {
         val dependencyModule = moduleMapping.getValue(dependency.moduleName)
         val newKind = when (dependencyModule.moduleKind) {
-            TestModuleKind.Source, TestModuleKind.LibrarySource, TestModuleKind.ScriptSource -> DependencyKind.Source
-            TestModuleKind.LibraryBinary -> DependencyKind.Binary
-            // There is no explicit module kind, so the dependency already has the right kind
-            null -> return dependency
+            TestModuleKind.Source,
+            TestModuleKind.LibrarySource,
+            TestModuleKind.ScriptSource,
+            TestModuleKind.CodeFragment -> {
+                DependencyKind.Source
+            }
+
+            TestModuleKind.LibraryBinary -> {
+                DependencyKind.Binary
+            }
+
+            null -> {
+                // There is no explicit module kind, so the dependency already has the right kind
+                return dependency
+            }
         }
 
         return dependency.copy(kind = newKind)

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/test/configurators/TestModuleKind.kt

@@ -38,6 +38,11 @@ enum class TestModuleKind(val suffix: String) {
      * @see org.jetbrains.kotlin.analysis.test.framework.project.structure.KtScriptModuleFactory
      */
     ScriptSource("ScriptSource"),
+
+    /**
+     * @see org.jetbrains.kotlin.analysis.test.framework.project.structure.KtCodeFragmentModuleFactory
+     */
+    CodeFragment("CodeFragment")
 }
 
 val TestModule.moduleKind: TestModuleKind?