[LL API] Add in-body modification tests for code fragments

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localClass.fragment.kt

@@ -0,0 +1,7 @@
+class Foo(val x: Int) {
+    fun foo() {
+        <expr>foo()</expr>
+    }
+}
+
+Foo(1)

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localClass.kt

@@ -0,0 +1,3 @@
+fun test() {
+    <caret>Unit
+}

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localClass.txt

@@ -0,0 +1,21 @@
+BEFORE MODIFICATION:
+CODE FRAGMENT: {
+    local final [ResolvedTo(RAW_FIR)] class Foo : R|kotlin/Any| {
+        public [ResolvedTo(RAW_FIR)] [ContainingClassKey=Foo] constructor([ResolvedTo(RAW_FIR)] [CorrespondingProperty=<local>/x] x: R|kotlin/Int|): R|Foo| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final [ResolvedTo(RAW_FIR)] [IsFromPrimaryConstructor=true] val x: R|kotlin/Int| = R|<local>/x|
+            public [ResolvedTo(RAW_FIR)] [ContainingClassKey=Foo] get(): R|kotlin/Int|
+
+        public final [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| {
+            this@R|/Foo|.R|<local>/foo|()
+        }
+
+    }
+
+    R|<local>/Foo.Foo|(Int(1))
+}
+
+AFTER MODIFICATION:
+CODE FRAGMENT: { LAZY_BLOCK }

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localFunction.fragment.kt

@@ -0,0 +1,7 @@
+fun foo() {}
+
+fun bar() {
+    <expr>foo()</expr>
+}
+
+bar()

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localFunction.kt

@@ -0,0 +1,3 @@
+fun test() {
+    <caret>Unit
+}

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localFunction.txt

@@ -0,0 +1,14 @@
+BEFORE MODIFICATION:
+CODE FRAGMENT: {
+    local final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| {
+    }
+
+    local final? [ResolvedTo(RAW_FIR)] fun bar(): R|kotlin/Unit| {
+        R|/foo|()
+    }
+
+    R|/bar|()
+}
+
+AFTER MODIFICATION:
+CODE FRAGMENT: { LAZY_BLOCK }

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/simple.fragment.kt

@@ -0,0 +1 @@
+<expr>foo</expr>

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/simple.kt

@@ -0,0 +1,10 @@
+package test
+
+class Foo
+
+fun test() {
+    val foo = Foo()
+    <caret>consume(foo)
+}
+
+fun consume(foo: Foo) {}

analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/simple.txt

@@ -0,0 +1,7 @@
+BEFORE MODIFICATION:
+CODE FRAGMENT: {
+    R|<local>/foo|
+}
+
+AFTER MODIFICATION:
+CODE FRAGMENT: { LAZY_BLOCK }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractLowLevelApiCodeFragmentTest.kt

@@ -8,12 +8,12 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
-import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.test.directives.model.RegisteredDirectives
-import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.model.TestFile
-import org.jetbrains.kotlin.psi.KtPsiFactory
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.sourceFileProvider
 import java.io.File
 
 abstract class AbstractLowLevelApiCodeFragmentTest : AbstractLowLevelApiSingleFileTest() {
@@ -22,20 +22,30 @@ abstract class AbstractLowLevelApiCodeFragmentTest : AbstractLowLevelApiSingleFi
     abstract fun doTest(ktCodeFragment: KtCodeFragment, moduleStructure: TestModuleStructure, testServices: TestServices)
 
     final override fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices) {
-        val contextElement = testServices.expressionMarkerProvider.getElementOfTypeAtCaret<KtDeclaration>(ktFile)
+        val contextElement = testServices.expressionMarkerProvider.getElementOfTypeAtCaret<KtElement>(ktFile)
 
-        val fragmentText = moduleStructure.originalTestDataFiles.single()
+        val fragmentFile = moduleStructure.originalTestDataFiles.single()
             .run { File(parent, "$nameWithoutExtension.fragment.$extension") }
-            .readText()
 
-        val isBlockFragment = fragmentText.any { it == '\n' }
+        val fragmentTestFile = TestFile(
+            "fragment.kt",
+            fragmentFile.readText(),
+            fragmentFile,
+            startLineNumberInOriginalFile = 0,
+            isAdditional = true,
+            RegisteredDirectives.Empty
+        )
+
+        val processedFragmentText = testServices.sourceFileProvider.getContentOfSourceFile(fragmentTestFile)
+
+        val isBlockFragment = processedFragmentText.any { it == '\n' }
 
         val project = ktFile.project
         val factory = KtPsiFactory(project, markGenerated = false)
 
         val ktCodeFragment = when {
-            isBlockFragment -> factory.createBlockCodeFragment(fragmentText, contextElement)
+            isBlockFragment -> factory.createBlockCodeFragment(processedFragmentText, contextElement)
-            else -> factory.createExpressionCodeFragment(fragmentText, contextElement)
+            else -> factory.createExpressionCodeFragment(processedFragmentText, contextElement)
         }
 
         doTest(ktCodeFragment, moduleStructure, testServices)

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/AbstractCodeFragmentInBlockModificationTest.kt

@@ -0,0 +1,29 @@
+/*
+ * 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.low.level.api.fir.file.structure
+
+import org.jetbrains.kotlin.analysis.low.level.api.fir.AbstractLowLevelApiCodeFragmentTest
+import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
+import org.jetbrains.kotlin.psi.KtCodeFragment
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.test.services.TestModuleStructure
+import org.jetbrains.kotlin.test.services.TestServices
+import org.jetbrains.kotlin.test.services.assertions
+import kotlin.test.assertNull
+
+abstract class AbstractCodeFragmentInBlockModificationTest : AbstractLowLevelApiCodeFragmentTest() {
+    override fun doTest(ktCodeFragment: KtCodeFragment, moduleStructure: TestModuleStructure, testServices: TestServices) {
+        val targetElement = testServices.expressionMarkerProvider
+            .getBottommostSelectedElementOfType(ktCodeFragment, KtElement::class.java)
+
+        assertNull(targetElement.getNonLocalReanalyzableContainingDeclaration())
+
+        val (before, after) = testInBlockModification(ktCodeFragment, ktCodeFragment, testServices)
+        val actualText = "BEFORE MODIFICATION:\n$before\nAFTER MODIFICATION:\n$after"
+
+        testServices.assertions.assertEqualsToTestDataFileSibling(actualText)
+    }
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/AbstractInBlockModificationTest.kt

@@ -16,7 +16,8 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.Analys
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
-import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtAnnotated
+import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.kotlin.test.services.TestModuleStructure
@@ -32,43 +33,51 @@ abstract class AbstractInBlockModificationTest : AbstractLowLevelApiSingleFileTe
 
         val declaration = selectedElement.getNonLocalReanalyzableContainingDeclaration()
         val actual = if (declaration != null) {
-            resolveWithCaches(ktFile) { firSession ->
+            val (before, after) = testInBlockModification(ktFile, declaration, testServices)
-                val firDeclarationBefore = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
+            "BEFORE MODIFICATION:\n$before\nAFTER MODIFICATION:\n$after"
-                val declarationTextBefore = firDeclarationBefore.render()
-
-                declaration.modifyBody()
-                invalidateAfterInBlockModification(declaration)
-
-                val declarationTextAfterModification = firDeclarationBefore.render()
-                testServices.assertions.assertNotEquals(declarationTextBefore, declarationTextAfterModification) {
-                    "The declaration before and after modification must be in different state"
-                }
-
-                val firDeclarationAfter = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
-                testServices.assertions.assertEquals(firDeclarationBefore, firDeclarationAfter) {
-                    "The declaration before and after must be the same"
-                }
-
-                val declarationTextAfter = firDeclarationAfter.render()
-                testServices.assertions.assertEquals(declarationTextBefore, declarationTextAfter) {
-                    "The declaration must have the same in the resolved state"
-                }
-
-                "BEFORE MODIFICATION:\n$declarationTextBefore\nAFTER MODIFICATION:\n$declarationTextAfterModification"
-            }
         } else {
             "IN-BLOCK MODIFICATION IS NOT APPLICABLE FOR THIS PLACE"
         }
 
         testServices.assertions.assertEqualsToTestDataFileSibling(actual)
     }
+}
 
-    /**
+internal fun testInBlockModification(file: KtFile, declaration: KtAnnotated, testServices: TestServices): Pair<String, String> {
-     * Emulate modification inside the body
+    return resolveWithCaches(file) { firSession ->
-     */
+        val firDeclarationBefore = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
-    private fun KtDeclaration.modifyBody() {
+        val declarationTextBefore = firDeclarationBefore.render()
-        parentsWithSelf.filterIsInstance<ASTDelegatePsiElement>().forEach {
+
-            it.subtreeChanged()
+        declaration.modifyBody()
+        invalidateAfterInBlockModification(declaration)
+
+        val declarationTextAfterModification = firDeclarationBefore.render()
+        testServices.assertions.assertNotEquals(declarationTextBefore, declarationTextAfterModification) {
+            "The declaration before and after modification must be in different state"
+        }
+
+        val firDeclarationAfter = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
+        testServices.assertions.assertEquals(firDeclarationBefore, firDeclarationAfter) {
+            "The declaration before and after must be the same"
+        }
+
+        val declarationTextAfter = firDeclarationAfter.render()
+        testServices.assertions.assertEquals(declarationTextBefore, declarationTextAfter) {
+            "The declaration must have the same in the resolved state"
+        }
+
+        Pair(declarationTextBefore, declarationTextAfterModification)
+    }
+}
+
+/**
+ * Emulate modification inside the body
+ */
+private fun KtAnnotated.modifyBody() {
+    for (parent in parentsWithSelf) {
+        when (parent) {
+            is ASTDelegatePsiElement -> parent.subtreeChanged()
+            is KtCodeFragment -> parent.subtreeChanged()
         }
     }
 }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/CodeFragmentInBlockModificationTestGenerated.java

@@ -0,0 +1,44 @@
+/*
+ * 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.low.level.api.fir.file.structure;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.util.KtTestUtil;
+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/low-level-api-fir/testdata/inBlockModification/codeFragments")
+@TestDataPath("$PROJECT_ROOT")
+public class CodeFragmentInBlockModificationTestGenerated extends AbstractCodeFragmentInBlockModificationTest {
+    @Test
+    public void testAllFilesPresentInCodeFragments() throws Exception {
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/inBlockModification/codeFragments"), Pattern.compile("^([^.]+)\\.kt$"), null, false);
+    }
+
+    @Test
+    @TestMetadata("localClass.kt")
+    public void testLocalClass() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localClass.kt");
+    }
+
+    @Test
+    @TestMetadata("localFunction.kt")
+    public void testLocalFunction() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/localFunction.kt");
+    }
+
+    @Test
+    @TestMetadata("simple.kt")
+    public void testSimple() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/inBlockModification/codeFragments/simple.kt");
+    }
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/OutOfContentRootInBlockModificationTestGenerated.java

@@ -21,7 +21,7 @@ import java.util.regex.Pattern;
 public class OutOfContentRootInBlockModificationTestGenerated extends AbstractOutOfContentRootInBlockModificationTest {
     @Test
     public void testAllFilesPresentInInBlockModification() throws Exception {
-        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/inBlockModification"), Pattern.compile("^(.+)\\.(kt)$"), null, true);
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/inBlockModification"), Pattern.compile("^(.+)\\.(kt)$"), null, false);
     }
 
     @Test

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/ScriptInBlockModificationTestGenerated.java

@@ -21,7 +21,7 @@ import java.util.regex.Pattern;
 public class ScriptInBlockModificationTestGenerated extends AbstractScriptInBlockModificationTest {
     @Test
     public void testAllFilesPresentInInBlockModification() throws Exception {
-        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/inBlockModification"), Pattern.compile("^(.+)\\.(kts)$"), null, true);
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/inBlockModification"), Pattern.compile("^(.+)\\.(kts)$"), null, false);
     }
 
     @Test

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/SourceInBlockModificationTestGenerated.java

@@ -21,7 +21,7 @@ import java.util.regex.Pattern;
 public class SourceInBlockModificationTestGenerated extends AbstractSourceInBlockModificationTest {
     @Test
     public void testAllFilesPresentInInBlockModification() throws Exception {
-        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/inBlockModification"), Pattern.compile("^(.+)\\.(kt)$"), null, true);
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/inBlockModification"), Pattern.compile("^(.+)\\.(kt)$"), null, false);
     }
 
     @Test

compiler/fir/tree/src/org/jetbrains/kotlin/fir/renderer/FirRenderer.kt

@@ -202,6 +202,11 @@ class FirRenderer(
             printer.popIndent()
         }
 
+        override fun visitCodeFragment(codeFragment: FirCodeFragment) {
+            printer.print("CODE FRAGMENT:")
+            bodyRenderer?.renderBody(codeFragment.block)
+        }
+
         override fun visitFileAnnotationsContainer(fileAnnotationsContainer: FirFileAnnotationsContainer) {
             if (fileAnnotationsContainerRenderer != null) {
                 fileAnnotationsContainerRenderer.render(fileAnnotationsContainer)

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

@@ -15,12 +15,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostic.AbstractSource
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostic.compiler.based.AbstractDiagnosticCompilerTestDataTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostic.compiler.based.AbstractLLFirPreresolvedReversedDiagnosticCompilerTestDataSpecTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostic.compiler.based.AbstractLLFirPreresolvedReversedDiagnosticCompilerTestDataTest
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.AbstractOutOfContentRootFileStructureTest
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.*
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.AbstractOutOfContentRootInBlockModificationTest
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.AbstractScriptFileStructureTest
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.AbstractScriptInBlockModificationTest
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.AbstractSourceFileStructureTest
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.AbstractSourceInBlockModificationTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.AbstractErrorResistanceTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.AbstractOutOfContentRootWholeFileResolvePhaseTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolve.AbstractScriptWholeFileResolvePhaseTest
@@ -69,15 +64,19 @@ internal fun TestGroupSuite.generateFirLowLevelApiTests() {
         }
 
         testClass<AbstractSourceInBlockModificationTest> {
-            model("inBlockModification", pattern = TestGeneratorUtil.KT)
+            model("inBlockModification", recursive = false, pattern = TestGeneratorUtil.KT)
         }
 
         testClass<AbstractOutOfContentRootInBlockModificationTest> {
-            model("inBlockModification", pattern = TestGeneratorUtil.KT)
+            model("inBlockModification", recursive = false, pattern = TestGeneratorUtil.KT)
         }
 
         testClass<AbstractScriptInBlockModificationTest> {
-            model("inBlockModification", pattern = TestGeneratorUtil.KTS)
+            model("inBlockModification", recursive = false, pattern = TestGeneratorUtil.KTS)
+        }
+
+        testClass<AbstractCodeFragmentInBlockModificationTest> {
+            model("inBlockModification/codeFragments", recursive = false, pattern = TestGeneratorUtil.KT_WITHOUT_DOTS_IN_NAME)
         }
 
         testClass<AbstractSourceFileStructureTest> {