From 4503adb7fcaffa4d2e76f2e29fb1de3cf408b5af Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Thu, 20 Nov 2014 13:10:33 +0300
Subject: [PATCH] Partial body resolve filter: more correct lambda resolve for
 return

---
 .../resolve/lazy/PartialBodyResolveFilter.kt  | 37 ++++++++++---------
 .../partialBodyResolve/ReturnLambda.dump      |  3 ++
 .../partialBodyResolve/ReturnLambda.kt        |  9 +++++
 .../PartialBodyResolveTestGenerated.java      |  6 +++
 4 files changed, 37 insertions(+), 18 deletions(-)
 create mode 100644 idea/testData/resolve/partialBodyResolve/ReturnLambda.dump
 create mode 100644 idea/testData/resolve/partialBodyResolve/ReturnLambda.kt

diff --git a/idea/ide-common/src/org/jetbrains/jet/lang/resolve/lazy/PartialBodyResolveFilter.kt b/idea/ide-common/src/org/jetbrains/jet/lang/resolve/lazy/PartialBodyResolveFilter.kt
index 0e71ad18645..22fd17e9284 100644
--- a/idea/ide-common/src/org/jetbrains/jet/lang/resolve/lazy/PartialBodyResolveFilter.kt
+++ b/idea/ide-common/src/org/jetbrains/jet/lang/resolve/lazy/PartialBodyResolveFilter.kt
@@ -83,7 +83,8 @@ class PartialBodyResolveFilter(
                 val smartCastPlaces = potentialSmartCastPlaces(statement)
                 if (!smartCastPlaces.isEmpty()) {
                     addStatementToResolve(statement)
-                    statementsToResolve.addStatementsForPlaces(smartCastPlaces.values().flatMap { it })
+                    //TODO: do we really need correct resolve for ALL smart cast places?
+                    addStatementsToResolveForPlaces(smartCastPlaces.values().flatMap { it })
                 }
                 else if (statement is JetDeclaration) {
                     addStatementToResolve(statement)
@@ -94,25 +95,25 @@ class PartialBodyResolveFilter(
         addStatementsToResolve(parent)
     }
 
-    private fun addStatementToResolve(statement: JetExpression) {
+    private fun addStatementToResolve(statement: JetExpression): Boolean {
         assert(statement.getParent() is JetBlockExpression)
-        if (statementsToResolve.add(statement)) {
-            // search for code blocks inside statement and make sure its values will be resolved
-            statement.accept(object : JetVisitorVoid(){
-                override fun visitBlockExpression(expression: JetBlockExpression) {
-                    if (isValueNeeded(expression)) {
-                        val value = expression.lastStatement()
-                        if (value != null) {
-                            addStatementsToResolve(value)
-                        }
+        if (!statementsToResolve.add(statement)) return false
+        // search for code blocks inside statement and make sure its values will be resolved
+        statement.accept(object : JetVisitorVoid(){
+            override fun visitBlockExpression(expression: JetBlockExpression) {
+                if (isValueNeeded(expression)) {
+                    val value = expression.lastStatement()
+                    if (value != null) {
+                        addStatementsToResolve(value)
                     }
                 }
+            }
 
-                override fun visitJetElement(element: JetElement) {
-                    element.acceptChildren(this)
-                }
-            })
-        }
+            override fun visitJetElement(element: JetElement) {
+                element.acceptChildren(this)
+            }
+        })
+        return true
     }
 
     /**
@@ -351,11 +352,11 @@ class PartialBodyResolveFilter(
 
     private fun JetElement.noControlFlowInside() = this is JetFunction || this is JetClass || this is JetClassBody
 
-    private fun MutableSet<JetExpression>.addStatementsForPlaces(places: Collection<JetExpression>) {
+    private fun addStatementsToResolveForPlaces(places: Collection<JetExpression>) {
         for (place in places) {
             var parent: PsiElement = place
             while (true) {
-                if (parent.isStatement() && !add(parent as JetExpression)) break
+                if (parent.isStatement() && !addStatementToResolve(parent as JetExpression)) break
                 parent = parent.getParent()
             }
         }
diff --git a/idea/testData/resolve/partialBodyResolve/ReturnLambda.dump b/idea/testData/resolve/partialBodyResolve/ReturnLambda.dump
new file mode 100644
index 00000000000..5457feb8002
--- /dev/null
+++ b/idea/testData/resolve/partialBodyResolve/ReturnLambda.dump
@@ -0,0 +1,3 @@
+Resolve target: value-parameter val p: kotlin.String? smart-cast to kotlin.String
+Skipped statements:
+println()
diff --git a/idea/testData/resolve/partialBodyResolve/ReturnLambda.kt b/idea/testData/resolve/partialBodyResolve/ReturnLambda.kt
new file mode 100644
index 00000000000..7783bb1a6eb
--- /dev/null
+++ b/idea/testData/resolve/partialBodyResolve/ReturnLambda.kt
@@ -0,0 +1,9 @@
+fun foo(p: String?): () -> String {
+    if (p == null) {
+        return {
+            println()
+            "a"
+        }
+    }
+    <caret>p.size
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/jet/resolve/PartialBodyResolveTestGenerated.java b/idea/tests/org/jetbrains/jet/resolve/PartialBodyResolveTestGenerated.java
index 6e181d9f1b9..25a3b2424dd 100644
--- a/idea/tests/org/jetbrains/jet/resolve/PartialBodyResolveTestGenerated.java
+++ b/idea/tests/org/jetbrains/jet/resolve/PartialBodyResolveTestGenerated.java
@@ -264,6 +264,12 @@ public class PartialBodyResolveTestGenerated extends AbstractPartialBodyResolveT
         doTest(fileName);
     }
 
+    @TestMetadata("ReturnLambda.kt")
+    public void testReturnLambda() throws Exception {
+        String fileName = JetTestUtils.navigationMetadata("idea/testData/resolve/partialBodyResolve/ReturnLambda.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("Simple.kt")
     public void testSimple() throws Exception {
         String fileName = JetTestUtils.navigationMetadata("idea/testData/resolve/partialBodyResolve/Simple.kt");