Folding: fold function with expression body (KT-6316)

#KT-6316 Fixed

idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt

@@ -40,6 +40,7 @@ import org.jetbrains.kotlin.psi.psiUtil.allChildren
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.referenceExpression
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.psi.stubs.elements.KtFunctionElementType
 import org.jetbrains.kotlin.resolve.calls.components.isVararg
 
 class KotlinFoldingBuilder : CustomFoldingBuilder(), DumbAware {
@@ -105,6 +106,11 @@ class KotlinFoldingBuilder : CustomFoldingBuilder(), DumbAware {
         val type = node.elementType
         val parentType = node.treeParent?.elementType
         
+        if (type is KtFunctionElementType) {
+            val bodyExpression = (node.psi as? KtNamedFunction)?.bodyExpression
+            if (bodyExpression != null && bodyExpression !is KtBlockExpression) return true
+        }
+
         return type == KtNodeTypes.FUNCTION_LITERAL ||
                 (type == KtNodeTypes.BLOCK && parentType != KtNodeTypes.FUNCTION_LITERAL) ||
                 type == KtNodeTypes.CLASS_BODY || type == KtTokens.BLOCK_COMMENT || type == KDocTokens.KDOC ||
@@ -132,6 +138,13 @@ class KotlinFoldingBuilder : CustomFoldingBuilder(), DumbAware {
     }
 
     private fun getRangeToFold(node: ASTNode): TextRange {
+        if (node.elementType is KtFunctionElementType) {
+            val bodyExpression = (node.psi as? KtNamedFunction)?.bodyExpression
+            if (bodyExpression != null && bodyExpression !is KtBlockExpression) {
+                return bodyExpression.textRange
+            }
+        }
+        
         if (node.elementType == KtNodeTypes.FUNCTION_LITERAL) {
             val psi = node.psi as? KtFunctionLiteral
             val lbrace = psi?.lBrace

idea/testData/folding/checkCollapse/functionWithExpressionBody.kt

@@ -0,0 +1,4 @@
+fun test() = <fold text='{...}' expand='true'>if (true)
+    1
+else
+    0</fold>

idea/testData/folding/checkCollapse/functionWithExpressionBody2.kt

@@ -0,0 +1,5 @@
+fun test() =
+    <fold text='{...}' expand='true'>if (true)
+        1
+    else
+        0</fold>

idea/tests/org/jetbrains/kotlin/idea/folding/KotlinFoldingTestGenerated.java

@@ -118,6 +118,16 @@ public class KotlinFoldingTestGenerated extends AbstractKotlinFoldingTest {
             runTest("idea/testData/folding/checkCollapse/functionLiteral.kt");
         }
 
+        @TestMetadata("functionWithExpressionBody.kt")
+        public void testFunctionWithExpressionBody() throws Exception {
+            runTest("idea/testData/folding/checkCollapse/functionWithExpressionBody.kt");
+        }
+
+        @TestMetadata("functionWithExpressionBody2.kt")
+        public void testFunctionWithExpressionBody2() throws Exception {
+            runTest("idea/testData/folding/checkCollapse/functionWithExpressionBody2.kt");
+        }
+
         @TestMetadata("headerKDoc.kt")
         public void testHeaderKDoc() throws Exception {
             runTest("idea/testData/folding/checkCollapse/headerKDoc.kt");