[AA] handles FIR isUsedAsExpression for return within function block

For the following example:

```
fun foo(bar: Int) {
  <expr>if (bar == 4) return "Four"
  else return "Int"</expr>
}
```

AA FE1.0 `isUsedAsExpression` returns `false`.
Since the current AA FIR `isUsedAsExpression` returns `true` for the
above example, this commit fixes it.

analysis/analysis-api-fe10/tests-gen/org/jetbrains/kotlin/analysis/api/fe10/test/cases/generated/cases/components/expressionInfoProvider/Fe10IdeNormalAnalysisSourceModuleIsUsedAsExpressionTestGenerated.java

@@ -958,6 +958,18 @@ public class Fe10IdeNormalAnalysisSourceModuleIsUsedAsExpressionTestGenerated ex
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/functionLiteralExtensionReceiverType.kt");
     }
 
+    @Test
+    @TestMetadata("function_body_with_if.kt")
+    public void testFunction_body_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_body_with_if.kt");
+    }
+
+    @Test
+    @TestMetadata("function_equal_block_with_if.kt")
+    public void testFunction_equal_block_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_equal_block_with_if.kt");
+    }
+
     @Test
     @TestMetadata("ifBranches_unused.kt")
     public void testIfBranches_unused() throws Exception {
@@ -1510,12 +1522,24 @@ public class Fe10IdeNormalAnalysisSourceModuleIsUsedAsExpressionTestGenerated ex
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_implicit_unit.kt");
     }
 
+    @Test
+    @TestMetadata("return_inside_if.kt")
+    public void testReturn_inside_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_inside_if.kt");
+    }
+
     @Test
     @TestMetadata("return_value.kt")
     public void testReturn_value() throws Exception {
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_value.kt");
     }
 
+    @Test
+    @TestMetadata("return_with_if.kt")
+    public void testReturn_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_with_if.kt");
+    }
+
     @Test
     @TestMetadata("run_block.kt")
     public void testRun_block() throws Exception {

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

@@ -209,8 +209,6 @@ internal class KtFirExpressionInfoProvider(
             is KtFunctionLiteral ->
                 parent.bodyBlockExpression == child && !returnsUnit(parent)
 
-            // Named functions do not use their bodies if the function itself returns unit
-            // UNLESS it's an expression body/lambda of type Unit.
             /** See [doesNamedFunctionUseBody] */
             is KtNamedFunction ->
                 doesNamedFunctionUseBody(parent, child)
@@ -402,23 +400,26 @@ private fun doesPropertyAccessorUseBody(propertyAccessor: KtPropertyAccessor, bo
 }
 
 /**
- *  Named functions do not consider their bodies used if the function itself
+ * Returns whether the function uses its body as an expression (i.e., the function uses the result value of the expression) or not.
- *  returns Unit UNLESS the function body is an expression body and the body is
+ *
- *  of type Unit.
+ * Named functions do not consider their bodies used if
+ *  - the function body is a block e.g., `fun foo(): Int { return bar }` or
+ *  - the function itself returns Unit
  */
-private fun doesNamedFunctionUseBody(namedFunction: KtNamedFunction, body: PsiElement): Boolean =
+private fun doesNamedFunctionUseBody(namedFunction: KtNamedFunction, body: PsiElement): Boolean = when {
-    when {
+    // The body is a block expression e.g., fun foo(): Int { return bar }
-        !returnsUnit(namedFunction) ->
+    namedFunction.bodyBlockExpression == body ->
-            true
+        false
-        namedFunction.bodyBlockExpression == body ->
+    // Note that `namedFunction.hasBlockBody() == false` means the function definition uses `=` e.g., fun foo() = bar
-            false
+    !returnsUnit(namedFunction) ->
-        namedFunction.bodyExpression == body ->
+        true
-            analyze(namedFunction) {
+    namedFunction.bodyExpression == body ->
-                (body as KtExpression).getKtType()?.isUnit == true
+        analyze(namedFunction) {
-            }
+            (body as KtExpression).getKtType()?.isUnit == true
-        else ->
+        }
-            false
+    else ->
-    }
+        false
+}
 
 
 private fun KtAnalysisSession.isSimpleVariableAccessCall(reference: KtReferenceExpression): Boolean =

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

@@ -958,6 +958,18 @@ public class FirIdeDependentAnalysisSourceModuleIsUsedAsExpressionTestGenerated
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/functionLiteralExtensionReceiverType.kt");
     }
 
+    @Test
+    @TestMetadata("function_body_with_if.kt")
+    public void testFunction_body_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_body_with_if.kt");
+    }
+
+    @Test
+    @TestMetadata("function_equal_block_with_if.kt")
+    public void testFunction_equal_block_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_equal_block_with_if.kt");
+    }
+
     @Test
     @TestMetadata("ifBranches_unused.kt")
     public void testIfBranches_unused() throws Exception {
@@ -1510,12 +1522,24 @@ public class FirIdeDependentAnalysisSourceModuleIsUsedAsExpressionTestGenerated
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_implicit_unit.kt");
     }
 
+    @Test
+    @TestMetadata("return_inside_if.kt")
+    public void testReturn_inside_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_inside_if.kt");
+    }
+
     @Test
     @TestMetadata("return_value.kt")
     public void testReturn_value() throws Exception {
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_value.kt");
     }
 
+    @Test
+    @TestMetadata("return_with_if.kt")
+    public void testReturn_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_with_if.kt");
+    }
+
     @Test
     @TestMetadata("run_block.kt")
     public void testRun_block() throws Exception {

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

@@ -958,6 +958,18 @@ public class FirIdeNormalAnalysisSourceModuleIsUsedAsExpressionTestGenerated ext
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/functionLiteralExtensionReceiverType.kt");
     }
 
+    @Test
+    @TestMetadata("function_body_with_if.kt")
+    public void testFunction_body_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_body_with_if.kt");
+    }
+
+    @Test
+    @TestMetadata("function_equal_block_with_if.kt")
+    public void testFunction_equal_block_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_equal_block_with_if.kt");
+    }
+
     @Test
     @TestMetadata("ifBranches_unused.kt")
     public void testIfBranches_unused() throws Exception {
@@ -1510,12 +1522,24 @@ public class FirIdeNormalAnalysisSourceModuleIsUsedAsExpressionTestGenerated ext
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_implicit_unit.kt");
     }
 
+    @Test
+    @TestMetadata("return_inside_if.kt")
+    public void testReturn_inside_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_inside_if.kt");
+    }
+
     @Test
     @TestMetadata("return_value.kt")
     public void testReturn_value() throws Exception {
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_value.kt");
     }
 
+    @Test
+    @TestMetadata("return_with_if.kt")
+    public void testReturn_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_with_if.kt");
+    }
+
     @Test
     @TestMetadata("run_block.kt")
     public void testRun_block() throws Exception {

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

@@ -958,6 +958,18 @@ public class FirStandaloneNormalAnalysisSourceModuleIsUsedAsExpressionTestGenera
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/functionLiteralExtensionReceiverType.kt");
     }
 
+    @Test
+    @TestMetadata("function_body_with_if.kt")
+    public void testFunction_body_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_body_with_if.kt");
+    }
+
+    @Test
+    @TestMetadata("function_equal_block_with_if.kt")
+    public void testFunction_equal_block_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_equal_block_with_if.kt");
+    }
+
     @Test
     @TestMetadata("ifBranches_unused.kt")
     public void testIfBranches_unused() throws Exception {
@@ -1510,12 +1522,24 @@ public class FirStandaloneNormalAnalysisSourceModuleIsUsedAsExpressionTestGenera
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_implicit_unit.kt");
     }
 
+    @Test
+    @TestMetadata("return_inside_if.kt")
+    public void testReturn_inside_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_inside_if.kt");
+    }
+
     @Test
     @TestMetadata("return_value.kt")
     public void testReturn_value() throws Exception {
         runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_value.kt");
     }
 
+    @Test
+    @TestMetadata("return_with_if.kt")
+    public void testReturn_with_if() throws Exception {
+        runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_with_if.kt");
+    }
+
     @Test
     @TestMetadata("run_block.kt")
     public void testRun_block() throws Exception {

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_body_with_if.kt

@@ -0,0 +1,5 @@
+fun test(n: Int) =
+    <expr>if (n == 1)
+        "one"
+    else
+        "two"</expr>

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_body_with_if.txt

@@ -0,0 +1,6 @@
+expression: IF
+text: if (n == 1)
+        "one"
+    else
+        "two"
+isUsedAsExpression: true

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_equal_block_with_if.kt

@@ -0,0 +1,6 @@
+fun test(n: Int) = {
+    <expr>if (n == 1)
+        "one"
+    else
+        "two"</expr>
+}

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/function_equal_block_with_if.txt

@@ -0,0 +1,6 @@
+expression: IF
+text: if (n == 1)
+        "one"
+    else
+        "two"
+isUsedAsExpression: true

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_inside_if.kt

@@ -0,0 +1,6 @@
+fun test(n: Int): String {
+    if (n == 1)
+        <expr>return "one"</expr>
+    else
+        return "two"
+}

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_inside_if.txt

@@ -0,0 +1,3 @@
+expression: RETURN
+text: return "one"
+isUsedAsExpression: false

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_with_if.kt

@@ -0,0 +1,6 @@
+fun test(n: Int): String {
+    <expr>if (n == 1)
+        return "one"
+    else
+        return "two"</expr>
+}

analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/return_with_if.txt

@@ -0,0 +1,6 @@
+expression: IF
+text: if (n == 1)
+        return "one"
+    else
+        return "two"
+isUsedAsExpression: false