[LL API] Support code fragment analysis

2023-06-14 21:57:57 +09:00
parent 3fb038ffd0
commit 02af189066
32 changed files with 493 additions and 32 deletions
@@ -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");
+    }
+}
@@ -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)
+    }
+}
@@ -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");
+    }
+}
@@ -0,0 +1,6 @@
+class Foo {
+    val a = x
+    val b = y
+}
+
+Foo().b
@@ -0,0 +1,4 @@
+fun test() {
+    val x = 0
+    <caret>val y = 0
+}
@@ -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)
@@ -0,0 +1,5 @@
+fun foo(): Pair<Int, Int> {
+    return x to y
+}
+
+foo()
@@ -0,0 +1,6 @@
+// WITH_STDLIB
+
+fun test() {
+    val x = 0
+    <caret>val y = 0
+}
@@ -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)
@@ -0,0 +1 @@
+foo(x)
@@ -0,0 +1,6 @@
+fun test() {
+    val x = 0
+    <caret>val y = 1
+}
+
+fun foo(a: String) {}
@@ -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)
@@ -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)
        }
@@ -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.
@@ -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)
            }
@@ -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
@@ -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

 /**
@@ -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)
@@ -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,
@@ -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)
@@ -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
+    }
 }
@@ -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(
-                session = session,
-                baseScopeProvider = scopeProvider,
-                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()
-            }
+            val result = build(session, scopeProvider, designation, rootNonLocalDeclaration, replacementApplier = replacementApplier)
+            replacementApplier?.ensureApplied()
+            return result
        }

        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)
@@ -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)
    }
@@ -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
@@ -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 -> {}
-            else -> throwUnexpectedFirElementError(target)
+
+        when (target) {
+            is FirRegularClass, is FirScript, is FirCodeFragment -> {}
+            else -> {
+                if (!target.isRegularDeclarationWithAnnotation) {
+                    throwUnexpectedFirElementError(target)
+                }
+            }
        }

        requireIsInstance<FirAnnotationContainer>(target)
@@ -65,6 +65,7 @@ private class LLFirContractsTargetResolver(
            is FirAnonymousInitializer,
            is FirFile,
            is FirScript,
+            is FirCodeFragment,
            is FirFileAnnotationsContainer,
            is FirDanglingModifierList,
            -> {
@@ -78,6 +78,7 @@ internal class LLFirImplicitBodyTargetResolver(
            is FirRegularClass,
            is FirTypeAlias,
            is FirFile,
+            is FirCodeFragment,
            is FirAnonymousInitializer,
            is FirDanglingModifierList,
            is FirFileAnnotationsContainer,
@@ -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)
            }
@@ -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)) {
@@ -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
@@ -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

@@ -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