[LL API] Fix analysis of functions inside code fragments

FirDeclarationsResolveTransformer.transformSimpleFunction() performs
proper local function analysis only if the parent is properly set.

^KTIJ-26608 Fixed

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

@@ -91,6 +91,12 @@ public class FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerate
             runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/duplicateLibrary.kt");
         }
 
+        @Test
+        @TestMetadata("fragmentFunction.kt")
+        public void testFragmentFunction() throws Exception {
+            runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/fragmentFunction.kt");
+        }
+
         @Test
         @TestMetadata("imports.kt")
         public void testImports() throws Exception {

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

@@ -91,6 +91,12 @@ public class FirIdeNormalAnalysisSourceModuleCompilerFacilityTestGenerated exten
             runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/duplicateLibrary.kt");
         }
 
+        @Test
+        @TestMetadata("fragmentFunction.kt")
+        public void testFragmentFunction() throws Exception {
+            runTest("analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/fragmentFunction.kt");
+        }
+
         @Test
         @TestMetadata("imports.kt")
         public void testImports() throws Exception {

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/fragmentFunction.fragment.kt

@@ -0,0 +1,2 @@
+fun String.ext() = this + "EXT"
+"base".ext()

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/fragmentFunction.ir.txt

@@ -0,0 +1,19 @@
+MODULE_FRAGMENT
+  FILE fqName:<root> fileName:fragment.kt
+    CLASS CLASS name:CodeFragment modality:FINAL visibility:public superTypes:[kotlin.Any]
+      $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.CodeFragment
+      CONSTRUCTOR visibility:public <> () returnType:<root>.CodeFragment [primary]
+        BLOCK_BODY
+          DELEGATING_CONSTRUCTOR_CALL 'public constructor <init> () [primary] declared in kotlin.Any'
+      FUN name:run visibility:public modality:FINAL <> () returnType:kotlin.String
+        EXPRESSION_BODY
+          BLOCK type=kotlin.String origin=null
+            FUN LOCAL_FUNCTION name:ext visibility:local modality:FINAL <> ($receiver:kotlin.String) returnType:kotlin.String
+              $receiver: VALUE_PARAMETER name:<this> type:kotlin.String
+              BLOCK_BODY
+                RETURN type=kotlin.Nothing from='local final fun ext (): kotlin.String declared in <root>.CodeFragment.run'
+                  CALL 'public final fun plus (other: kotlin.Any?): kotlin.String [operator] declared in kotlin.String' type=kotlin.String origin=PLUS
+                    $this: GET_VAR '<this>: kotlin.String declared in <root>.CodeFragment.run.ext' type=kotlin.String origin=null
+                    other: CONST String type=kotlin.String value="EXT"
+            CALL 'local final fun ext (): kotlin.String declared in <root>.CodeFragment.run' type=kotlin.String origin=null
+              $receiver: CONST String type=kotlin.String value="base"

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/fragmentFunction.kt

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

analysis/analysis-api/testData/components/compilerFacility/compilation/codeFragments/fragmentFunction.txt

@@ -0,0 +1,6 @@
+public final class CodeFragment {
+    // source: 'fragment.kt'
+    public method <init>(): void
+    private final static method run$ext(p0: java.lang.String): java.lang.String
+    public final static method run(): java.lang.String
+}

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

@@ -1,7 +1,7 @@
-fun foo() {}
+fun foo(a: Int, b: String) {}
 
 fun bar() {
-    <expr>foo()</expr>
+    <expr>foo(1, "foo")</expr>
 }
 
 bar()

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

@@ -1,10 +1,10 @@
 BEFORE MODIFICATION:
 CODE FRAGMENT: {
-    local final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| {
+    local final [ResolvedTo(RAW_FIR)] fun foo([ResolvedTo(RAW_FIR)] a: R|kotlin/Int|, [ResolvedTo(RAW_FIR)] b: R|kotlin/String|): R|kotlin/Unit| {
     }
 
-    local final? [ResolvedTo(RAW_FIR)] fun bar(): R|kotlin/Unit| {
-        R|/foo|()
+    local final [ResolvedTo(RAW_FIR)] fun bar(): R|kotlin/Unit| {
+        R|/foo|(Int(1), String(foo))
     }
 
     R|/bar|()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -541,7 +541,8 @@ class BodyResolveContext(
         }
     }
 
-    fun <T> withScopesForCodeFragment(codeFragment: FirCodeFragment, holder: SessionHolder, f: () -> T): T {
+    @OptIn(PrivateForInline::class)
+    fun <T> withCodeFragment(codeFragment: FirCodeFragment, holder: SessionHolder, f: () -> T): T {
         val codeFragmentContext = codeFragment.codeFragmentContext ?: error("Context is not set for a code fragment")
         val towerDataContext = codeFragmentContext.towerDataContext
 
@@ -566,7 +567,7 @@ class BodyResolveContext(
         )
 
         return withTowerDataContexts(newContext) {
-            f()
+            withContainer(codeFragment, f)
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -631,7 +631,7 @@ open class FirDeclarationsResolveTransformer(
 
     override fun transformCodeFragment(codeFragment: FirCodeFragment, data: ResolutionMode): FirCodeFragment {
         dataFlowAnalyzer.enterCodeFragment(codeFragment)
-        context.withScopesForCodeFragment(codeFragment, components) {
+        context.withCodeFragment(codeFragment, components) {
             transformBlock(codeFragment.block, data)
         }
         dataFlowAnalyzer.exitCodeFragment()