From e52b524d9ab3ebef0c1372fd27eb1ab5f31b6fc6 Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Sat, 11 Jul 2015 14:28:41 +0300
Subject: [PATCH] Fixed rename of operator functions

---
 .../intentions/OperatorToFunctionIntention.kt | 28 ++++++++++++++---
 .../idea/references/JetSimpleNameReference.kt | 31 ++++++-------------
 .../kotlin/idea/references/referenceUtil.kt   |  8 ++---
 3 files changed, 36 insertions(+), 31 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/intentions/OperatorToFunctionIntention.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/intentions/OperatorToFunctionIntention.kt
index 5b9f0734b0e..7e932fa33af 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/intentions/OperatorToFunctionIntention.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/intentions/OperatorToFunctionIntention.kt
@@ -22,8 +22,9 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 
 public class OperatorToFunctionIntention : JetSelfTargetingIntention<JetExpression>(javaClass(), "Replace overloaded operator with function call") {
@@ -126,7 +127,7 @@ public class OperatorToFunctionIntention : JetSelfTargetingIntention<JetExpressi
                 return element
             }
 
-            val context = element.analyze()
+            val context = element.analyze(BodyResolveMode.PARTIAL)
             val functionCandidate = element.getResolvedCall(context)
             val functionName = functionCandidate?.getCandidateDescriptor()?.getName().toString()
             val elemType = context.getType(left)
@@ -195,14 +196,31 @@ public class OperatorToFunctionIntention : JetSelfTargetingIntention<JetExpressi
             return callee.getParent()!!.replace(transformed) as JetExpression
         }
 
-        public fun convert(element: JetExpression): JetExpression {
-            return when (element) {
+        public fun convert(element: JetExpression): Pair<JetExpression, JetSimpleNameExpression> {
+            val result = when (element) {
                 is JetPrefixExpression -> convertPrefix(element)
                 is JetPostfixExpression -> convertPostFix(element)
                 is JetBinaryExpression -> convertBinary(element)
                 is JetArrayAccessExpression -> convertArrayAccess(element)
                 is JetCallExpression -> convertCall(element)
-                else -> element
+                else -> throw IllegalArgumentException(element.toString())
+            }
+
+            return result to findCallName(result)
+        }
+
+        private fun findCallName(result: JetExpression): JetSimpleNameExpression {
+            return when (result) {
+                is JetBinaryExpression -> {
+                    if (JetPsiUtil.isAssignment(result))
+                        findCallName(result.getRight()!!)
+                    else
+                        findCallName(result.getLeft()!!)
+                }
+
+                is JetUnaryExpression -> findCallName(result.getBaseExpression()!!)
+
+                else -> result.getQualifiedElementSelector() as JetSimpleNameExpression
             }
         }
     }
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetSimpleNameReference.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetSimpleNameReference.kt
index b4453c21d1f..eaaeb37fba7 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetSimpleNameReference.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/JetSimpleNameReference.kt
@@ -33,10 +33,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.plugin.references.SimpleNameReferenceExtension
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getParentOfTypeAndBranch
-import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElement
-import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
-import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.dataClassUtils.isComponentLike
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
@@ -98,26 +95,16 @@ class JetSimpleNameReference(expression: JetSimpleNameExpression) : JetSimpleRef
             }
         }
 
-        var nameElement = expression.getReferencedNameElement()
+        val nameElement = expression.getReferencedNameElement()
 
-        val elementType = nameElement.getNode()?.getElementType()
-        val opExpression = PsiTreeUtil.getParentOfType<JetExpression>(expression, javaClass<JetUnaryExpression>(), javaClass<JetBinaryExpression>())
-        if (elementType is JetToken && OperatorConventions.getNameForOperationSymbol(elementType) != null && opExpression != null) {
-            val bindingContext = expression.analyze(BodyResolveMode.PARTIAL)
-            val oldDescriptor = bindingContext[BindingContext.REFERENCE_TARGET, expression]
-            val newExpression = OperatorToFunctionIntention.convert(opExpression)
-            newExpression.accept(object : JetTreeVisitorVoid() {
-                override fun visitCallExpression(expression: JetCallExpression) {
-                    val callee = expression.getCalleeExpression() as? JetSimpleNameExpression
-                    if (callee != null && bindingContext[BindingContext.REFERENCE_TARGET, callee] == oldDescriptor) {
-                        nameElement = callee.getReferencedNameElement()
-                    }
-                    else {
-                        super.visitCallExpression(expression)
-                    }
-                }
+        val elementType = nameElement.getNode().getElementType()
+        if (elementType is JetToken && OperatorConventions.getNameForOperationSymbol(elementType) != null) {
+            val opExpression = expression.getParent() as? JetOperationExpression
+            if (opExpression != null) {
+                val (newExpression, newNameElement) = OperatorToFunctionIntention.convert(opExpression)
+                newNameElement.replace(element)
+                return newExpression
             }
-            )
         }
 
         nameElement.replace(element)
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
index d10c7cb6c76..ce5caa01a32 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt
@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.getQualifiedElementSelector
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstance
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
@@ -112,10 +113,9 @@ private fun PsiElement.isConstructorOf(unwrappedCandidate: PsiElement) =
 fun AbstractJetReference<out JetExpression>.renameImplicitConventionalCall(newName: String?): JetExpression {
     if (newName == null) return expression
 
-    val expr = OperatorToFunctionIntention.convert(expression) as JetQualifiedExpression
-    val callee = (expr.getSelectorExpression() as JetCallExpression).getCalleeExpression() as JetSimpleNameExpression
-    val newCallee = callee.mainReference.handleElementRename(newName)
-    return newCallee.getStrictParentOfType<JetQualifiedExpression>() as JetExpression
+    val (newExpression, newNameElement) = OperatorToFunctionIntention.convert(expression)
+    newNameElement.mainReference.handleElementRename(newName)
+    return newExpression
 }
 
 val JetSimpleNameExpression.mainReference: JetSimpleNameReference