From 72648d305e031b572e147e646b101aa0b23e8876 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Tue, 21 Apr 2015 20:47:55 +0200
Subject: [PATCH] better folding for function literals  #KT-4090 Fixed

---
 .../kotlin/idea/KotlinFoldingBuilder.kt       | 37 ++++++++++++++-----
 .../folding/checkCollapse/functionLiteral.kt  | 12 ++++++
 .../folding/KotlinFoldingTestGenerated.java   |  6 +++
 3 files changed, 46 insertions(+), 9 deletions(-)
 create mode 100644 idea/testData/folding/checkCollapse/functionLiteral.kt

diff --git a/idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt b/idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt
index 93e12655b70..194e5a16685 100644
--- a/idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/KotlinFoldingBuilder.kt
@@ -25,14 +25,12 @@ import com.intellij.openapi.project.DumbAware
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiWhiteSpace
-import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.JetNodeTypes
 import org.jetbrains.kotlin.kdoc.lexer.KDocTokens
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.JetFile
-import org.jetbrains.kotlin.psi.JetImportDirective
+import org.jetbrains.kotlin.psi.JetFunctionLiteral
 import org.jetbrains.kotlin.psi.JetImportList
-
 import java.util.ArrayList
 
 public class KotlinFoldingBuilder : FoldingBuilderEx(), DumbAware {
@@ -61,19 +59,40 @@ public class KotlinFoldingBuilder : FoldingBuilderEx(), DumbAware {
     }
 
     private fun appendDescriptors(node: ASTNode, document: Document, descriptors: MutableList<FoldingDescriptor>) {
-        val textRange = node.getTextRange()
-        val type = node.getElementType()
-        if ((type == JetNodeTypes.BLOCK || type == JetNodeTypes.CLASS_BODY || type == JetTokens.BLOCK_COMMENT || type == KDocTokens.KDOC) &&
-                !isOneLine(textRange, document)) {
-            descriptors.add(FoldingDescriptor(node, textRange))
+        if (needFolding(node)) {
+            val textRange = getRangeToFold(node)
+            if (!isOneLine(textRange, document)) {
+                descriptors.add(FoldingDescriptor(node, textRange))
+            }
         }
-        var child: ASTNode? = node.getFirstChildNode()
+
+        var child = node.getFirstChildNode()
         while (child != null) {
             appendDescriptors(child, document, descriptors)
             child = child.getTreeNext()
         }
     }
 
+    private fun needFolding(node: ASTNode): Boolean {
+        val type = node.getElementType()
+        val parentType = node.getTreeParent()?.getElementType()
+        return type == JetNodeTypes.FUNCTION_LITERAL ||
+               (type == JetNodeTypes.BLOCK && parentType != JetNodeTypes.FUNCTION_LITERAL) ||
+               type == JetNodeTypes.CLASS_BODY || type == JetTokens.BLOCK_COMMENT || type == KDocTokens.KDOC
+    }
+
+    private fun getRangeToFold(node: ASTNode): TextRange {
+        if (node.getElementType() == JetNodeTypes.FUNCTION_LITERAL) {
+            val psi = node.getPsi() as? JetFunctionLiteral
+            val lbrace = psi?.getLBrace()
+            val rbrace = psi?.getRBrace()
+            if (lbrace != null && rbrace != null) {
+                return TextRange(lbrace.getTextRange().getStartOffset(), rbrace.getTextRange().getEndOffset())
+            }
+        }
+        return node.getTextRange()
+    }
+
     private fun isOneLine(textRange: TextRange, document: Document) =
         document.getLineNumber(textRange.getStartOffset()) == document.getLineNumber(textRange.getEndOffset())
 
diff --git a/idea/testData/folding/checkCollapse/functionLiteral.kt b/idea/testData/folding/checkCollapse/functionLiteral.kt
new file mode 100644
index 00000000000..0386d9ebf67
--- /dev/null
+++ b/idea/testData/folding/checkCollapse/functionLiteral.kt
@@ -0,0 +1,12 @@
+fun some(body: () -> Unit) = body()
+
+fun oneMore() <fold text='{...}' expand='true'>{
+    some <fold text='{...}' expand='true'>{
+        // this is a comment
+        val v1 = "Body"
+        val v2 = "of"
+        val v3 = "function"
+    }</fold>
+}</fold>
+
+// Generated from: idea/testData/folding/checkCollapse/block.kt
\ 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 7e7071e4c9f..08bb9703742 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/folding/KotlinFoldingTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/folding/KotlinFoldingTestGenerated.java
@@ -88,6 +88,12 @@ public class KotlinFoldingTestGenerated extends AbstractKotlinFoldingTest {
             JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/folding/checkCollapse"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("functionLiteral.kt")
+        public void testFunctionLiteral() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/folding/checkCollapse/functionLiteral.kt");
+            doSettingsFoldingTest(fileName);
+        }
+
         @TestMetadata("headerKDoc.kt")
         public void testHeaderKDoc() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/folding/checkCollapse/headerKDoc.kt");