Fixed optimize imports for operators

compiler/frontend/src/org/jetbrains/kotlin/psi/KtOperationReferenceExpression.kt

@@ -23,16 +23,12 @@ import org.jetbrains.kotlin.lexer.KtSingleValueToken
 import org.jetbrains.kotlin.lexer.KtToken
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.parsing.KotlinExpressionParsing
+import org.jetbrains.kotlin.resolve.constants.evaluate.binaryOperations
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 
 class KtOperationReferenceExpression(node: ASTNode) : KtSimpleNameExpressionImpl(node) {
     override fun getReferencedNameElement() = findChildByType<PsiElement?>(KotlinExpressionParsing.ALL_OPERATIONS) ?: this
 
-    fun getNameForConventionalOperation(unaryOperations: Boolean = true, binaryOperations: Boolean = true): Name? {
-        val operator = (firstChild as? TreeElement)?.elementType as? KtToken ?: return null
-        return OperatorConventions.getNameForOperationSymbol(operator, unaryOperations, binaryOperations)
-    }
-
-    fun isPredefinedOperator() = (firstChild as? TreeElement)?.elementType is KtSingleValueToken
-
+    val operationSignTokenType: KtSingleValueToken?
+        get() = (firstChild as? TreeElement)?.elementType as? KtSingleValueToken
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallResolverUtil.kt

@@ -131,13 +131,13 @@ fun isConventionCall(call: Call): Boolean {
     val callElement = call.callElement
     if (callElement is KtArrayAccessExpression || callElement is KtDestructuringDeclarationEntry) return true
     val calleeExpression = call.calleeExpression as? KtOperationReferenceExpression ?: return false
-    return calleeExpression.getNameForConventionalOperation() != null
+    return calleeExpression.operationSignTokenType != null
 }
 
 fun isInfixCall(call: Call): Boolean {
     val operationRefExpression = call.calleeExpression as? KtOperationReferenceExpression ?: return false
     val binaryExpression = operationRefExpression.parent as? KtBinaryExpression ?: return false
-    return binaryExpression.operationReference === operationRefExpression && !operationRefExpression.isPredefinedOperator()
+    return binaryExpression.operationReference === operationRefExpression && operationRefExpression.operationSignTokenType == null
 }
 
 fun isInvokeCallOnVariable(call: Call): Boolean {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/OperatorCallChecker.kt

@@ -57,7 +57,7 @@ class OperatorCallChecker : CallChecker {
             return
         }
 
-        val isConventionOperator = element is KtOperationReferenceExpression && element.getNameForConventionalOperation() != null
+        val isConventionOperator = element is KtOperationReferenceExpression && element.operationSignTokenType != null
         if (isConventionOperator || element is KtArrayAccessExpression) {
             if (!functionDescriptor.isOperator) {
                 report(reportOn, functionDescriptor, context.trace)

idea/ide-common/src/org/jetbrains/kotlin/idea/util/OptimizedImportsBuilder.kt

@@ -70,23 +70,36 @@ class OptimizedImportsBuilder(
 
     private val importInsertHelper = ImportInsertHelper.getInstance(file.project)
 
-    private val lockedImports = HashSet<ImportPath>()
+    private sealed class LockedImport {
+        // force presence of this import
+        data class Positive(val importPath: ImportPath) : LockedImport() {
+            override fun toString() = importPath.toString()
+        }
+
+        // force absence of this import
+        data class Negative(val importPath: ImportPath) : LockedImport() {
+            override fun toString() = "-" + importPath.toString()
+        }
+    }
+
+    private val lockedImports = HashSet<LockedImport>()
 
     fun buildOptimizedImports(): List<ImportPath>? {
         // TODO: should we drop unused aliases?
         // keep all non-trivial aliases
         file.importDirectives
                 .mapNotNull { it.importPath }
-                .filterTo(lockedImports) {
+                .filter {
                     val aliasName = it.alias
                     aliasName != null && aliasName != it.fqnPart().shortName()
                 }
+                .mapTo(lockedImports) { LockedImport.Positive(it) }
 
         while (true) {
             val lockedImportsBefore = lockedImports.size
             val result = tryBuildOptimizedImports()
             if (lockedImports.size == lockedImportsBefore) return result
-            testLog?.append("Trying to build import list again with locked imports: ${lockedImports.joinToString { it.pathStr }}\n")
+            testLog?.append("Trying to build import list again with locked imports: ${lockedImports.joinToString()}\n")
         }
     }
 
@@ -95,6 +108,7 @@ class OptimizedImportsBuilder(
         return when {
             parent is KtQualifiedExpression && element == parent.selectorExpression -> parent
             parent is KtCallExpression && element == parent.calleeExpression -> getExpressionToAnalyze(parent)
+            parent is KtOperationExpression && element == parent.operationReference -> parent
             parent is KtUserType -> null //TODO: is it always correct?
             else -> element as? KtExpression //TODO: what if not expression? Example: KtPropertyDelegationMethodsReference
         }
@@ -102,17 +116,21 @@ class OptimizedImportsBuilder(
 
     private fun tryBuildOptimizedImports(): List<ImportPath>? {
         val importsToGenerate = HashSet<ImportPath>()
-        importsToGenerate.addAll(lockedImports)
+        lockedImports
+                .filterIsInstance<LockedImport.Positive>()
+                .mapTo(importsToGenerate) { it.importPath }
 
         val descriptorsByParentFqName = HashMap<FqName, MutableSet<DeclarationDescriptor>>()
         for (descriptor in data.descriptorsToImport) {
             val fqName = descriptor.importableFqName!!
 
             val explicitImportPath = ImportPath(fqName, false)
-            if (explicitImportPath in lockedImports) continue
+            if (explicitImportPath in importsToGenerate) continue
 
-            if (canUseStarImport(descriptor, fqName)) {
-                descriptorsByParentFqName.getOrPut(fqName.parent()) { HashSet() }.add(descriptor)
+            val parentFqName = fqName.parent()
+            val starImportPath = ImportPath(parentFqName, true)
+            if (canUseStarImport(descriptor, fqName) && !starImportPath.isNegativeLocked()) {
+                descriptorsByParentFqName.getOrPut(parentFqName) { HashSet() }.add(descriptor)
             }
             else {
                 importsToGenerate.add(explicitImportPath)
@@ -123,12 +141,13 @@ class OptimizedImportsBuilder(
 
         for (parentFqName in descriptorsByParentFqName.keys) {
             val starImportPath = ImportPath(parentFqName, true)
-            if (starImportPath in lockedImports) continue
+            if (starImportPath in importsToGenerate) continue
 
             val descriptors = descriptorsByParentFqName[parentFqName]!!
             val fqNames = descriptors.map { it.importableFqName!! }.toSet()
             val nameCountToUseStar = descriptors.first().nameCountToUseStar()
             val useExplicitImports = fqNames.size < nameCountToUseStar && !options.isInPackagesToUseStarImport(parentFqName)
+                                     || starImportPath.isNegativeLocked()
             if (useExplicitImports) {
                 fqNames
                         .filter { !isImportedByDefault(it) }
@@ -190,11 +209,15 @@ class OptimizedImportsBuilder(
         val fqName = descriptor.importableFqName ?: return
         val explicitImportPath = ImportPath(fqName, false)
         val starImportPath = ImportPath(fqName.parent(), true)
-        if (file.importDirectives.any { it.importPath == explicitImportPath }) {
-            lockedImports.add(explicitImportPath)
+        val importPaths = file.importDirectives.map { it.importPath }
+        if (explicitImportPath in importPaths) {
+            lockedImports.add(LockedImport.Positive(explicitImportPath))
         }
-        else if (file.importDirectives.any { it.importPath == starImportPath }) {
-            lockedImports.add(starImportPath)
+        else if (starImportPath in importPaths) {
+            lockedImports.add(LockedImport.Positive(starImportPath))
+        }
+        else { // there is no import for this descriptor in the original import list, so do not allow to import it by star-import
+            lockedImports.add(LockedImport.Negative(starImportPath))
         }
     }
 
@@ -298,4 +321,6 @@ class OptimizedImportsBuilder(
         return descriptors1.size == descriptors2.size &&
                descriptors1.zip(descriptors2).all { it.first.importableFqName == it.second.importableFqName } //TODO: can have different order?
     }
+
+    private fun ImportPath.isNegativeLocked(): Boolean = lockedImports.any { it is LockedImport.Negative && it.importPath == this }
 }

idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/KtSimpleNameReference.kt

@@ -206,5 +206,27 @@ class KtSimpleNameReference(expression: KtSimpleNameExpression) : KtSimpleRefere
     override fun getCanonicalText(): String = expression.text
 
     override val resolvesByNames: Collection<String>
-        get() = listOf(element.getReferencedName())
+        get() {
+            val element = element
+
+            if (element is KtOperationReferenceExpression) {
+                val tokenType = element.operationSignTokenType
+                if (tokenType != null) {
+                    val name = OperatorConventions.getNameForOperationSymbol(
+                            tokenType, element.parent is KtUnaryExpression, element.parent is KtBinaryExpression)?.asString()
+                    if (name != null) { // can it be null ever?
+                        val counterpart = OperatorConventions.ASSIGNMENT_OPERATION_COUNTERPARTS[tokenType]
+                        if (counterpart != null) {
+                            val counterpartName = OperatorConventions.getNameForOperationSymbol(counterpart, false, true)!!.asString()
+                            return listOf(name, counterpartName)
+                        }
+                        else {
+                            return listOf(name)
+                        }
+                    }
+                }
+            }
+
+            return listOf(element.getReferencedName())
+        }
 }

idea/testData/editor/optimizeImports/Operators.dependency.kt

@@ -0,0 +1,7 @@
+package p1
+
+operator fun Runnable.plus(p: Int): Runnable = TODO()
+operator fun Runnable.minus(p: Int): Runnable = TODO()
+operator fun Runnable.timesAssign(p: Int): Unit = TODO()
+operator fun Runnable.div(p: Int): Runnable = TODO()
+operator fun Runnable.unaryMinus(): Runnable = TODO()

idea/testData/editor/optimizeImports/Operators.kt

@@ -0,0 +1,9 @@
+// NAME_COUNT_TO_USE_STAR_IMPORT: 5
+import p1.*
+
+fun f(runnable: Runnable) {
+    var r = -runnable
+    print(r + 1)
+    r /= 2
+    r *= 3
+}

idea/testData/editor/optimizeImports/Operators.kt.after

@@ -0,0 +1,12 @@
+// NAME_COUNT_TO_USE_STAR_IMPORT: 5
+import p1.div
+import p1.plus
+import p1.timesAssign
+import p1.unaryMinus
+
+fun f(runnable: Runnable) {
+    var r = -runnable
+    print(r + 1)
+    r /= 2
+    r *= 3
+}

idea/testData/editor/optimizeImports/PlusAndPlusAssign.dependency1.kt

@@ -0,0 +1,3 @@
+package p1
+
+operator fun Runnable.plus(p: Int): Runnable = TODO()

idea/testData/editor/optimizeImports/PlusAndPlusAssign.dependency2.kt

@@ -0,0 +1,7 @@
+package p2
+
+operator fun Runnable.plusAssign(p: Int): Unit = TODO()
+
+val v1 = 1
+val v2 = 1
+val v3 = 1

idea/testData/editor/optimizeImports/PlusAndPlusAssign.kt

@@ -0,0 +1,11 @@
+// NAME_COUNT_TO_USE_STAR_IMPORT: 2
+import p1.plus
+import p2.v1
+import p2.v2
+import p2.v3
+
+fun f(runnable: Runnable) {
+    var r = runnable
+    r += 1
+    v1 + v2 + v3
+}

idea/testData/editor/optimizeImports/PlusAndPlusAssign.kt.after

@@ -0,0 +1,11 @@
+// NAME_COUNT_TO_USE_STAR_IMPORT: 2
+import p1.plus
+import p2.v1
+import p2.v2
+import p2.v3
+
+fun f(runnable: Runnable) {
+    var r = runnable
+    r += 1
+    v1 + v2 + v3
+}

idea/testData/editor/optimizeImports/PlusAndPlusAssign.kt.log

@@ -0,0 +1,3 @@
+Additional checking of reference KtSimpleNameReference: +=
+Changed resolve of KtSimpleNameReference: +=
+Trying to build import list again with locked imports: -p2.*, p1.plus

idea/tests/org/jetbrains/kotlin/idea/imports/OptimizeImportsTestGenerated.java

@@ -227,12 +227,24 @@ public class OptimizeImportsTestGenerated extends AbstractOptimizeImportsTest {
         doTest(fileName);
     }
 
+    @TestMetadata("Operators.kt")
+    public void testOperators() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/Operators.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("Overloads.kt")
     public void testOverloads() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/Overloads.kt");
         doTest(fileName);
     }
 
+    @TestMetadata("PlusAndPlusAssign.kt")
+    public void testPlusAndPlusAssign() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/PlusAndPlusAssign.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("RemoveImportsIfGeneral.kt")
     public void testRemoveImportsIfGeneral() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/editor/optimizeImports/RemoveImportsIfGeneral.kt");