diff --git a/idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt b/idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt
index 9cd4ec6eb9c..5d7f4b9e023 100644
--- a/idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt
+++ b/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 {
@@ -104,6 +105,11 @@ class KotlinFoldingBuilder : CustomFoldingBuilder(), DumbAware {
private fun needFolding(node: ASTNode, document: Document): Boolean {
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) ||
@@ -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
diff --git a/idea/testData/folding/checkCollapse/functionWithExpressionBody.kt b/idea/testData/folding/checkCollapse/functionWithExpressionBody.kt
new file mode 100644
index 00000000000..ac4d2851af7
--- /dev/null
+++ b/idea/testData/folding/checkCollapse/functionWithExpressionBody.kt
@@ -0,0 +1,4 @@
+fun test() = if (true)
+ 1
+else
+ 0
\ No newline at end of file
diff --git a/idea/testData/folding/checkCollapse/functionWithExpressionBody2.kt b/idea/testData/folding/checkCollapse/functionWithExpressionBody2.kt
new file mode 100644
index 00000000000..39480cacad1
--- /dev/null
+++ b/idea/testData/folding/checkCollapse/functionWithExpressionBody2.kt
@@ -0,0 +1,5 @@
+fun test() =
+ if (true)
+ 1
+ else
+ 0
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/folding/KotlinFoldingTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/folding/KotlinFoldingTestGenerated.java
index 10d7bf3fcc9..da539d96a62 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/folding/KotlinFoldingTestGenerated.java
+++ b/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");