[LL API] Support code fragment analysis

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/components/diagnosticProvider/FirIdeNormalAnalysisSourceModuleCodeFragmentCollectDiagnosticsTestGenerated.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2023 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.api.fir.test.cases.generated.cases.components.diagnosticProvider;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.util.KtTestUtil;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.analysis.api.fir.test.configurators.AnalysisApiFirTestConfiguratorFactory;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfiguratorFactoryData;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfigurator;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.FrontendKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisSessionMode;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiMode;
+import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.diagnosticProvider.AbstractCodeFragmentCollectDiagnosticsTest;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.analysis.api.GenerateAnalysisApiTestsKt}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics")
+@TestDataPath("$PROJECT_ROOT")
+public class FirIdeNormalAnalysisSourceModuleCodeFragmentCollectDiagnosticsTestGenerated extends AbstractCodeFragmentCollectDiagnosticsTest {
+    @NotNull
+    @Override
+    public AnalysisApiTestConfigurator getConfigurator() {
+        return AnalysisApiFirTestConfiguratorFactory.INSTANCE.createConfigurator(
+            new AnalysisApiTestConfiguratorFactoryData(
+                FrontendKind.Fir,
+                TestModuleKind.Source,
+                AnalysisSessionMode.Normal,
+                AnalysisApiMode.Ide
+            )
+        );
+    }
+
+    @Test
+    public void testAllFilesPresentInCodeFragmentDiagnostics() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+    }
+
+    @Test
+    @TestMetadata("localClass.kt")
+    public void testLocalClass() throws Exception {
+        runTest("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localClass.kt");
+    }
+
+    @Test
+    @TestMetadata("localFunction.kt")
+    public void testLocalFunction() throws Exception {
+        runTest("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localFunction.kt");
+    }
+
+    @Test
+    @TestMetadata("typeMismatch.kt")
+    public void testTypeMismatch() throws Exception {
+        runTest("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/typeMismatch.kt");
+    }
+}

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/diagnosticProvider/AbstractCodeFragmentCollectDiagnosticsTest.kt

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2023 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.api.impl.base.test.cases.components.diagnosticProvider
+
+import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtPsiFactory
+import org.jetbrains.kotlin.test.model.TestModule
+import org.jetbrains.kotlin.test.services.TestServices
+import java.io.File
+
+abstract class AbstractCodeFragmentCollectDiagnosticsTest : AbstractCollectDiagnosticsTest() {
+    override fun doTestByFileStructure(ktFile: KtFile, module: TestModule, testServices: TestServices) {
+        val contextElement = testServices.expressionMarkerProvider.getElementOfTypeAtCaret<KtElement>(ktFile)
+
+        val fragmentText = module.files.single().originalFile
+            .run { File(parent, "$nameWithoutExtension.fragment.$extension") }
+            .readText()
+
+        val isBlockFragment = fragmentText.any { it == '\n' }
+
+        val project = ktFile.project
+        val factory = KtPsiFactory(project, markGenerated = false)
+
+        val codeFragment = when {
+            isBlockFragment -> factory.createBlockCodeFragment(fragmentText, contextElement)
+            else -> factory.createExpressionCodeFragment(fragmentText, contextElement)
+        }
+
+        super.doTestByFileStructure(codeFragment, module, testServices)
+    }
+}

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/components/diagnosticProvider/FirStandaloneNormalAnalysisSourceModuleCodeFragmentCollectDiagnosticsTestGenerated.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2023 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.api.standalone.fir.test.cases.generated.cases.components.diagnosticProvider;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.util.KtTestUtil;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.analysis.api.standalone.fir.test.AnalysisApiFirStandaloneModeTestConfiguratorFactory;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfiguratorFactoryData;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestConfigurator;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.TestModuleKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.FrontendKind;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisSessionMode;
+import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiMode;
+import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.diagnosticProvider.AbstractCodeFragmentCollectDiagnosticsTest;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.analysis.api.GenerateAnalysisApiTestsKt}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics")
+@TestDataPath("$PROJECT_ROOT")
+public class FirStandaloneNormalAnalysisSourceModuleCodeFragmentCollectDiagnosticsTestGenerated extends AbstractCodeFragmentCollectDiagnosticsTest {
+    @NotNull
+    @Override
+    public AnalysisApiTestConfigurator getConfigurator() {
+        return AnalysisApiFirStandaloneModeTestConfiguratorFactory.INSTANCE.createConfigurator(
+            new AnalysisApiTestConfiguratorFactoryData(
+                FrontendKind.Fir,
+                TestModuleKind.Source,
+                AnalysisSessionMode.Normal,
+                AnalysisApiMode.Standalone
+            )
+        );
+    }
+
+    @Test
+    public void testAllFilesPresentInCodeFragmentDiagnostics() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+    }
+
+    @Test
+    @TestMetadata("localClass.kt")
+    public void testLocalClass() throws Exception {
+        runTest("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localClass.kt");
+    }
+
+    @Test
+    @TestMetadata("localFunction.kt")
+    public void testLocalFunction() throws Exception {
+        runTest("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localFunction.kt");
+    }
+
+    @Test
+    @TestMetadata("typeMismatch.kt")
+    public void testTypeMismatch() throws Exception {
+        runTest("analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/typeMismatch.kt");
+    }
+}

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localClass.fragment.kt

@@ -0,0 +1,6 @@
+class Foo {
+    val a = x
+    val b = y
+}
+
+Foo().b

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localClass.kt

@@ -0,0 +1,4 @@
+fun test() {
+    val x = 0
+    <caret>val y = 0
+}

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localClass.txt

@@ -0,0 +1,4 @@
+Diagnostics from elements:
+  for PSI element of type KtNameReferenceExpression at (3,13-14)
+    UNRESOLVED_REFERENCE      text ranges: [(38,39)]
+      PSI: KtNameReferenceExpression at (3,13-14)

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localFunction.fragment.kt

@@ -0,0 +1,5 @@
+fun foo(): Pair<Int, Int> {
+    return x to y
+}
+
+foo()

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localFunction.kt

@@ -0,0 +1,6 @@
+// WITH_STDLIB
+
+fun test() {
+    val x = 0
+    <caret>val y = 0
+}

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/localFunction.txt

@@ -0,0 +1,4 @@
+Diagnostics from elements:
+  for PSI element of type KtNameReferenceExpression at (2,17-18)
+    UNRESOLVED_REFERENCE      text ranges: [(44,45)]
+      PSI: KtNameReferenceExpression at (2,17-18)

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/typeMismatch.fragment.kt

@@ -0,0 +1 @@
+foo(x)

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/typeMismatch.kt

@@ -0,0 +1,6 @@
+fun test() {
+    val x = 0
+    <caret>val y = 1
+}
+
+fun foo(a: String) {}

analysis/analysis-api/testData/components/diagnosticsProvider/codeFragmentDiagnostics/typeMismatch.txt

@@ -0,0 +1,4 @@
+Diagnostics from elements:
+  for PSI element of type KtNameReferenceExpression at (1,5-6)
+    ARGUMENT_TYPE_MISMATCH      text ranges: [(4,5)]
+      PSI: KtNameReferenceExpression at (1,5-6)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/FirDesignation.kt

@@ -108,7 +108,7 @@ private fun collectDesignationPath(target: FirElementWithResolveState): List<Fir
             return if (target.diagnostic == ConeDestructuringDeclarationsOnTopLevel) emptyList() else null
         }
 
-        is FirScript -> {
+        is FirScript, is FirCodeFragment -> {
             return emptyList()
         }
 
@@ -270,7 +270,7 @@ fun FirElementWithResolveState.tryCollectDesignation(): FirDesignation? =
 
 fun FirElementWithResolveState.tryCollectDesignationWithFile(): FirDesignationWithFile? {
     return when (this) {
-        is FirScript, is FirFileAnnotationsContainer -> {
+        is FirScript, is FirCodeFragment, is FirFileAnnotationsContainer -> {
             val firFile = getContainingFile() ?: return null
             FirDesignationWithFile(path = emptyList(), this, firFile)
         }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostics/LLFirDiagnosticVisitor.kt

@@ -5,13 +5,21 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+import org.jetbrains.kotlin.analysis.utils.printer.parentsOfType
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContextForProvider
 import org.jetbrains.kotlin.fir.analysis.collectors.CheckerRunningDiagnosticCollectorVisitor
 import org.jetbrains.kotlin.fir.analysis.collectors.DiagnosticCollectorComponents
+import org.jetbrains.kotlin.fir.declarations.FirCodeFragment
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.psi
+import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtDeclaration
 
 internal open class LLFirDiagnosticVisitor(
     context: CheckerContextForProvider,
@@ -44,6 +52,34 @@ internal open class LLFirDiagnosticVisitor(
         }
     }
 
+    override fun visitCodeFragment(codeFragment: FirCodeFragment, data: Nothing?) {
+        val contextElement = (codeFragment.source.psi as? KtCodeFragment)?.context
+        if (contextElement != null) {
+            fun process(containingSymbols: List<FirDeclaration>) {
+                if (containingSymbols.isEmpty()) {
+                    super.visitCodeFragment(codeFragment, data)
+                } else {
+                    withDeclaration(containingSymbols.first()) {
+                        process(containingSymbols.subList(1, containingSymbols.size))
+                    }
+                }
+            }
+
+            val project = contextElement.project
+            val module = ProjectStructureProvider.getModule(project, contextElement, contextualModule = null)
+            val resolveSession = module.getFirResolveSession(project)
+
+            // Register containing declarations of a context element
+            contextElement.parentsOfType<KtDeclaration>().toList().asReversed()
+                .map { it.resolveToFirSymbol(resolveSession).fir }
+                .run(::process)
+
+            return
+        }
+
+        super.visitCodeFragment(codeFragment, data)
+    }
+
     /**
      * Some FirClassChecker may report diagnostics on class member headers.
      * That diagnostics should be suppressed if we have a `@Suppress` annotation on class member.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostics/fir/ContextCollectingDiagnosticCollectorVisitor.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollectorV
 import org.jetbrains.kotlin.fir.containingClass
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.classId
-import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.SessionHolder
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
@@ -67,7 +66,7 @@ internal object PersistenceContextCollector {
             is FirCallableDeclaration -> declaration.symbol.callableId.isLocal
             is FirDanglingModifierList -> declaration.containingClass()?.classId?.isLocal == true
             is FirAnonymousInitializer -> declaration.containingClass().classId.isLocal
-            is FirScript -> false
+            is FirScript, is FirCodeFragment -> false
             else -> errorWithAttachment("Unsupported declaration ${declaration::class}") {
                 withFirEntry("declaration", declaration)
             }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt

@@ -262,7 +262,8 @@ internal fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDec
                 !notNullCandidate.isPartOf(parent) ||
                 parent is KtClassInitializer ||
                 parent is KtObjectLiteralExpression ||
-                parent is KtCallElement
+                parent is KtCallElement ||
+                parent is KtCodeFragment
             ) {
                 // Candidate turned out to be local. Let's find another one.
                 candidate = null

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirFileBuilder.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.builder.PsiRawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
 
 /**

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructure.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 import java.util.concurrent.ConcurrentHashMap
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.canBePartOfParentDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.fir.correspondingProperty
 import org.jetbrains.kotlin.fir.declarations.impl.FirPrimaryConstructor
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
@@ -189,6 +190,13 @@ internal class FileStructure private constructor(
     }
 
     private fun createStructureElement(container: KtElement): FileStructureElement = when {
+        container is KtCodeFragment -> {
+            val firCodeFragment = firFile.codeFragment
+            firCodeFragment.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+
+            val timestamp = container.modificationStamp
+            ReanalyzableCodeFragmentStructureElement(firFile, container, firCodeFragment.symbol, timestamp, moduleComponents)
+        }
         container is KtFile -> {
             val firFile = moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
             firFile.lazyResolveToPhase(FirResolvePhase.IMPORTS)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt

@@ -9,6 +9,9 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiErrorElement
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.ClassDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.FileDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.FileStructureElementDiagnostics
@@ -20,7 +23,9 @@ import org.jetbrains.kotlin.fir.correspondingProperty
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirErrorConstructor
 import org.jetbrains.kotlin.fir.declarations.impl.FirPrimaryConstructor
+import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirCodeFragmentSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
@@ -106,12 +111,12 @@ internal class KtToFirMapping(firElement: FirElement, recorder: FirElementsRecor
     }
 }
 
-internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S : FirBasedSymbol<*>>(
+internal sealed class ReanalyzableStructureElement<KT : KtAnnotated, S : FirBasedSymbol<*>>(
     firFile: FirFile,
     val firSymbol: S,
     moduleComponents: LLFirModuleResolveComponents,
 ) : FileStructureElement(firFile, moduleComponents) {
-    abstract override val psi: KtDeclaration
+    abstract override val psi: KtAnnotated
     abstract val timestamp: Long
 
     /**
@@ -132,6 +137,45 @@ internal sealed class ReanalyzableStructureElement<KT : KtDeclaration, S : FirBa
     }
 }
 
+internal class ReanalyzableCodeFragmentStructureElement(
+    firFile: FirFile,
+    override val psi: KtCodeFragment,
+    firSymbol: FirCodeFragmentSymbol,
+    override val timestamp: Long,
+    moduleComponents: LLFirModuleResolveComponents,
+) : ReanalyzableStructureElement<KtCodeFragment, FirCodeFragmentSymbol>(firFile, firSymbol, moduleComponents) {
+    override val mappings = KtToFirMapping(firSymbol.fir, recorder)
+
+    override fun reanalyze(): ReanalyzableStructureElement<KtCodeFragment, FirCodeFragmentSymbol> {
+        val originalCodeFragment = firSymbol.fir
+        val originalDesignation = originalCodeFragment.collectDesignation()
+
+        firFile.transformDeclarations(object : FirTransformer<Nothing?>() {
+            override fun <E : FirElement> transformElement(element: E, data: Nothing?): E = element
+
+            override fun transformCodeFragment(codeFragment: FirCodeFragment, data: Nothing?): FirCodeFragment {
+                return RawFirNonLocalDeclarationBuilder.build(
+                    session = originalCodeFragment.moduleData.session,
+                    scopeProvider = originalCodeFragment.moduleData.session.kotlinScopeProvider,
+                    designation = originalDesignation,
+                    rootNonLocalDeclaration = psi,
+                ) as FirCodeFragment
+            }
+        }, null)
+
+        val newCodeFragment = firFile.codeFragment
+        newCodeFragment.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+
+        return ReanalyzableCodeFragmentStructureElement(
+            firFile,
+            psi,
+            newCodeFragment.symbol,
+            psi.modificationStamp,
+            moduleComponents
+        )
+    }
+}
+
 internal class ReanalyzableFunctionStructureElement(
     firFile: FirFile,
     override val psi: KtNamedFunction,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/inBlockModification.kt

@@ -8,9 +8,10 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
 import com.intellij.openapi.application.ApplicationManager
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.fir.contracts.impl.FirEmptyContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
@@ -21,23 +22,44 @@ import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildLazyBlock
 import org.jetbrains.kotlin.fir.expressions.builder.buildLazyExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
+import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtElement
+
+@Deprecated(
+    "Temporarily left for binary compatibility. Use invalidateAfterInBlockModification(KtElement) instead.",
+    replaceWith = ReplaceWith("invalidateAfterInBlockModification(declaration)", "org.jetbrains.kotlin.psi.KtElement"),
+    level = DeprecationLevel.HIDDEN,
+)
+@LLFirInternals
+@Suppress("unused")
+fun invalidateAfterInBlockModification(declaration: KtDeclaration): Boolean {
+    return invalidateAfterInBlockModification(declaration as KtElement)
+}
 
 /**
  * Must be called in a write action.
  * @return **false** if it is not in-block modification
  */
 @LLFirInternals
-fun invalidateAfterInBlockModification(declaration: KtDeclaration): Boolean {
+fun invalidateAfterInBlockModification(declaration: KtElement): Boolean {
     ApplicationManager.getApplication().assertIsWriteThread()
 
     val project = declaration.project
     val ktModule = ProjectStructureProvider.getModule(project, declaration, contextualModule = null)
     val resolveSession = ktModule.getFirResolveSession(project)
-    return when (val firDeclaration = declaration.resolveToFirSymbol(resolveSession).fir) {
+
+    val firDeclaration = when (declaration) {
+        is KtCodeFragment -> declaration.getOrBuildFirFile(resolveSession).codeFragment
+        is KtDeclaration -> declaration.resolveToFirSymbol(resolveSession).fir
+        else -> errorWithFirSpecificEntries("Unexpected declaration kind: ${declaration.javaClass.simpleName}", psi = declaration)
+    }
+
+    return when (firDeclaration) {
         is FirSimpleFunction -> firDeclaration.inBodyInvalidation()
         is FirPropertyAccessor -> firDeclaration.inBodyInvalidation()
         is FirProperty -> firDeclaration.inBodyInvalidation()
+        is FirCodeFragment -> firDeclaration.inBodyInvalidation()
         else -> errorWithFirSpecificEntries("Unknown declaration with body", fir = firDeclaration, psi = declaration)
     }
 }
@@ -145,6 +167,17 @@ private fun FirPropertyAccessor.inBodyInvalidation(): Boolean {
     return true
 }
 
+private fun FirCodeFragment.inBodyInvalidation(): Boolean {
+    if (block is FirLazyBlock) {
+        return true
+    }
+
+    decreasePhase(FirResolvePhase.BODY_RESOLVE.previous)
+    replaceBlock(buildLazyBlock())
+
+    return true
+}
+
 private fun FirProperty.invalidateInitializer(): PropertyExpressionState = replaceWithLazyExpressionIfNeeded(::initializer, ::replaceInitializer)
 
 private fun FirProperty.invalidateDelegate(): PropertyExpressionState = replaceWithLazyExpressionIfNeeded(::delegate, ::replaceDelegate)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -87,7 +87,7 @@ private inline fun <reified T : FirDeclaration> revive(
         session = session,
         scopeProvider = session.kotlinScopeProvider,
         designation = designation,
-        rootNonLocalDeclaration = psiFactory(designation) as KtDeclaration,
+        rootNonLocalDeclaration = psiFactory(designation) as KtElement,
     ) as T
 }
 
@@ -265,6 +265,14 @@ private fun needCalculatingLazyBodyForProperty(firProperty: FirProperty): Boolea
             || firProperty.delegate is FirLazyExpression
             || firProperty.getExplicitBackingField()?.initializer is FirLazyExpression
 
+private fun calculateLazyBodyForCodeFragment(designation: FirDesignation) {
+    val codeFragment = designation.target as FirCodeFragment
+    require(codeFragment.block is FirLazyBlock)
+
+    val newCodeFragment = revive<FirCodeFragment>(designation)
+    codeFragment.replaceBlock(newCodeFragment.block)
+}
+
 private enum class FirLazyAnnotationTransformerScope {
     ALL_ANNOTATIONS,
     COMPILER_ONLY;
@@ -472,4 +480,13 @@ private abstract class FirLazyBodiesCalculatorTransformer : FirTransformer<Persi
 
         return anonymousInitializer
     }
+
+    override fun transformCodeFragment(codeFragment: FirCodeFragment, data: PersistentList<FirRegularClass>): FirCodeFragment {
+        if (codeFragment.block is FirLazyBlock) {
+            val designation = FirDesignation(data, codeFragment)
+            calculateLazyBodyForCodeFragment(designation)
+        }
+
+        return codeFragment
+    }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt

@@ -36,7 +36,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
     session: FirSession,
     baseScopeProvider: FirScopeProvider,
     private val originalDeclaration: FirDeclaration,
-    private val declarationToBuild: KtDeclaration,
+    private val declarationToBuild: KtElement,
     private val functionsToRebind: Set<FirFunction>? = null,
     private val replacementApplier: RawFirReplacement.Applier? = null,
 ) : PsiRawFirBuilder(session, baseScopeProvider, bodyBuildingMode = BodyBuildingMode.NORMAL) {
@@ -58,24 +58,16 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             replacement: RawFirReplacement?,
         ): FirDeclaration {
             val replacementApplier = replacement?.Applier()
-            val builder = RawFirNonLocalDeclarationBuilder(
+            val result = build(session, scopeProvider, designation, rootNonLocalDeclaration, replacementApplier = replacementApplier)
-                session = session,
+            replacementApplier?.ensureApplied()
-                baseScopeProvider = scopeProvider,
+            return result
-                originalDeclaration = designation.target as FirDeclaration,
-                declarationToBuild = rootNonLocalDeclaration,
-                replacementApplier = replacementApplier
-            )
-            builder.context.packageFqName = rootNonLocalDeclaration.containingKtFile.packageFqName
-            return builder.moveNext(designation.path.iterator(), containingClass = null).also {
-                replacementApplier?.ensureApplied()
-            }
         }
 
         fun buildWithFunctionSymbolRebind(
             session: FirSession,
             scopeProvider: FirScopeProvider,
             designation: FirDesignation,
-            rootNonLocalDeclaration: KtDeclaration,
+            rootNonLocalDeclaration: KtElement,
         ): FirDeclaration {
             val functionsToRebind = when (val originalDeclaration = designation.target) {
                 is FirFunction -> setOf(originalDeclaration)
@@ -83,12 +75,26 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
                 else -> null
             }
 
+            return build(session, scopeProvider, designation, rootNonLocalDeclaration, functionsToRebind)
+        }
+
+        fun build(
+            session: FirSession,
+            scopeProvider: FirScopeProvider,
+            designation: FirDesignation,
+            rootNonLocalDeclaration: KtElement,
+            functionsToRebind: Set<FirFunction>? = null,
+            replacementApplier: RawFirReplacement.Applier? = null
+        ): FirDeclaration {
+            check(rootNonLocalDeclaration is KtDeclaration || rootNonLocalDeclaration is KtCodeFragment)
+
             val builder = RawFirNonLocalDeclarationBuilder(
                 session = session,
                 baseScopeProvider = scopeProvider,
                 originalDeclaration = designation.target as FirDeclaration,
                 declarationToBuild = rootNonLocalDeclaration,
                 functionsToRebind = functionsToRebind,
+                replacementApplier = replacementApplier
             )
             builder.context.packageFqName = rootNonLocalDeclaration.containingKtFile.packageFqName
             return builder.moveNext(designation.path.iterator(), containingClass = null)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt

@@ -162,7 +162,7 @@ internal fun LLFirAbstractBodyTargetResolver.transformAnnotations(target: FirEle
             target.transformSingle(transformer, ResolutionMode.ContextIndependent)
         }
 
-        target is FirFile -> {}
+        target is FirCodeFragment || target is FirFile -> {}
 
         else -> throwUnexpectedFirElementError(target)
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -5,11 +5,18 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.transformers
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignationWithFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.asResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirElementError
+import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkDelegatedConstructorIsResolved
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.*
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
@@ -25,19 +32,28 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildLazyDelegatedConstructo
 import org.jetbrains.kotlin.fir.expressions.builder.buildMultiDelegatedConstructorCall
 import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
 import org.jetbrains.kotlin.fir.expressions.impl.FirLazyDelegatedConstructorCall
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.references.FirThisReference
 import org.jetbrains.kotlin.fir.references.builder.buildExplicitSuperReference
 import org.jetbrains.kotlin.fir.references.builder.buildExplicitThisReference
+import org.jetbrains.kotlin.fir.resolve.FirCodeFragmentContext
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.dfa.FirControlFlowGraphReferenceImpl
+import org.jetbrains.kotlin.fir.resolve.codeFragmentContext
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.contracts.FirContractsDslNames
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.isUsedInControlFlowGraphBuilderForClass
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import org.jetbrains.kotlin.utils.exceptions.buildErrorWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment
 import org.jetbrains.kotlin.utils.findIsInstanceAnd
 
@@ -100,6 +116,14 @@ private class LLFirBodyTargetResolver(
                     calculateControlFlowGraph(target)
                 }
 
+                return true
+            }
+            is FirCodeFragment -> {
+                resolveCodeFragmentContext(target)
+                performCustomResolveUnderLock(target) {
+                    resolve(target, BodyStateKeepers.CODE_FRAGMENT)
+                }
+
                 return true
             }
         }
@@ -136,9 +160,57 @@ private class LLFirBodyTargetResolver(
         }
     }
 
+    private fun resolveCodeFragmentContext(firCodeFragment: FirCodeFragment) {
+        val ktCodeFragment = firCodeFragment.psi as? KtCodeFragment
+            ?: errorWithAttachment("Code fragment source not found") {
+                withFirEntry("firCodeFragment", firCodeFragment)
+            }
+
+        val module = firCodeFragment.llFirModuleData.ktModule
+        val resolveSession = module.getFirResolveSession(ktCodeFragment.project) as LLFirResolvableResolveSession
+
+        val contextPsiElement = ktCodeFragment.context
+        val contextKtFile = contextPsiElement?.containingFile as? KtFile
+
+        firCodeFragment.codeFragmentContext = if (contextKtFile != null) {
+            val contextFirFile = resolveSession.getOrBuildFirFile(contextKtFile)
+            val target = computeCodeFragmentContextResolveTarget(contextFirFile, contextPsiElement)
+
+            val contextParentKtElements = contextPsiElement.parentsWithSelf.toSet()
+
+            val contextProvider = ContextCollector.process(target, transformer.components) { candidate ->
+                when (candidate) {
+                    contextPsiElement -> ContextCollector.FilterResponse.STOP
+                    in contextParentKtElements -> ContextCollector.FilterResponse.CONTINUE
+                    else -> ContextCollector.FilterResponse.SKIP
+                }
+            }
+
+            val elementContext = contextProvider[contextPsiElement, ContextCollector.ContextKind.BODY]
+                ?: contextParentKtElements.firstNotNullOf { contextProvider[it, ContextCollector.ContextKind.SELF] }
+
+            LLFirCodeFragmentContext(elementContext.towerDataContext, elementContext.smartCasts)
+        } else {
+            val towerDataContext = FirTowerDataContext()
+            LLFirCodeFragmentContext(towerDataContext, emptyMap())
+        }
+    }
+
+    private fun computeCodeFragmentContextResolveTarget(contextFirFile: FirFile, contextPsiElement: PsiElement): LLFirResolveTarget {
+        val contextKtDeclaration = contextPsiElement.getNonLocalContainingOrThisDeclaration()
+        if (contextKtDeclaration != null) {
+            val designationPath = FirElementFinder.collectDesignationPath(contextFirFile, contextKtDeclaration)
+            if (designationPath != null) {
+                return FirDesignationWithFile(designationPath.path, designationPath.target, contextFirFile).asResolveTarget()
+            }
+        }
+
+        return LLFirWholeFileResolveTarget(contextFirFile)
+    }
+
     override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
         when (target) {
-            is FirRegularClass -> error("Should have been resolved in ${::doResolveWithoutLock.name}")
+            is FirRegularClass, is FirCodeFragment -> error("Should have been resolved in ${::doResolveWithoutLock.name}")
             is FirConstructor -> resolve(target, BodyStateKeepers.CONSTRUCTOR)
             is FirFunction -> resolve(target, BodyStateKeepers.FUNCTION)
             is FirProperty -> resolve(target, BodyStateKeepers.PROPERTY)
@@ -209,6 +281,10 @@ internal object BodyStateKeepers {
         )
     }
 
+    val CODE_FRAGMENT: StateKeeper<FirCodeFragment, FirDesignationWithFile> = stateKeeper { _, _ ->
+        add(FirCodeFragment::block, FirCodeFragment::replaceBlock, ::blockGuard)
+    }
+
     val ANONYMOUS_INITIALIZER: StateKeeper<FirAnonymousInitializer, FirDesignationWithFile> = stateKeeper { _, _ ->
         add(FirAnonymousInitializer::body, FirAnonymousInitializer::replaceBody, ::blockGuard)
         add(FirAnonymousInitializer::controlFlowGraphReference, FirAnonymousInitializer::replaceControlFlowGraphReference)
@@ -409,4 +485,9 @@ private fun requireSameSize(old: List<FirStatement>, new: List<FirStatement>) {
             new.forEachIndexed { index, statement -> withFirEntry("statement$index", statement) }
         }
     }
-}
+}
+
+private class LLFirCodeFragmentContext(
+    override val towerDataContext: FirTowerDataContext,
+    override val variables: Map<FirBasedSymbol<*>, Set<ConeKotlinType>>
+) : FirCodeFragmentContext

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirCompilerAnnotationsLazyResolver.kt

@@ -103,9 +103,14 @@ private class LLFirCompilerRequiredAnnotationsTargetResolver(
 
     override fun doResolveWithoutLock(target: FirElementWithResolveState): Boolean {
         if (target is FirFile) return false
-        when {
+
-            target is FirRegularClass || target is FirScript || target.isRegularDeclarationWithAnnotation -> {}
+        when (target) {
-            else -> throwUnexpectedFirElementError(target)
+            is FirRegularClass, is FirScript, is FirCodeFragment -> {}
+            else -> {
+                if (!target.isRegularDeclarationWithAnnotation) {
+                    throwUnexpectedFirElementError(target)
+                }
+            }
         }
 
         requireIsInstance<FirAnnotationContainer>(target)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirContractsLazyResolver.kt

@@ -65,6 +65,7 @@ private class LLFirContractsTargetResolver(
             is FirAnonymousInitializer,
             is FirFile,
             is FirScript,
+            is FirCodeFragment,
             is FirFileAnnotationsContainer,
             is FirDanglingModifierList,
             -> {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirImplicitTypesLazyResolver.kt

@@ -78,6 +78,7 @@ internal class LLFirImplicitBodyTargetResolver(
             is FirRegularClass,
             is FirTypeAlias,
             is FirFile,
+            is FirCodeFragment,
             is FirAnonymousInitializer,
             is FirDanglingModifierList,
             is FirFileAnnotationsContainer,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirTypeLazyResolver.kt

@@ -131,7 +131,7 @@ private class LLFirTypeTargetResolver(
             is FirRegularClass -> {
                 resolveClassTypes(target)
             }
-            is FirFile, is FirAnonymousInitializer -> {}
+            is FirFile, is FirAnonymousInitializer, is FirCodeFragment -> {}
             else -> errorWithAttachment("Unknown declaration ${target::class}") {
                 withFirEntry("declaration", target)
             }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirElementFinder.kt

@@ -105,6 +105,15 @@ internal object FirElementFinder {
                         continue
                     }
 
+                    subDeclaration is FirCodeFragment -> {
+                        val codeFragmentDeclarations = subDeclaration.block.statements.asSequence().filterIsInstance<FirDeclaration>()
+                        if (find(codeFragmentDeclarations.asIterable(), classIdPathIndex)) {
+                            return true
+                        }
+
+                        continue
+                    }
+
                     subDeclaration is FirRegularClass && currentClassSegment == subDeclaration.symbol.name -> {
                         path += subDeclaration
                         if (find(subDeclaration.declarations, classIdPathIndex + 1)) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/containingFileUtils.kt

@@ -26,7 +26,7 @@ fun FirElementWithResolveState.getContainingFile(): FirFile? {
         is FirCallableDeclaration -> provider.getFirCallableContainerFile(symbol)
         is FirClassLikeDeclaration -> provider.getFirClassifierContainerFileIfAny(symbol)
         is FirAnonymousInitializer -> containingClass().getContainingFile()
-        is FirDanglingModifierList -> {
+        is FirDanglingModifierList, is FirCodeFragment -> {
             val ktFile = psi?.containingFile as? KtFile
                 ?: error("File for dangling modifier list cannot be null")
             val moduleComponents = llFirResolvableSession?.moduleComponents

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/declarationUtils.kt

@@ -196,6 +196,12 @@ private fun KtClassLikeDeclaration.findFir(provider: FirProvider): FirClassLikeD
     }
 }
 
+val FirFile.codeFragment: FirCodeFragment
+    get() {
+        return declarations.singleOrNull() as? FirCodeFragment
+            ?: errorWithFirSpecificEntries("Code fragment not found in a FirFile", fir = this)
+    }
+
 val FirDeclaration.isGeneratedDeclaration
     get() = realPsi == null
 

generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/analysisApi.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.compile
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.containingDeclarationProvider.AbstractContainingDeclarationProviderByDelegatedMemberScopeTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.containingDeclarationProvider.AbstractContainingDeclarationProviderByMemberScopeTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.containingDeclarationProvider.AbstractContainingDeclarationProviderByPsiTest
+import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.diagnosticProvider.AbstractCodeFragmentCollectDiagnosticsTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.diagnosticProvider.AbstractCollectDiagnosticsTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.expressionInfoProvider.AbstractIsUsedAsExpressionTest
 import org.jetbrains.kotlin.analysis.api.impl.base.test.cases.components.expressionInfoProvider.AbstractReturnTargetSymbolTest
@@ -265,6 +266,10 @@ private fun AnalysisApiTestGroup.generateAnalysisApiComponentsTests() {
         test(AbstractCollectDiagnosticsTest::class) {
             model("diagnostics")
         }
+
+        test(AbstractCodeFragmentCollectDiagnosticsTest::class, filter = frontendIs(FrontendKind.Fir)) {
+            model("codeFragmentDiagnostics", pattern = TestGeneratorUtil.KT_WITHOUT_DOTS_IN_NAME)
+        }
     }
 
     // for K1, symbols do not have a proper equality implementation, so the tests are failing