Partial body resolve filter: changed algorithm so that all statements to resolve are found when constructing it

idea/ide-common/src/org/jetbrains/jet/lang/resolve/lazy/PartialBodyResolveFilter.kt

@@ -38,7 +38,6 @@ class PartialBodyResolveFilter(
 ) : (JetElement) -> Boolean {
 
     private val statementsToResolve = HashSet<JetExpression>()
-    private val processedBlocks = HashSet<JetBlockExpression>()
 
     private val nothingFunctionNames = HashSet(probablyNothingCallableNamesService.functionNames())
     private val nothingPropertyNames = probablyNothingCallableNamesService.propertyNames()
@@ -63,13 +62,7 @@ class PartialBodyResolveFilter(
         addStatementsToResolve(elementToResolve)
     }
 
-    override fun invoke(statement: JetElement): Boolean {
+    override fun invoke(statement: JetElement) = statement in statementsToResolve
-        val block = statement.getParent() as JetBlockExpression
-        if (block !in processedBlocks) {
-            processBlock(block)
-        }
-        return statement in statementsToResolve
-    }
 
     tailRecursive
     private fun addStatementsToResolve(element: JetElement) {
@@ -77,11 +70,10 @@ class PartialBodyResolveFilter(
         val parent = element.getParent() as? JetElement ?: return
 
         if (parent is JetBlockExpression) {
-            processBlock(parent)
             if (element in statementsToResolve) return // already processed
 
             if (element is JetExpression) {
-                statementsToResolve.add(element)
+                addStatementToResolve(element)
             }
 
             for (statement in element.siblings(forward = false, withItself = false)) {
@@ -90,11 +82,11 @@ class PartialBodyResolveFilter(
 
                 val smartCastPlaces = potentialSmartCastPlaces(statement)
                 if (!smartCastPlaces.isEmpty()) {
-                    statementsToResolve.add(statement)
+                    addStatementToResolve(statement)
                     statementsToResolve.addStatementsForPlaces(smartCastPlaces.values().flatMap { it })
                 }
                 else if (statement is JetDeclaration) {
-                    statementsToResolve.add(statement)
+                    addStatementToResolve(statement)
                 }
             }
         }
@@ -102,12 +94,24 @@ class PartialBodyResolveFilter(
         addStatementsToResolve(parent)
     }
 
-    private fun processBlock(block: JetBlockExpression) {
+    private fun addStatementToResolve(statement: JetExpression) {
-        if (processedBlocks.add(block)) {
+        assert(statement.getParent() is JetBlockExpression)
-            val lastStatement = block.lastStatement()
+        if (statementsToResolve.add(statement)) {
-            if (lastStatement != null && lastStatement !in statementsToResolve && isValueNeeded(block)) {
+            // search for code blocks inside statement and make sure its values will be resolved
-                addStatementsToResolve(lastStatement)
+            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)
+                }
+            })
         }
     }