ExpressionUtils.getAssignmentTo

java/java-impl/src/com/intellij/codeInspection/streamMigration/ReplaceWithFindFirstFix.java

@@ -86,15 +86,12 @@ class ReplaceWithFindFirstFix extends MigrateToStreamFix {
           return replaceInitializer(loopStatement, var, initializer, replacementText, status);
         }
       }
-      PsiAssignmentExpression previousAssignment =
-        ExpressionUtils.getAssignment(PsiTreeUtil.skipSiblingsBackward(loopStatement, PsiWhiteSpace.class, PsiComment.class));
-      if(previousAssignment != null) {
-        PsiExpression prevRValue = previousAssignment.getRExpression();
-        if(prevRValue != null && ExpressionUtils.isReferenceTo(previousAssignment.getLExpression(), var)) {
-          previousAssignment.delete();
-          return loopStatement.replace(elementFactory.createStatementFromText(
-            var.getName() + " = " + generateOptionalUnwrap(stream, tb, value, prevRValue, var.getType()) + ";", loopStatement));
-        }
+      PsiElement maybeAssignment = PsiTreeUtil.skipSiblingsBackward(loopStatement, PsiWhiteSpace.class, PsiComment.class);
+      PsiExpression prevRValue = ExpressionUtils.getAssignmentTo(maybeAssignment, var);
+      if(prevRValue != null) {
+        maybeAssignment.delete();
+        return loopStatement.replace(elementFactory.createStatementFromText(
+          var.getName() + " = " + generateOptionalUnwrap(stream, tb, value, prevRValue, var.getType()) + ";", loopStatement));
       }
       return loopStatement.replace(elementFactory.createStatementFromText(
         var.getName() + " = " + generateOptionalUnwrap(stream, tb, value, lValue, var.getType()) + ";", loopStatement));

java/java-impl/src/com/intellij/codeInspection/streamMigration/StreamApiMigrationInspection.java

@@ -1438,14 +1438,11 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo
             }
           }
         }
-        PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(first);
-        if(assignment != null) {
-          PsiExpression rValue = assignment.getRExpression();
-          if(rValue != null && ExpressionUtils.isReferenceTo(assignment.getLExpression(), myVariable)) {
-            PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
-            MapOp op = new MapOp(myPreviousOp, rValue, myVariable, myVariable.getType());
-            return new TerminalBlock(op, myVariable, leftOver);
-          }
+        PsiExpression rValue = ExpressionUtils.getAssignmentTo(first, myVariable);
+        if(rValue != null) {
+          PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
+          MapOp op = new MapOp(myPreviousOp, rValue, myVariable, myVariable.getType());
+          return new TerminalBlock(op, myVariable, leftOver);
         }
       }
       return null;

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

@@ -736,6 +736,25 @@ public class ExpressionUtils {
     return null;
   }
 
+  /**
+   * Returns an expression assigned to the target variable if supplied element is
+   * either simple (non-compound) assignment expression or an expression statement containing assignment expression
+   * and the corresponding assignment l-value is the reference to target variable.
+   *
+   * @param element element to get assignment expression from
+   * @param target a variable to extract an assignment to
+   * @return extracted assignment r-value or null if assignment is not found or assignment is compound or it's an assignment
+   * to the wrong variable
+   */
+  @Contract("null, _ -> null; _, null -> null")
+  public static PsiExpression getAssignmentTo(PsiElement element, PsiVariable target) {
+    PsiAssignmentExpression assignment = getAssignment(element);
+    if(assignment != null && isReferenceTo(assignment.getLExpression(), target)) {
+      return assignment.getRExpression();
+    }
+    return null;
+  }
+
   @Contract("null, _ -> false")
   public static boolean isLiteral(PsiElement element, Object value) {
     return element instanceof PsiLiteralExpression && value.equals(((PsiLiteralExpression)element).getValue());