From b721b0211a9f23d1494e8fc101bf3018875ecda0 Mon Sep 17 00:00:00 2001
From: Evgeny Gerashchenko <Evgeny.Gerashchenko@jetbrains.com>
Date: Tue, 23 Jul 2013 17:42:16 +0400
Subject: [PATCH] Fixed higlighting of replaced expressions.

 #KT-2637 fixed
---
 .../jetbrains/jet/lang/psi/JetExpressionImpl.java | 13 +++++++------
 .../inline/KotlinInlineLocalHandler.java          | 15 +++++++++++----
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetExpressionImpl.java b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetExpressionImpl.java
index bc1584cd381..c0002d7c9c8 100644
--- a/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetExpressionImpl.java
+++ b/compiler/frontend/src/org/jetbrains/jet/lang/psi/JetExpressionImpl.java
@@ -18,7 +18,7 @@ package org.jetbrains.jet.lang.psi;
 
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.impl.source.tree.TreeElement;
+import com.intellij.util.IncorrectOperationException;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.jet.JetNodeType;
 
@@ -45,11 +45,12 @@ abstract class JetExpressionImpl extends JetElementImpl implements JetExpression
     }
 
     @Override
-    public void replaceChildInternal(PsiElement child, TreeElement newElement) {
-        PsiElement newPsi = newElement.getPsi();
-        if (newPsi instanceof JetExpression && JetPsiUtil.areParenthesesNecessary((JetExpression) newPsi, ((JetExpression) child), this)) {
-            newElement = (TreeElement) JetPsiFactory.createExpression(getProject(), "(" + newElement.getText() + ")").getNode();
+    public PsiElement replace(@NotNull PsiElement newElement) throws IncorrectOperationException {
+        PsiElement parent = getParent();
+        if (parent instanceof JetExpression && newElement instanceof JetExpression &&
+            JetPsiUtil.areParenthesesNecessary((JetExpression) newElement, this, (JetExpression) parent)) {
+            return super.replace(JetPsiFactory.createExpression(getProject(), "(" + newElement.getText() + ")"));
         }
-        super.replaceChildInternal(child, newElement);
+        return super.replace(newElement);
     }
 }
diff --git a/idea/src/org/jetbrains/jet/plugin/refactoring/inline/KotlinInlineLocalHandler.java b/idea/src/org/jetbrains/jet/plugin/refactoring/inline/KotlinInlineLocalHandler.java
index c453a375e95..bc296fe9122 100644
--- a/idea/src/org/jetbrains/jet/plugin/refactoring/inline/KotlinInlineLocalHandler.java
+++ b/idea/src/org/jetbrains/jet/plugin/refactoring/inline/KotlinInlineLocalHandler.java
@@ -67,7 +67,7 @@ public class KotlinInlineLocalHandler extends InlineActionHandler {
     }
 
     @Override
-    public void inlineElement(Project project, Editor editor, PsiElement element) {
+    public void inlineElement(Project project, final Editor editor, PsiElement element) {
         final JetProperty val = (JetProperty) element;
         String name = val.getName();
 
@@ -110,10 +110,12 @@ public class KotlinInlineLocalHandler extends InlineActionHandler {
 
         PsiReference[] referencesArray = references.toArray(references.toArray(new PsiReference[references.size()]));
 
+        EditorColorsManager editorColorsManager = EditorColorsManager.getInstance();
+        final TextAttributes searchResultsAttributes = editorColorsManager.getGlobalScheme().getAttributes(EditorColors.SEARCH_RESULT_ATTRIBUTES);
+        final HighlightManager highlightManager = HighlightManager.getInstance(project);
+
         if (editor != null && !ApplicationManager.getApplication().isUnitTestMode()) {
-            EditorColorsManager editorColorsManager = EditorColorsManager.getInstance();
-            TextAttributes attributes = editorColorsManager.getGlobalScheme().getAttributes(EditorColors.SEARCH_RESULT_ATTRIBUTES);
-            HighlightManager.getInstance(project).addOccurrenceHighlights(editor, referencesArray, attributes, true, null);
+            highlightManager.addOccurrenceHighlights(editor, referencesArray, searchResultsAttributes, true, null);
             if (!showDialog(project, name, references)) {
                 StatusBar statusBar = WindowManager.getInstance().getStatusBar(project);
                 if (statusBar != null) {
@@ -154,6 +156,11 @@ public class KotlinInlineLocalHandler extends InlineActionHandler {
                                 if (parametersForFunctionLiteral != null) {
                                     addFunctionLiteralParameterTypes(parametersForFunctionLiteral, inlinedExpressions);
                                 }
+
+                                if (editor != null && !ApplicationManager.getApplication().isUnitTestMode()) {
+                                    highlightManager.addOccurrenceHighlights(
+                                            editor, inlinedExpressions.toArray(new PsiElement[inlinedExpressions.size()]), searchResultsAttributes, true, null);
+                                }
                             }
                         });
                     }