Introduce Variable: Proper occurrence search in generated block expression

#KT-7836 Fixed
 #KT-8443 Fixed
 #EA-71411 Fixed

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.java

@@ -24,10 +24,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.Pass;
 import com.intellij.openapi.util.Ref;
-import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.*;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.refactoring.HelpID;
 import com.intellij.refactoring.introduce.inplace.OccurrencesChooser;
@@ -67,10 +64,7 @@ import org.jetbrains.kotlin.resolve.scopes.JetScope;
 import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.checker.JetTypeChecker;
 
-import java.util.ArrayList;
+import java.util.*;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
 
 import static org.jetbrains.kotlin.psi.PsiPackage.JetPsiFactory;
 
@@ -253,6 +247,8 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
         }
     }
 
+    private static Key<Boolean> OCCURRENCE = Key.create("OCCURRENCE");
+
     private static Runnable introduceVariable(
             final JetExpression expression,
             final String nameSuggestion,
@@ -342,7 +338,8 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
 
                     if (replaceOccurrence && commonContainer != null) {
                         for (JetExpression replace : allReplaces) {
-                            JetExpression exprAfterReplace = replaceExpression(replace);
+                            JetExpression exprAfterReplace = replaceExpression(replace, false);
+                            exprAfterReplace.putCopyableUserData(OCCURRENCE, true);
                             if (anchor == replace) {
                                 anchor = exprAfterReplace;
                             }
@@ -384,6 +381,18 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                         if (elem != null) {
                             reference.set((JetExpression)elem);
                         }
+
+                        emptyBody.accept(
+                                new JetTreeVisitorVoid() {
+                                    @Override
+                                    public void visitSimpleNameExpression(@NotNull JetSimpleNameExpression expression) {
+                                        if (expression.getCopyableUserData(OCCURRENCE) == null) return;
+
+                                        expression.putCopyableUserData(OCCURRENCE, null);
+                                        references.add(expression);
+                                    }
+                                }
+                        );
                     }
                     else {
                         PsiElement parent = anchor.getParent();
@@ -423,7 +432,7 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                         JetExpression replace = allReplaces.get(i);
 
                         if (i != 0 ? replaceOccurrence : shouldReplaceOccurrence(replace, bindingContext, commonContainer)) {
-                            replaceExpression(replace);
+                            replaceExpression(replace, true);
                         }
                         else {
                             PsiElement sibling = PsiTreeUtil.skipSiblingsBackward(replace, PsiWhiteSpace.class);
@@ -450,7 +459,7 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                 return elem;
             }
 
-            private JetExpression replaceExpression(JetExpression replace) {
+            private JetExpression replaceExpression(JetExpression replace, boolean addToReferences) {
                 boolean isActualExpression = expression == replace;
 
                 JetExpression replacement = psiFactory.createExpression(nameSuggestion);
@@ -474,7 +483,9 @@ public class KotlinIntroduceVariableHandler extends KotlinIntroduceHandlerBase {
                     result = newEntry.getExpression();
                 }
 
-                references.add(result);
+                if (addToReferences) {
+                    references.add(result);
+                }
                 if (isActualExpression) reference.set(result);
 
                 return result;

idea/testData/refactoring/introduceVariable/IntroduceAndCreateBlock.kt

@@ -0,0 +1,3 @@
+fun test(n: Int) {
+    val t = if (n > 0) <selection>n + 1</selection> else n
+}

idea/testData/refactoring/introduceVariable/IntroduceAndCreateBlock.kt.after

@@ -0,0 +1,6 @@
+fun test(n: Int) {
+    val t = if (n > 0) {
+        val i = n + 1
+        i
+    } else n
+}

idea/testData/refactoring/introduceVariable/IntroduceLambdaAndCreateBlock.kt

@@ -0,0 +1,5 @@
+// WITH_RUNTIME
+
+fun foo(p: Int, list: List<Int>) {
+    if (p > 0) list.filter <selection>{ it > 0 }</selection>
+}

idea/testData/refactoring/introduceVariable/IntroduceLambdaAndCreateBlock.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+
+fun foo(p: Int, list: List<Int>) {
+    if (p > 0) {
+        val function: (Int) -> Boolean = { it > 0 }
+        list.filter(function)
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/introduce/JetExtractionTestGenerated.java

@@ -139,6 +139,18 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
             doIntroduceVariableTest(fileName);
         }
 
+        @TestMetadata("IntroduceAndCreateBlock.kt")
+        public void testIntroduceAndCreateBlock() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/introduceVariable/IntroduceAndCreateBlock.kt");
+            doIntroduceVariableTest(fileName);
+        }
+
+        @TestMetadata("IntroduceLambdaAndCreateBlock.kt")
+        public void testIntroduceLambdaAndCreateBlock() throws Exception {
+            String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/introduceVariable/IntroduceLambdaAndCreateBlock.kt");
+            doIntroduceVariableTest(fileName);
+        }
+
         @TestMetadata("It.kt")
         public void testIt() throws Exception {
             String fileName = JetTestUtils.navigationMetadata("idea/testData/refactoring/introduceVariable/It.kt");