From e81fbe0a0532cb0ac06944f99efc5eca6767d49e Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Sat, 14 Sep 2019 10:22:08 +0900
Subject: [PATCH] Folding: fold function with expression body (KT-6316)

#KT-6316 Fixed
---
 .../jetbrains/kotlin/idea/KotlinFoldingBuilder.kt   | 13 +++++++++++++
 .../checkCollapse/functionWithExpressionBody.kt     |  4 ++++
 .../checkCollapse/functionWithExpressionBody2.kt    |  5 +++++
 .../idea/folding/KotlinFoldingTestGenerated.java    | 10 ++++++++++
 4 files changed, 32 insertions(+)
 create mode 100644 idea/testData/folding/checkCollapse/functionWithExpressionBody.kt
 create mode 100644 idea/testData/folding/checkCollapse/functionWithExpressionBody2.kt

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() = <fold text='{...}' expand='true'>if (true)
+    1
+else
+    0</fold>
\ 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() =
+    <fold text='{...}' expand='true'>if (true)
+        1
+    else
+        0</fold>
\ 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");