From f9c4ceda7c99d1cd62738e77b8d16be4b7f410c4 Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Tue, 22 Aug 2017 16:54:32 +0700
Subject: [PATCH] AtomicConversionRule: also support non-void context for
 updateAndGet (IDEA-CR-24025)

---
 .../rules/AtomicConversionRule.java           | 54 +++++++++----------
 .../intentions/atomic/afterInLambdas.java     | 10 +++-
 .../intentions/atomic/beforeInLambdas.java    |  7 +++
 3 files changed, 40 insertions(+), 31 deletions(-)

diff --git a/java/typeMigration/src/com/intellij/refactoring/typeMigration/rules/AtomicConversionRule.java b/java/typeMigration/src/com/intellij/refactoring/typeMigration/rules/AtomicConversionRule.java
index acefde2c937e..940bd7f73819 100644
--- a/java/typeMigration/src/com/intellij/refactoring/typeMigration/rules/AtomicConversionRule.java
+++ b/java/typeMigration/src/com/intellij/refactoring/typeMigration/rules/AtomicConversionRule.java
@@ -209,36 +209,32 @@ public class AtomicConversionRule extends TypeConversionRule {
       boolean voidContext = parent instanceof PsiExpressionStatement ||
                   (parent instanceof PsiLambdaExpression && PsiType.VOID.equals(LambdaUtil.getFunctionalInterfaceReturnType(
                     (PsiFunctionalExpression)parent)));
+      if (operationSign == JavaTokenType.EQ) {
+        if (!voidContext) return null;
+        final PsiExpression lExpression = ((PsiAssignmentExpression)context).getLExpression();
+        if (lExpression instanceof PsiReferenceExpression) {
+          final PsiElement element = ((PsiReferenceExpression)lExpression).resolve();
+          if (element instanceof PsiVariable && ((PsiVariable)element).hasModifierProperty(PsiModifier.FINAL)) {
+            return wrapWithNewExpression(to, from, ((PsiAssignmentExpression)context).getRExpression(), element);
+          }
+        }
+        return new TypeConversionDescriptor("$qualifier$ = $val$", "$qualifier$.set($val$)");
+      }
+      if (PsiUtil.isLanguageLevel8OrHigher(context)) {
+        final String name =
+          JavaCodeStyleManager.getInstance(context.getProject()).suggestUniqueVariableName("v", context, false);
+        return new TypeConversionDescriptor("$qualifier$" + sign + "$val$",
+                                            "$qualifier$.updateAndGet(" +
+                                            name +
+                                            " -> " +
+                                            getBoxedWrapper(from, to, name + " " + sign.charAt(0) + " $val$)"));
+      }
       if (voidContext) {
-        if (operationSign == JavaTokenType.EQ) {
-          final PsiExpression lExpression = ((PsiAssignmentExpression)context).getLExpression();
-          if (lExpression instanceof PsiReferenceExpression) {
-            final PsiElement element = ((PsiReferenceExpression)lExpression).resolve();
-            if (element instanceof PsiVariable && ((PsiVariable)element).hasModifierProperty(PsiModifier.FINAL)) {
-              return wrapWithNewExpression(to, from, ((PsiAssignmentExpression)context).getRExpression(), element);
-            }
-          }
-          return new TypeConversionDescriptor("$qualifier$ = $val$", "$qualifier$.set($val$)");
-        }
-        else {
-          if (PsiUtil.isLanguageLevel8OrHigher(context)) {
-            final String name =
-              JavaCodeStyleManager.getInstance(context.getProject()).suggestUniqueVariableName("v", context, false);
-            return new TypeConversionDescriptor("$qualifier$" + sign + "$val$", "$qualifier$.updateAndGet("
-                                                                                + name + " -> " + getBoxedWrapper(from, to, name + " " + sign.charAt(0) + " $val$)")); }
-          else {
-            if (context.getParent() instanceof PsiStatement) {
-              return new TypeConversionDescriptor("$qualifier$" + sign + "$val$", "$qualifier$.set(" +
-                                                                                  getBoxedWrapper(from, to, "$qualifier$.get() " +
-                                                                                                            sign.charAt(0) +
-                                                                                                            " $val$") +
-                                                                                  ")");
-            } else {
-              return null;
-            }
-          }
-        }
-      } //else should be a conflict
+        return new TypeConversionDescriptor("$qualifier$" + sign + "$val$",
+                                            "$qualifier$.set(" +
+                                            getBoxedWrapper(from, to, "$qualifier$.get() " + sign.charAt(0) + " $val$") +
+                                            ")");
+      }
     }
     else if (context instanceof PsiPostfixExpression) {
       final String sign = ((PsiPostfixExpression)context).getOperationSign().getText();
diff --git a/java/typeMigration/testData/intentions/atomic/afterInLambdas.java b/java/typeMigration/testData/intentions/atomic/afterInLambdas.java
index fc7061887455..1d959dd9c3f0 100644
--- a/java/typeMigration/testData/intentions/atomic/afterInLambdas.java
+++ b/java/typeMigration/testData/intentions/atomic/afterInLambdas.java
@@ -1,6 +1,8 @@
-import java.util.concurrent.atomic.AtomicInteger;
-
 // "Convert to atomic" "true"
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.*;
+
 public class InLambdas
 {
   public void test()
@@ -10,5 +12,9 @@ public class InLambdas
     Runnable r2 = () -> x.addAndGet(2);
     Runnable r3 = () -> x.updateAndGet(v -> v * 2);
     Runnable r4 = () -> x.set(5);
+    System.out.println(x.updateAndGet(v -> v / 3));
+    IntSupplier s = () -> {
+      return x.updateAndGet(v -> v * 2);
+    };
   }
 }
\ No newline at end of file
diff --git a/java/typeMigration/testData/intentions/atomic/beforeInLambdas.java b/java/typeMigration/testData/intentions/atomic/beforeInLambdas.java
index 0a947749547e..a8c9de45e6e0 100644
--- a/java/typeMigration/testData/intentions/atomic/beforeInLambdas.java
+++ b/java/typeMigration/testData/intentions/atomic/beforeInLambdas.java
@@ -1,4 +1,7 @@
 // "Convert to atomic" "true"
+
+import java.util.function.*;
+
 public class InLambdas
 {
   public void test()
@@ -8,5 +11,9 @@ public class InLambdas
     Runnable r2 = () -> x+=2;
     Runnable r3 = () -> x*=2;
     Runnable r4 = () -> x = 5;
+    System.out.println(x /= 3);
+    IntSupplier s = () -> {
+      return x *= 2;
+    };
   }
 }
\ No newline at end of file