From 431f3ad18580d3621e90f80fbd40ad1eb98cb464 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <Nikolay.Krasko@jetbrains.com>
Date: Mon, 23 Jun 2014 18:54:50 +0400
Subject: [PATCH] Smart enter in function literal expression

 #KT-3600 In Progress
---
 .../plugin/editor/KotlinSmartEnterHandler.kt  | 19 ++++--
 .../codeInsight/smartEnter/SmartEnterTest.kt  | 58 ++++++++++++++++++-
 2 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt b/idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt
index f3232244132..f2ca2c0eaa6 100644
--- a/idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt
+++ b/idea/src/org/jetbrains/jet/plugin/editor/KotlinSmartEnterHandler.kt
@@ -32,6 +32,7 @@ import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
 import org.jetbrains.jet.lang.psi.JetIfExpression
 import org.jetbrains.jet.lang.psi.JetForExpression
 import org.jetbrains.jet.lang.psi.JetParameter
+import org.jetbrains.jet.lang.psi.JetFunctionLiteral
 import com.intellij.psi.tree.TokenSet
 import org.jetbrains.jet.JetNodeTypes
 import org.jetbrains.jet.lang.psi.JetLoopExpression
@@ -64,12 +65,17 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
         if (atCaret is PsiWhiteSpace) return null
 
         while (atCaret != null) {
-            if (atCaret?.isJetStatement() == true) return atCaret
-
-            if (atCaret is JetDeclaration &&
-                    atCaret !is JetParameter &&
-                    (atCaret?.getParent() !is JetForExpression)) {
-                return atCaret
+            when {
+                atCaret?.isJetStatement() == true -> return atCaret
+                atCaret?.getParent() is JetFunctionLiteral -> return atCaret
+                atCaret is JetDeclaration -> {
+                    val declaration = atCaret!!
+                    when {
+                        declaration is JetParameter && !declaration.isInLambdaExpression() -> {/* proceed to function declaration */}
+                        declaration.getParent() is JetForExpression -> {/* skip variable declaration in 'for' expression */}
+                        else -> return atCaret
+                    }
+                }
             }
 
             atCaret = atCaret?.getParent()
@@ -152,3 +158,4 @@ public class KotlinSmartEnterHandler: SmartEnterProcessorWithFixers() {
 }
 
 private val IF_BRANCHES_CONTAINERS = TokenSet.create(JetNodeTypes.THEN, JetNodeTypes.ELSE)
+private fun JetParameter.isInLambdaExpression() = this.getParent()?.getParent() is JetFunctionLiteral
diff --git a/idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt b/idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt
index f2ed2b6f204..1b0b73feff2 100644
--- a/idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt
+++ b/idea/tests/org/jetbrains/jet/plugin/codeInsight/smartEnter/SmartEnterTest.kt
@@ -553,8 +553,8 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             ,
             """
             when ( {
-            }
                 <caret>
+            }
             """
     )
 
@@ -841,6 +841,62 @@ class SmartEnterTest : JetLightCodeInsightFixtureTestCase() {
             """
     )
 
+    fun testInLambda1() = doFunTest(
+            """
+            some {
+                p -><caret>
+            }
+            """,
+            """
+            some {
+                p ->
+                <caret>
+            }
+            """
+    )
+
+    fun testInLambda2() = doFunTest(
+            """
+            some {
+                p<caret> ->
+            }
+            """,
+            """
+            some {
+                p ->
+                <caret>
+            }
+            """
+    )
+
+    fun testInLambda3() = doFunTest(
+            """
+            some {
+                (<caret>p: Int) : Int ->
+            }
+            """,
+            """
+            some {
+                (p: Int) : Int ->
+                <caret>
+            }
+            """
+    )
+
+    fun testInLambda4() = doFunTest(
+            """
+            some {
+                (p: <caret>Int) : Int ->
+            }
+            """,
+            """
+            some {
+                (p: Int) : Int ->
+                <caret>
+            }
+            """
+    )
+
     fun doFunTest(before: String, after: String) {
         fun String.withFunContext(): String {
             val bodyText = "//----${this.trimIndent()}//----"