VariableExtractor: support ternary conditions in assignment and return (IDEA-202048)

java/java-impl/src/com/intellij/refactoring/introduceVariable/VariableExtractor.java

@@ -244,12 +244,10 @@ class VariableExtractor {
         }
       }
     }
-    if (firstOccurrence != null && ControlFlowUtils.canExtractStatement(firstOccurrence)) {
-      PsiExpression ancestorCandidate = anchor instanceof PsiIfStatement ? ((PsiIfStatement)anchor).getCondition() :
-                                        anchor instanceof PsiReturnStatement ? ((PsiReturnStatement)anchor).getReturnValue():
-                                        anchor instanceof PsiExpression ? (PsiExpression)anchor :
-                                        null;
-      if (PsiTreeUtil.isAncestor(ancestorCandidate, firstOccurrence, false) &&
+    if (firstOccurrence != null && ControlFlowUtils.canExtractStatement(firstOccurrence) && 
+        !PsiUtil.isAccessedForWriting(firstOccurrence)) {
+      PsiExpression ancestorCandidate = ExpressionUtils.getTopLevelExpression(firstOccurrence);
+      if (PsiTreeUtil.isAncestor(anchor, ancestorCandidate, false) &&
           ReorderingUtils.canExtract(ancestorCandidate, firstOccurrence) == ThreeState.NO) {
         return firstOccurrence;
       }

java/java-impl/src/com/intellij/refactoring/util/EnsureCodeBlockImpl.java

@@ -54,6 +54,11 @@ class EnsureCodeBlockImpl {
       if (parent == null) return null;
       return replace(expression, parent, (oldParent, copy) -> extractFieldInitializer((PsiField)oldParent, (PsiField)copy));
     }
+    
+    PsiConditionalExpression ternary = findSurroundingTernary(expression);
+    if (ternary != null && parent instanceof PsiStatement) {
+      return replace(expression, parent, (oldParent, copy) -> replaceTernaryWithIf((PsiStatement)oldParent, ternary));
+    }
 
     PsiPolyadicExpression condition = findSurroundingConditionChain(expression);
     PsiExpression operand;
@@ -183,6 +188,28 @@ class EnsureCodeBlockImpl {
     Objects.requireNonNull(whileStatement.getCondition()).replace(lOperands);
     return newParent;
   }
+  private static PsiElement replaceTernaryWithIf(PsiStatement statement, PsiConditionalExpression ternary) {
+    PsiElementFactory factory = JavaPsiFacade.getElementFactory(statement.getProject());
+    PsiIfStatement ifStatement =
+      (PsiIfStatement)factory.createStatementFromText("if(" + ternary.getCondition().getText() + ") {} else {}", statement);
+    Object mark = new Object();
+    PsiTreeUtil.mark(ternary, mark);
+    PsiStatement thenStatement = (PsiStatement)statement.copy();
+    PsiConditionalExpression thenTernary = Objects.requireNonNull((PsiConditionalExpression)PsiTreeUtil.releaseMark(thenStatement, mark));
+    PsiExpression branch1 = ternary.getThenExpression();
+    if (branch1 != null) {
+      thenTernary.replace(branch1);
+    }
+    PsiStatement elseStatement = (PsiStatement)statement.copy();
+    PsiConditionalExpression elseTernary = Objects.requireNonNull((PsiConditionalExpression)PsiTreeUtil.releaseMark(elseStatement, mark));
+    PsiExpression branch = ternary.getElseExpression();
+    if (branch != null) {
+      elseTernary.replace(branch);
+    }
+    ((PsiBlockStatement)Objects.requireNonNull(ifStatement.getThenBranch())).getCodeBlock().add(thenStatement);
+    ((PsiBlockStatement)Objects.requireNonNull(ifStatement.getElseBranch())).getCodeBlock().add(elseStatement);
+    return statement.replace(ifStatement);
+  }
 
   private static PsiElement splitIf(PsiIfStatement outerIf, PsiPolyadicExpression andChain, PsiExpression operand) {
     PsiExpression lOperands = SplitConditionUtil.getLOperands(andChain, andChain.getTokenBeforeOperand(operand));
@@ -227,6 +254,18 @@ class EnsureCodeBlockImpl {
     return polyadicExpression;
   }
 
+  @Nullable
+  private static PsiConditionalExpression findSurroundingTernary(@NotNull PsiExpression expression) {
+    PsiExpression current = expression;
+    PsiConditionalExpression ternary;
+    do {
+      current = ternary = PsiTreeUtil.getParentOfType(current, PsiConditionalExpression.class, true,
+                                                      PsiStatement.class, PsiLambdaExpression.class);
+    }
+    while (ternary != null && PsiTreeUtil.isAncestor(ternary.getCondition(), expression, false));
+    return ternary;
+  }
+
   private static boolean hasNameCollision(PsiElement declaration, PsiElement context) {
     if (declaration instanceof PsiDeclarationStatement) {
       PsiResolveHelper helper = JavaPsiFacade.getInstance(context.getProject()).getResolveHelper();

java/java-tests/testData/refactoring/introduceVariable/AssignTernary.after.java

@@ -0,0 +1,12 @@
+class Test {
+    String x(String s) {
+        String x;
+        if (s == null) {
+            x = "";
+        } else {
+            String temp = s.trim();
+            x = temp;
+        }
+        return x;
+    }
+}

java/java-tests/testData/refactoring/introduceVariable/AssignTernary.java

@@ -0,0 +1,7 @@
+class Test {
+    String x(String s) {
+        String x;
+        x = s == null ? "" : <selection>s.trim()</selection>;
+        return x;
+    }
+}

java/java-tests/testData/refactoring/introduceVariable/ReturnTernary.after.java

@@ -0,0 +1,10 @@
+class Test {
+    String x(String s) {
+        if (s == null) {
+            return "";
+        } else {
+            String temp = s.trim();
+            return temp;
+        }
+    }
+}

java/java-tests/testData/refactoring/introduceVariable/ReturnTernary.java

@@ -0,0 +1,5 @@
+class Test {
+    String x(String s) {
+        return s == null ? "" : <selection>s.trim()</selection>;
+    }
+}

java/java-tests/testSrc/com/intellij/java/refactoring/IntroduceVariableTest.java

@@ -238,6 +238,14 @@ public class IntroduceVariableTest extends LightCodeInsightTestCase {
   public void testReturnOrChain() {
     doTest(new MockIntroduceVariableHandler("temp", true, false, false, CommonClassNames.JAVA_LANG_STRING));
   }
+  
+  public void testReturnTernary() {
+    doTest(new MockIntroduceVariableHandler("temp", true, false, false, CommonClassNames.JAVA_LANG_STRING));
+  }
+  
+  public void testAssignTernary() {
+    doTest(new MockIntroduceVariableHandler("temp", true, false, false, CommonClassNames.JAVA_LANG_STRING));
+  }
 
   public void testLambdaAndChain() {
     doTest(new MockIntroduceVariableHandler("temp", true, false, false, CommonClassNames.JAVA_LANG_STRING));

plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ControlFlowUtils.java

@@ -928,7 +928,10 @@ public class ControlFlowUtils {
         }
       }
       if (parent instanceof PsiConditionalExpression && ((PsiConditionalExpression)parent).getCondition() != cur) {
-        return false;
+        PsiElement ternaryParent = PsiUtil.skipParenthesizedExprUp(parent.getParent());
+        return ternaryParent instanceof PsiReturnStatement ||
+               (ternaryParent instanceof PsiAssignmentExpression && ternaryParent.getParent() instanceof PsiExpressionStatement &&
+                PsiUtil.skipParenthesizedExprDown(((PsiAssignmentExpression)ternaryParent).getRExpression()) == parent);
       }
       if(parent instanceof PsiMethodCallExpression) {
         PsiReferenceExpression methodExpression = ((PsiMethodCallExpression)parent).getMethodExpression();