Fix caret positioning in unfolded assignments

idea/src/org/jetbrains/jet/plugin/codeInsight/codeTransformations/branchedTransformations/BranchedUnfoldingUtils.java

@@ -18,6 +18,7 @@ package org.jetbrains.jet.plugin.codeInsight.codeTransformations.branchedTransfo
 
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.psi.*;
@@ -79,7 +80,9 @@ public class BranchedUnfoldingUtils {
         thenExpr.replace(JetPsiFactory.createBinaryExpression(project, lhs, op, thenExpr));
         elseExpr.replace(JetPsiFactory.createBinaryExpression(project, lhs, op, elseExpr));
 
-        assignment.replace(newIfExpression);
+        PsiElement resultElement = assignment.replace(newIfExpression);
+
+        editor.getCaretModel().moveToOffset(resultElement.getTextOffset());
     }
 
     public static void unfoldAssignmentToWhen(@NotNull JetBinaryExpression assignment, @NotNull Editor editor) {
@@ -102,7 +105,9 @@ public class BranchedUnfoldingUtils {
             currExpr.replace(JetPsiFactory.createBinaryExpression(project, lhs, op, currExpr));
         }
 
-        assignment.replace(newWhenExpression);
+        PsiElement resultElement = assignment.replace(newWhenExpression);
+
+        editor.getCaretModel().moveToOffset(resultElement.getTextOffset());
     }
 
     public static void unfoldReturnToIf(@NotNull JetReturnExpression returnExpression) {

idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToIf/simpleIfWithComplexAssignmentLHS.kt

@@ -0,0 +1,7 @@
+fun test(n: Int): Array<String> {
+    var x: Array<String> = Array<String>(1, {""})
+
+    x[0] = <caret>if (n > 5) "A" else "B"
+
+    return x
+}

idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToIf/simpleIfWithComplexAssignmentLHS.kt.after

@@ -0,0 +1,7 @@
+fun test(n: Int): Array<String> {
+    var x: Array<String> = Array<String>(1, {""})
+
+    <caret>if (n > 5) x[0] = "A" else x[0] = "B"
+
+    return x
+}

idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToWhen/simpleWhenWithComplexAssignmentLHS.kt

@@ -0,0 +1,12 @@
+fun test(n: Int): Array<String> {
+    var x: Array<String> = Array<String>(1, {""})
+
+    x[0] = <caret>when(n) {
+        in 0..10 -> "small"
+        in 10..100 -> "average"
+        in 100..1000 -> "big"
+        else -> "unknown"
+    }
+
+    return x
+}

idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToWhen/simpleWhenWithComplexAssignmentLHS.kt.after

@@ -0,0 +1,12 @@
+fun test(n: Int): Array<String> {
+    var x: Array<String> = Array<String>(1, {""})
+
+    <caret>when(n) {
+        in 0..10 -> x[0] = "small"
+        in 10..100 -> x[0] = "average"
+        in 100..1000 -> x[0] = "big"
+        else -> x[0] = "unknown"
+    }
+
+    return x
+}

idea/tests/org/jetbrains/jet/plugin/codeInsight/codeTransformations/CodeTransformationsTestGenerated.java

@@ -223,6 +223,11 @@ public class CodeTransformationsTestGenerated extends AbstractCodeTransformation
             doTestUnfoldAssignmentToIf("idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToIf/simpleIfWithBlocks.kt");
         }
         
+        @TestMetadata("simpleIfWithComplexAssignmentLHS.kt")
+        public void testSimpleIfWithComplexAssignmentLHS() throws Exception {
+            doTestUnfoldAssignmentToIf("idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToIf/simpleIfWithComplexAssignmentLHS.kt");
+        }
+        
         @TestMetadata("simpleIfWithoutAssignment.kt")
         public void testSimpleIfWithoutAssignment() throws Exception {
             doTestUnfoldAssignmentToIf("idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToIf/simpleIfWithoutAssignment.kt");
@@ -251,8 +256,13 @@ public class CodeTransformationsTestGenerated extends AbstractCodeTransformation
             doTestUnfoldAssignmentToWhen("idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToWhen/simpleWhenWithBlocks.kt");
         }
         
+        @TestMetadata("simpleWhenWithComplexAssignmentLHS.kt")
+        public void testSimpleWhenWithComplexAssignmentLHS() throws Exception {
+            doTestUnfoldAssignmentToWhen("idea/testData/codeInsight/codeTransformations/branched/unfolding/assignmentToWhen/simpleWhenWithComplexAssignmentLHS.kt");
+        }
+        
     }
-
+    
     @TestMetadata("idea/testData/codeInsight/codeTransformations/branched/unfolding/returnToIf")
     public static class ReturnToIf extends AbstractCodeTransformationTest {
         public void testAllFilesPresentInReturnToIf() throws Exception {