Code cleanup to work more reliably with nested problems

idea/src/org/jetbrains/kotlin/idea/inspections/KotlinCleanupInspection.kt

@@ -34,9 +34,11 @@ import org.jetbrains.kotlin.idea.quickfix.ReplaceObsoleteLabelSyntaxFix
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.psi.JetAnnotationEntry
 import org.jetbrains.kotlin.psi.JetFile
-import org.jetbrains.kotlin.psi.JetTreeVisitorVoid
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.ErrorsJvm
+import kotlin.properties.Delegates
 
 public class KotlinCleanupInspection(): LocalInspectionTool(), CleanupLocalInspectionTool {
     // required to simplify the inspection registration in tests
@@ -60,17 +62,27 @@ public class KotlinCleanupInspection(): LocalInspectionTool(), CleanupLocalInspe
         }
 
         val diagnostics = analysisResult.bindingContext.getDiagnostics()
-        val problemDescriptors = arrayListOf<ProblemDescriptor>()
-        file.acceptChildren(object: JetTreeVisitorVoid() {
-            override fun visitElement(element: PsiElement) {
-                super.visitElement(element)
-                diagnostics.forElement(element)
-                   .filter { it.isCleanup() }
-                   .map { it.toProblemDescriptor(file, manager) }
-                   .filterNotNullTo(problemDescriptors)
+
+        class ProblemData(val problemDescriptor: ProblemDescriptor, elementToBeInvalidated: PsiElement) {
+            val depth = elementToBeInvalidated.parents(withItself = true).takeWhile { it !is PsiFile }.count()
+        }
+
+        val problems = arrayListOf<ProblemData>()
+        file.forEachDescendantOfType<PsiElement> { element ->
+            for (diagnostic in diagnostics.forElement(element)) {
+                if (diagnostic.isCleanup()) {
+                    val fixes = diagnostic.toCleanupFixes()
+                    if (fixes.isNotEmpty()) {
+                        val problemDescriptor = diagnostic.toProblemDescriptor(fixes, file, manager)
+                        //TODO: not quite correct for multiple
+                        val elementToBeInvalidated = fixes.map { it.elementToBeInvalidated() }.filterNotNull().firstOrNull() ?: element
+                        problems.add(ProblemData(problemDescriptor, elementToBeInvalidated))
+                    }
+                }
             }
-        })
-        return problemDescriptors.toTypedArray()
+        }
+
+        return problems.sortBy { it.depth }.map { it.problemDescriptor }.toTypedArray()
     }
 
     private fun checkJavaFile(file: PsiJavaFile, manager: InspectionManager): Array<out ProblemDescriptor>? {
@@ -102,6 +114,10 @@ public class KotlinCleanupInspection(): LocalInspectionTool(), CleanupLocalInspe
         return ReplaceObsoleteLabelSyntaxFix.looksLikeObsoleteLabel(annotationEntry)
     }
 
+    private fun Diagnostic.toCleanupFixes(): Collection<CleanupFix> {
+        return JetPsiChecker.createQuickfixes(this).filterIsInstance<CleanupFix>()
+    }
+
     private class Wrapper(val intention: IntentionAction, file: JetFile) : IntentionWrapper(intention, file) {
         override fun invoke(project: Project, editor: Editor?, file: PsiFile?) {
             if (intention.isAvailable(project, editor, file)) { // we should check isAvailable here because some elements may get invalidated (or other conditions may change)
@@ -110,17 +126,11 @@ public class KotlinCleanupInspection(): LocalInspectionTool(), CleanupLocalInspe
         }
     }
 
-    private fun Diagnostic.toProblemDescriptor(file: JetFile, manager: InspectionManager): ProblemDescriptor? {
-        val quickFixes = JetPsiChecker.createQuickfixes(this)
-                .filter { it is CleanupFix }
-                .map { Wrapper(it, file) }
-
-        if (quickFixes.isEmpty()) return null
-
+    private fun Diagnostic.toProblemDescriptor(fixes: Collection<CleanupFix>, file: JetFile, manager: InspectionManager): ProblemDescriptor {
         return manager.createProblemDescriptor(getPsiElement(),
                                                 DefaultErrorMessages.render(this),
                                                 false,
-                                                quickFixes.toTypedArray(),
+                                                fixes.map { Wrapper(it, file) }.toTypedArray(),
                                                 ProblemHighlightType.GENERIC_ERROR_OR_WARNING)
     }
 }

idea/src/org/jetbrains/kotlin/idea/quickfix/CleanupFix.kt

@@ -16,10 +16,15 @@
 
 package org.jetbrains.kotlin.idea.quickfix
 
+import com.intellij.codeInsight.intention.IntentionAction
+import com.intellij.psi.PsiElement
+
 /**
  * Marker interface for quickfixes that can be used as part of the "Cleanup Code" action. The diagnostics
  * that produce these quickfixes need to be added to KotlinCleanupInspection.cleanupDiagnosticsFactories.
  */
-public interface CleanupFix
+public interface CleanupFix : IntentionAction {
+    public fun elementToBeInvalidated(): PsiElement? = null
+}
 // TODO(yole): add isSafeToApply() method here to get rid of filtering by diagnostics factories in
 // KotlinCleanupInspection

idea/src/org/jetbrains/kotlin/idea/quickfix/DeprecatedSymbolUsageFix.kt

@@ -20,10 +20,11 @@ import com.intellij.codeInsight.intention.HighPriorityAction
 import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
-import org.jetbrains.kotlin.psi.JetSimpleNameExpression
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getCalleeExpressionIfAny
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
@@ -50,6 +51,19 @@ public class DeprecatedSymbolUsageFix(
         editor?.moveCaret(offset)
     }
 
+    override fun elementToBeInvalidated(): PsiElement? {
+        val parent = element.getParent()
+        return when (parent) {
+            is JetCallExpression -> {
+                val qualified = parent.getParent() as? JetQualifiedExpression
+                if (parent == qualified?.getSelectorExpression()) qualified else parent
+            }
+            is JetQualifiedExpression -> if (element == parent.getSelectorExpression()) parent else element
+            is JetOperationExpression -> if (element == parent.getOperationReference()) parent else element
+            else -> element
+        }
+    }
+
     companion object : JetSingleIntentionActionFactory() {
         override fun createAction(diagnostic: Diagnostic): IntentionAction? {
             val nameExpression = diagnostic.getPsiElement() as? JetSimpleNameExpression ?: return null

idea/src/org/jetbrains/kotlin/idea/quickfix/RemoveRightPartOfBinaryExpressionFix.java

@@ -21,6 +21,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.diagnostics.Diagnostic;
 import org.jetbrains.kotlin.idea.JetBundle;
 import org.jetbrains.kotlin.idea.core.quickfix.QuickFixUtil;
@@ -90,5 +91,12 @@ public class RemoveRightPartOfBinaryExpressionFix<T extends JetExpression> exten
             }
         };
     }
+
+    // TODO: drop it when converted to Kotlin
+    @Nullable
+    @Override
+    public PsiElement elementToBeInvalidated() {
+        return null;
+    }
 }
 

idea/testData/inspections/cleanup/cleanup.kt.after

@@ -28,7 +28,7 @@ val x = fun(x: String) { }
 fun foo() {
     loop@
     for (i in 1..100) {
-        val v = oldFun2(i as Int)
+        val v = bar(i + 2)
         /* comment */
         continue@loop
     }