FIR IDE: fix anonymous function symbol retrieval from function literal

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt

@@ -119,7 +119,7 @@ internal class FirModuleResolveStateImpl(
         val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(nonLocalFirForNamedDeclaration) { firDeclaration ->
             when (val realPsi = firDeclaration.realPsi) {
                 is KtObjectLiteralExpression -> realPsi.objectDeclaration == ktDeclaration
-                is KtFunctionLiteral -> realPsi.parent == ktDeclaration
+                is KtLambdaExpression -> realPsi.functionLiteral == ktDeclaration
                 else -> realPsi == ktDeclaration
             }
         }

idea/idea-frontend-fir/testData/symbols/symbolByPsi/implicitReturnInLambda.kt

@@ -0,0 +1,124 @@
+fun foo() {
+    val lam1 = { a: Int ->
+        val b = 1
+        a + b
+    }
+
+    val lam2 = { a: Int ->
+        val c = 1
+        if (a > 0)
+            a + c
+        else
+            a - c
+    }
+}
+
+// RESULT
+/*
+KtFirValueParameterSymbol:
+  annotatedType: [] kotlin/Int
+  annotationClassIds: []
+  annotations: []
+  callableIdIfNonLocal: null
+  hasDefaultValue: false
+  isExtension: false
+  isVararg: false
+  name: a
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: b
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirAnonymousFunctionSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+  valueParameters: [KtFirValueParameterSymbol(a)]
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Function1<kotlin/Int, kotlin/Int>
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: lam1
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirValueParameterSymbol:
+  annotatedType: [] kotlin/Int
+  annotationClassIds: []
+  annotations: []
+  callableIdIfNonLocal: null
+  hasDefaultValue: false
+  isExtension: false
+  isVararg: false
+  name: a
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: c
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirAnonymousFunctionSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+  valueParameters: [KtFirValueParameterSymbol(a)]
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Function1<kotlin/Int, kotlin/Int>
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: lam2
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirFunctionSymbol:
+  annotatedType: [] kotlin/Unit
+  annotationClassIds: []
+  annotations: []
+  callableIdIfNonLocal: /foo
+  dispatchType: null
+  isExtension: false
+  isExternal: false
+  isInfix: false
+  isInline: false
+  isOperator: false
+  isOverride: false
+  isStatic: false
+  isSuspend: false
+  modality: FINAL
+  name: foo
+  origin: SOURCE
+  receiverType: null
+  symbolKind: TOP_LEVEL
+  typeParameters: []
+  valueParameters: []
+  visibility: Public
+ */

idea/idea-frontend-fir/testData/symbols/symbolByPsi/implicitReturnInLambda.txt

@@ -0,0 +1,106 @@
+KtFirValueParameterSymbol:
+  annotatedType: [] kotlin/Int
+  annotationClassIds: []
+  annotations: []
+  callableIdIfNonLocal: null
+  hasDefaultValue: false
+  isExtension: false
+  isVararg: false
+  name: a
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: b
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirAnonymousFunctionSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+  valueParameters: [KtFirValueParameterSymbol(a)]
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Function1<kotlin/Int, kotlin/Int>
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: lam1
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirValueParameterSymbol:
+  annotatedType: [] kotlin/Int
+  annotationClassIds: []
+  annotations: []
+  callableIdIfNonLocal: null
+  hasDefaultValue: false
+  isExtension: false
+  isVararg: false
+  name: a
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: c
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirAnonymousFunctionSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+  valueParameters: [KtFirValueParameterSymbol(a)]
+
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Function1<kotlin/Int, kotlin/Int>
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: lam2
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+
+KtFirFunctionSymbol:
+  annotatedType: [] kotlin/Unit
+  annotationClassIds: []
+  annotations: []
+  callableIdIfNonLocal: /foo
+  dispatchType: null
+  isExtension: false
+  isExternal: false
+  isInfix: false
+  isInline: false
+  isOperator: false
+  isOverride: false
+  isStatic: false
+  isSuspend: false
+  modality: FINAL
+  name: foo
+  origin: SOURCE
+  receiverType: null
+  symbolKind: TOP_LEVEL
+  typeParameters: []
+  valueParameters: []
+  visibility: Public

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/symbols/SymbolByPsiTestGenerated.java

@@ -96,6 +96,12 @@ public class SymbolByPsiTestGenerated extends AbstractSymbolByPsiTest {
         runTest("idea/idea-frontend-fir/testData/symbols/symbolByPsi/implicitReturn.kt");
     }
 
+    @Test
+    @TestMetadata("implicitReturnInLambda.kt")
+    public void testImplicitReturnInLambda() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/symbols/symbolByPsi/implicitReturnInLambda.kt");
+    }
+
     @Test
     @TestMetadata("localDeclarations.kt")
     public void testLocalDeclarations() throws Exception {