better folding for function literals

#KT-4090 Fixed

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()
+        if (needFolding(node)) {
-        val type = node.getElementType()
+            val textRange = getRangeToFold(node)
-        if ((type == JetNodeTypes.BLOCK || type == JetNodeTypes.CLASS_BODY || type == JetTokens.BLOCK_COMMENT || type == KDocTokens.KDOC) &&
+            if (!isOneLine(textRange, document)) {
-                !isOneLine(textRange, document)) {
+                descriptors.add(FoldingDescriptor(node, textRange))
-            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())
 

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

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");