Fix header compilation for function / class declarations nested inside inline functions.

^KT-78422

Author: aviroop-123

compiler/fir/analysis-tests/testData/resolve/headerMode/classDeclaration.fir.txt

@@ -22,7 +22,9 @@ FILE: classDeclaration.kt
         public final fun funD(): R|kotlin/Int|
 
         public final inline fun funE(): R|kotlin/String| {
-            local final fun funF(): R|kotlin/String|
+            local final fun funF(): R|kotlin/String| {
+                ^funF String(funF body)
+            }
 
             ^funE R|<local>/funF|()
         }

compiler/fir/analysis-tests/testData/resolve/headerMode/enumClassDeclaration.fir.txt

@@ -24,7 +24,9 @@ FILE: enumClassDeclaration.kt
                 super<R|B|>()
             }
 
-            public open override fun getString(): R|kotlin/String|
+            public open override fun getString(): R|kotlin/String| {
+                ^getString String(north)
+            }
 
         }
 
@@ -33,7 +35,9 @@ FILE: enumClassDeclaration.kt
                 super<R|B|>()
             }
 
-            public open override fun getString(): R|kotlin/String|
+            public open override fun getString(): R|kotlin/String| {
+                ^getString String(south)
+            }
 
         }
 

compiler/fir/analysis-tests/testData/resolve/headerMode/functionDeclaration.fir.txt

@@ -12,7 +12,9 @@ FILE: functionDeclaration.kt
     private final fun funC(): R|kotlin/String|
     public final fun funD(): R|kotlin/Int|
     public final inline fun funE(): R|kotlin/String| {
-        local final fun funF(): R|kotlin/String|
+        local final fun funF(): R|kotlin/String| {
+            ^funF String(funF body)
+        }
 
         ^funE R|<local>/funF|()
     }
@@ -22,7 +24,9 @@ FILE: functionDeclaration.kt
                 super<R|kotlin/Any|>()
             }
 
-            public final fun funH(): R|kotlin/String|
+            public final fun funH(): R|kotlin/String| {
+                ^funH String(funH body)
+            }
 
         }
 

compiler/fir/analysis-tests/testData/resolve/headerMode/objectDeclaration.fir.txt

@@ -22,7 +22,9 @@ FILE: objectDeclaration.kt
         public final fun funD(): R|kotlin/Int|
 
         public final inline fun funE(): R|kotlin/String| {
-            local final fun funF(): R|kotlin/String|
+            local final fun funF(): R|kotlin/String| {
+                ^funF String(funF body)
+            }
 
             ^funE R|<local>/funF|()
         }

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/LightTreeRawFirDeclarationBuilder.kt

@@ -2058,7 +2058,8 @@ class LightTreeRawFirDeclarationBuilder(
                     }
 
                     val allowLegacyContractDescription = outerContractDescription == null
-                    val bodyWithContractDescription = withForcedLocalContext {
+                    val forceKeepingTheBodyInHeaderMode = context.forceKeepingTheBodyInHeaderMode || functionBuilder.status.isInline
+                    val bodyWithContractDescription = withForcedLocalContext(forceKeepingTheBodyInHeaderMode) {
                         convertFunctionBody(block, expression, allowLegacyContractDescription)
                     }
                     this.body = bodyWithContractDescription.first

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/AbstractRawFirBuilder.kt

@@ -126,7 +126,9 @@ abstract class AbstractRawFirBuilder<T : Any>(val baseSession: FirSession, val c
         }
     }
 
-    inline fun <R> withForcedLocalContext(block: () -> R): R {
+    inline fun <R> withForcedLocalContext(forceKeepingTheBodyInHeaderMode: Boolean = false, block: () -> R): R {
+        val oldForceKeepingTheBodyInHeaderMode = context.forceKeepingTheBodyInHeaderMode
+        context.forceKeepingTheBodyInHeaderMode = oldForceKeepingTheBodyInHeaderMode || forceKeepingTheBodyInHeaderMode
         val oldForcedLocalContext = context.inLocalContext
         context.inLocalContext = true
         val oldClassNameBeforeLocalContext = context.classNameBeforeLocalContext
@@ -141,6 +143,7 @@ abstract class AbstractRawFirBuilder<T : Any>(val baseSession: FirSession, val c
             context.classNameBeforeLocalContext = oldClassNameBeforeLocalContext
             context.inLocalContext = oldForcedLocalContext
             context.className = oldClassName
+            context.forceKeepingTheBodyInHeaderMode = oldForceKeepingTheBodyInHeaderMode
         }
     }
 

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/Context.kt

@@ -59,6 +59,8 @@ class Context<T> {
     val dispatchReceiverTypesStack: MutableList<ConeClassLikeType> = mutableListOf()
     var containerIsExpect: Boolean = false
 
+    var forceKeepingTheBodyInHeaderMode: Boolean = false
+
     var containingScriptSymbol: FirScriptSymbol? = null
     var containingReplSymbol: FirReplSnippetSymbol? = null
 

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

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.fir.declarations.utils.hasExplicitBackingField
 import org.jetbrains.kotlin.fir.declarations.utils.isConst
 import org.jetbrains.kotlin.fir.declarations.utils.isInline
 import org.jetbrains.kotlin.fir.declarations.utils.isLocal
+import org.jetbrains.kotlin.fir.declarations.utils.isNonLocal
 import org.jetbrains.kotlin.fir.declarations.utils.isScriptTopLevelDeclaration
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
@@ -1039,7 +1040,7 @@ open class FirDeclarationsResolveTransformer(
                 }
             result.transformReturnTypeRef(transformer, ResolutionMode.UpdateImplicitTypeRef(returnTypeRef))
         }
-        if (session.languageVersionSettings.getFlag(AnalysisFlags.headerMode) && !function.isInline) {
+        if (session.languageVersionSettings.getFlag(AnalysisFlags.headerMode) && !function.isInline && function.isNonLocal) {
             // Header mode: once the return type for non-inline function is known, the body can be removed.
             result.replaceBody(null)
         }