[Analysis API] Fix 'getContainingDeclaration()' for code fragments

^KT-65075 Fixed

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolContainingDeclarationProvider.kt

@@ -153,8 +153,8 @@ internal class KtFirSymbolContainingDeclarationProvider(
         }
     }
 
-    fun getContainingDeclarationByPsi(symbol: KtSymbol): KtDeclarationSymbol {
-        val containingDeclaration = getContainingPsi(symbol)
+    fun getContainingDeclarationByPsi(symbol: KtSymbol): KtDeclarationSymbol? {
+        val containingDeclaration = getContainingPsi(symbol) ?: return null
         return with(analysisSession) { containingDeclaration.getSymbol() }
     }
 
@@ -221,7 +221,7 @@ internal class KtFirSymbolContainingDeclarationProvider(
         return symbol.getContainingKtModule(analysisSession.firResolveSession)
     }
 
-    private fun getContainingPsi(symbol: KtSymbol): KtDeclaration {
+    private fun getContainingPsi(symbol: KtSymbol): KtDeclaration? {
         val source = symbol.firSymbol.source
             ?: errorWithAttachment("PSI should present for declaration built by Kotlin code") {
                 withSymbolAttachment("symbolForContainingPsi", symbol, analysisSession)
@@ -245,10 +245,21 @@ internal class KtFirSymbolContainingDeclarationProvider(
         }
 
         if (isOrdinarySymbolWithSource(symbol)) {
-            return psi.getContainingKtDeclaration()
-                ?: errorWithAttachment("Containing declaration should present for nested declaration ${psi::class}") {
+            val result = psi.getContainingPsiDeclaration()
+
+            if (result == null) {
+                val containingFile = psi.containingFile
+                if (containingFile is KtCodeFragment) {
+                    // All content inside a code fragment is implicitly local, but there is no non-local parent
+                    return null
+                }
+
+                errorWithAttachment("Containing declaration should present for nested declaration ${psi::class}") {
                     withSymbolAttachment("symbolForContainingPsi", symbol, analysisSession)
                 }
+            }
+
+            return result
         }
 
         errorWithAttachment("Unsupported declaration origin ${symbol.origin} ${psi::class}") {

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

@@ -52,6 +52,12 @@ public class FirIdeDependentAnalysisSourceModuleContainingDeclarationProviderByP
         runTest("analysis/analysis-api/testData/components/containingDeclarationProvider/containingDeclarationByPsi/classes.kt");
     }
 
+    @Test
+    @TestMetadata("codeFragments.kt")
+    public void testCodeFragments() throws Exception {
+        runTest("analysis/analysis-api/testData/components/containingDeclarationProvider/containingDeclarationByPsi/codeFragments.kt");
+    }
+
     @Test
     @TestMetadata("deeplyNestedCode.kt")
     public void testDeeplyNestedCode() throws Exception {

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

@@ -52,6 +52,12 @@ public class FirIdeNormalAnalysisSourceModuleContainingDeclarationProviderByPsiT
         runTest("analysis/analysis-api/testData/components/containingDeclarationProvider/containingDeclarationByPsi/classes.kt");
     }
 
+    @Test
+    @TestMetadata("codeFragments.kt")
+    public void testCodeFragments() throws Exception {
+        runTest("analysis/analysis-api/testData/components/containingDeclarationProvider/containingDeclarationByPsi/codeFragments.kt");
+    }
+
     @Test
     @TestMetadata("deeplyNestedCode.kt")
     public void testDeeplyNestedCode() throws Exception {

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/components/containingDeclarationProvider/AbstractContainingDeclarationProviderByPsiTest.kt

@@ -20,7 +20,7 @@ abstract class AbstractContainingDeclarationProviderByPsiTest : AbstractAnalysis
     override fun doTestByMainFile(mainFile: KtFile, mainModule: TestModule, testServices: TestServices) {
         val currentPath = mutableListOf<KtDeclaration>()
         val ktClasses = mutableListOf<KtClassOrObject>()
-        analyseForTest(mainFile.declarations.first()) {
+        analyseForTest(mainFile) {
             val expectedFileSymbol = mainFile.getFileSymbol()
             mainFile.accept(object : KtVisitorVoid() {
                 override fun visitElement(element: PsiElement) {

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

@@ -52,6 +52,12 @@ public class FirStandaloneNormalAnalysisSourceModuleContainingDeclarationProvide
         runTest("analysis/analysis-api/testData/components/containingDeclarationProvider/containingDeclarationByPsi/classes.kt");
     }
 
+    @Test
+    @TestMetadata("codeFragments.kt")
+    public void testCodeFragments() throws Exception {
+        runTest("analysis/analysis-api/testData/components/containingDeclarationProvider/containingDeclarationByPsi/codeFragments.kt");
+    }
+
     @Test
     @TestMetadata("deeplyNestedCode.kt")
     public void testDeeplyNestedCode() throws Exception {

analysis/analysis-api/testData/components/containingDeclarationProvider/containingDeclarationByPsi/codeFragments.kt

@@ -0,0 +1,22 @@
+// MODULE: context
+
+// FILE: context.kt
+fun test() {
+    <caret_context>Unit
+}
+
+fun foo() {}
+
+
+// MODULE: main
+// MODULE_KIND: CodeFragment
+
+// FILE: main.kt
+// CODE_FRAGMENT_KIND: EXPRESSION
+object {
+    fun foo() {}
+
+    fun bar() {
+        foo()
+    }
+}