From 2e1c8b6067d9b4e586f9c210f226df60f9615786 Mon Sep 17 00:00:00 2001 From: Tagir Valeev Date: Mon, 19 Sep 2016 14:33:14 +0700 Subject: [PATCH] IDEA-161260 convert to findFirst() if assignment is followed by break --- .../StreamApiMigrationInspection.java | 130 +++++++++++++----- .../Java8CollectionsApiInspection.java | 19 +-- .../afterFindFirstAssignment.java | 12 ++ .../afterFindFirstAssignmentIntermediate.java | 15 ++ .../afterFindFirstAssignmentWithLabel.java | 13 ++ .../beforeFindFirstAssignment.java | 17 +++ ...beforeFindFirstAssignmentIntermediate.java | 19 +++ .../beforeFindFirstAssignmentNested.java | 22 +++ .../beforeFindFirstAssignmentWithLabel.java | 22 +++ .../siyeh/ig/psiutils/ExpressionUtils.java | 20 ++- 10 files changed, 242 insertions(+), 47 deletions(-) create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignment.java create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentIntermediate.java create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentWithLabel.java create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignment.java create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentIntermediate.java create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentNested.java create mode 100644 java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentWithLabel.java diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java b/java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java index da494ef99498..0c329d21620c 100644 --- a/java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java +++ b/java/java-analysis-impl/src/com/intellij/codeInspection/StreamApiMigrationInspection.java @@ -186,6 +186,25 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo if(nonFinalVariables.isEmpty() && tb.getSingleStatement() instanceof PsiReturnStatement) { handleSingleReturn(statement, tb, operations); } + if(nonFinalVariables.size() == 1) { + PsiStatement[] statements = tb.getStatements(); + if (statements.length == 2) { + PsiStatement breakStatement = statements[1]; + if (!(breakStatement instanceof PsiBreakStatement) || + ((PsiBreakStatement)breakStatement).findExitedStatement() != statement) { + return; + } + PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(statements[0]); + if(assignment == null) return; + PsiExpression lValue = assignment.getLExpression(); + if (!(lValue instanceof PsiReferenceExpression)) return; + PsiElement var = ((PsiReferenceExpression)lValue).resolve(); + if(!(var instanceof PsiVariable) || !nonFinalVariables.contains(var)) return; + PsiExpression rValue = assignment.getRExpression(); + if(rValue == null || isVariableReferenced((PsiVariable)var, rValue)) return; + registerProblem(holder, isOnTheFly, statement, "findFirst", new ReplaceWithFindFirstFix()); + } + } } } } @@ -327,12 +346,12 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo } @Nullable - private static PsiLocalVariable extractAccumulator(PsiAssignmentExpression assignment) { + private static PsiVariable extractAccumulator(PsiAssignmentExpression assignment) { if(!(assignment.getLExpression() instanceof PsiReferenceExpression)) return null; PsiReferenceExpression lExpr = (PsiReferenceExpression)assignment.getLExpression(); PsiElement accumulator = lExpr.resolve(); - if(!(accumulator instanceof PsiLocalVariable)) return null; - PsiLocalVariable var = (PsiLocalVariable)accumulator; + if(!(accumulator instanceof PsiVariable)) return null; + PsiVariable var = (PsiVariable)accumulator; if(JavaTokenType.PLUSEQ.equals(assignment.getOperationTokenType())) { return var; } else if(JavaTokenType.EQ.equals(assignment.getOperationTokenType())) { @@ -392,7 +411,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo } @Nullable - private static PsiLocalVariable getAccumulatedVariable(TerminalBlock tb, + private static PsiVariable getAccumulatedVariable(TerminalBlock tb, List operations, List variables) { // have only one non-final variable @@ -400,7 +419,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo PsiAssignmentExpression assignment = tb.getSingleExpression(PsiAssignmentExpression.class); if(assignment == null) return null; - PsiLocalVariable var = extractAccumulator(assignment); + PsiVariable var = extractAccumulator(assignment); // the referred variable is the same as non-final variable if(var == null || !variables.contains(var)) return null; @@ -572,10 +591,10 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo @NotNull List replacements); static void replaceWithNumericAddition(@NotNull Project project, - PsiForeachStatement foreachStatement, - PsiLocalVariable var, - StringBuilder builder, - String expressionType) { + PsiForeachStatement foreachStatement, + PsiVariable var, + StringBuilder builder, + String expressionType) { PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project); restoreComments(foreachStatement, foreachStatement.getBody()); if (isDeclarationJustBefore(var, foreachStatement)) { @@ -584,7 +603,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo String typeStr = var.getType().getCanonicalText(); String replacement = (typeStr.equals(expressionType) ? "" : "(" + typeStr + ") ") + builder; initializer.replace(elementFactory.createExpressionFromText(replacement, foreachStatement)); - foreachStatement.delete(); + removeLoop(foreachStatement); simplifyAndFormat(project, var); return; } @@ -630,6 +649,15 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo iteratedValue instanceof PsiQualifiedExpression || iteratedValue instanceof PsiParenthesizedExpression ? iteratedValue.getText() : "(" + iteratedValue.getText() + ")"; } + + static void removeLoop(@NotNull PsiForeachStatement statement) { + PsiElement parent = statement.getParent(); + if(parent instanceof PsiLabeledStatement) { + parent.delete(); + } else { + statement.delete(); + } + } } private static class ReplaceWithForeachCallFix extends MigrateToStreamFix { @@ -769,7 +797,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo .append(")").toString(); PsiElement result = initializer.replace(elementFactory.createExpressionFromText(callText, null)); simplifyAndFormat(project, result); - foreachStatement.delete(); + removeLoop(foreachStatement); return; } final String qualifierText = qualifierExpression != null ? qualifierExpression.getText() + "." : ""; @@ -942,25 +970,57 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo @NotNull PsiStatement body, @NotNull TerminalBlock tb, @NotNull List intermediateOps) { - PsiReturnStatement returnStatement = (PsiReturnStatement)tb.getSingleStatement(); - PsiExpression value = returnStatement.getReturnValue(); - if(value == null) return; - PsiReturnStatement nextReturnStatement = getNextReturnStatement(foreachStatement); - if(nextReturnStatement == null) return; - PsiExpression orElseExpression = nextReturnStatement.getReturnValue(); - if(!ExpressionUtils.isSimpleExpression(orElseExpression)) return; - final PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project); - StringBuilder builder = generateStream(iteratedValue, intermediateOps).append(".findFirst()"); - if (!(value instanceof PsiReferenceExpression) || ((PsiReferenceExpression)value).resolve() != tb.getVariable()) { - builder.append(".map(").append(tb.getVariable().getName()).append(" -> ").append(value.getText()).append(")"); + PsiStatement statement = tb.getSingleStatement(); + PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project); + if(statement instanceof PsiReturnStatement) { + PsiReturnStatement returnStatement = (PsiReturnStatement)statement; + PsiExpression value = returnStatement.getReturnValue(); + if (value == null) return; + PsiReturnStatement nextReturnStatement = getNextReturnStatement(foreachStatement); + if (nextReturnStatement == null) return; + PsiExpression orElseExpression = nextReturnStatement.getReturnValue(); + if (!ExpressionUtils.isSimpleExpression(orElseExpression)) return; + StringBuilder builder = generateStream(iteratedValue, intermediateOps).append(".findFirst()"); + if (!(value instanceof PsiReferenceExpression) || ((PsiReferenceExpression)value).resolve() != tb.getVariable()) { + builder.append(".map(").append(tb.getVariable().getName()).append(" -> ").append(value.getText()).append(")"); + } + builder.append(".orElse(").append(orElseExpression.getText()).append(")"); + boolean siblings = nextReturnStatement.getParent() == foreachStatement.getParent(); + PsiElement result = foreachStatement.replace(elementFactory.createStatementFromText("return " + builder + ";", foreachStatement)); + if (siblings) { + nextReturnStatement.delete(); + } + simplifyAndFormat(project, result); + } else { + PsiStatement[] statements = tb.getStatements(); + if(statements.length != 2) return; + PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(statements[0]); + if(assignment == null) return; + PsiExpression lValue = assignment.getLExpression(); + if(!(lValue instanceof PsiReferenceExpression)) return; + PsiElement element = ((PsiReferenceExpression)lValue).resolve(); + if(!(element instanceof PsiVariable)) return; + PsiVariable var = (PsiVariable)element; + PsiExpression value = assignment.getRExpression(); + if(value == null) return; + StringBuilder builder = generateStream(iteratedValue, intermediateOps).append(".findFirst()"); + if (!(value instanceof PsiReferenceExpression) || ((PsiReferenceExpression)value).resolve() != tb.getVariable()) { + builder.append(".map(").append(tb.getVariable().getName()).append(" -> ").append(value.getText()).append(")"); + } + if(isDeclarationJustBefore(var, foreachStatement)) { + PsiExpression initializer = var.getInitializer(); + if(initializer != null) { + PsiElement result = + initializer.replace(elementFactory.createExpressionFromText(builder + ".orElse(" + initializer.getText() + ")", initializer)); + removeLoop(foreachStatement); + simplifyAndFormat(project, result); + return; + } + } + PsiElement result = foreachStatement.replace( + elementFactory.createStatementFromText(var.getName() + " = " + builder + ".orElse(" + var.getName() + ");", foreachStatement)); + simplifyAndFormat(project, result); } - builder.append(".orElse(").append(orElseExpression.getText()).append(")"); - boolean siblings = nextReturnStatement.getParent() == foreachStatement.getParent(); - PsiElement result = foreachStatement.replace(elementFactory.createStatementFromText("return " + builder + ";", foreachStatement)); - if(siblings) { - nextReturnStatement.delete(); - } - simplifyAndFormat(project, result); } } @@ -982,7 +1042,7 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo @NotNull List intermediateOps) { PsiAssignmentExpression assignment = tb.getSingleExpression(PsiAssignmentExpression.class); if (assignment == null) return; - PsiLocalVariable var = extractAccumulator(assignment); + PsiVariable var = extractAccumulator(assignment); if (var == null) return; PsiExpression addend = extractAddend(assignment); @@ -1008,8 +1068,12 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo } } - private static boolean isDeclarationJustBefore(PsiLocalVariable var, PsiStatement nextStatement) { + private static boolean isDeclarationJustBefore(PsiVariable var, PsiStatement nextStatement) { PsiElement declaration = var.getParent(); + PsiElement nextStatementParent = nextStatement.getParent(); + if(nextStatementParent instanceof PsiLabeledStatement) { + nextStatement = (PsiStatement)nextStatementParent; + } if(declaration instanceof PsiDeclarationStatement) { PsiElement[] elements = ((PsiDeclarationStatement)declaration).getDeclaredElements(); if (ArrayUtil.getLastElement(elements) == var && nextStatement.equals( @@ -1122,6 +1186,10 @@ public class StreamApiMigrationInspection extends BaseJavaBatchLocalInspectionTo return myStatements.length == 1 ? myStatements[0] : null; } + PsiStatement[] getStatements() { + return myStatements; + } + @Nullable T getSingleExpression(Class wantedType) { PsiStatement statement = getSingleStatement(); diff --git a/java/java-analysis-impl/src/com/intellij/codeInspection/java18api/Java8CollectionsApiInspection.java b/java/java-analysis-impl/src/com/intellij/codeInspection/java18api/Java8CollectionsApiInspection.java index d861bd7ac5d1..471d2301e39a 100644 --- a/java/java-analysis-impl/src/com/intellij/codeInspection/java18api/Java8CollectionsApiInspection.java +++ b/java/java-analysis-impl/src/com/intellij/codeInspection/java18api/Java8CollectionsApiInspection.java @@ -135,7 +135,7 @@ public class Java8CollectionsApiInspection extends BaseJavaBatchLocalInspectionT PsiExpression[] getArguments = getCall.getArgumentList().getExpressions(); if(getArguments.length != 1) return; PsiStatement thenBranch = ControlFlowUtils.stripBraces(statement.getThenBranch()); - PsiAssignmentExpression assignment = getAssignment(thenBranch); + PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(thenBranch); EquivalenceChecker equivalence = EquivalenceChecker.getCanonicalPsiEquivalence(); if(assignment != null) { /* @@ -160,7 +160,7 @@ public class Java8CollectionsApiInspection extends BaseJavaBatchLocalInspectionT PsiExpression key = getArguments[0]; PsiStatement[] statements = ((PsiBlockStatement)thenBranch).getCodeBlock().getStatements(); if(statements.length != 2) return; - assignment = getAssignment(statements[0]); + assignment = ExpressionUtils.getAssignment(statements[0]); if(assignment == null) return; PsiExpression lambdaCandidate = assignment.getRExpression(); if (lambdaCandidate == null || @@ -205,17 +205,6 @@ public class Java8CollectionsApiInspection extends BaseJavaBatchLocalInspectionT return (PsiReferenceExpression)value; } - @Contract("null -> null") - private static PsiAssignmentExpression getAssignment(PsiElement element) { - if(element instanceof PsiExpressionStatement) { - element = ((PsiExpressionStatement)element).getExpression(); - } - if (element instanceof PsiAssignmentExpression) { - return (PsiAssignmentExpression)element; - } - return null; - } - @Nullable @Contract("_, null -> null") static PsiMethodCallExpression tryExtractMapGetCall(PsiReferenceExpression target, PsiElement element) { @@ -234,7 +223,7 @@ public class Java8CollectionsApiInspection extends BaseJavaBatchLocalInspectionT } } } - PsiAssignmentExpression assignment = getAssignment(element); + PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(element); if(assignment != null) { PsiExpression lValue = assignment.getLExpression(); if (lValue instanceof PsiReferenceExpression && @@ -513,7 +502,7 @@ public class Java8CollectionsApiInspection extends BaseJavaBatchLocalInspectionT } else if(thenBranch instanceof PsiBlockStatement) { PsiStatement[] statements = ((PsiBlockStatement)thenBranch).getCodeBlock().getStatements(); if(statements.length != 2) return; - PsiAssignmentExpression assignment = getAssignment(statements[0]); + PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(statements[0]); if(assignment == null) return; PsiExpression lambdaCandidate = assignment.getRExpression(); if(lambdaCandidate == null) return; diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignment.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignment.java new file mode 100644 index 000000000000..6d2085e355d2 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignment.java @@ -0,0 +1,12 @@ +// "Replace with findFirst()" "true" + +import java.util.List; +import java.util.Map; +import java.util.Objects; + +public class Main { + public void testMap(Map> map) throws Exception { + int firstSize = map.values().stream().filter(Objects::nonNull).findFirst().map(List::size).orElse(0); + System.out.println(firstSize); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentIntermediate.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentIntermediate.java new file mode 100644 index 000000000000..8de459b0cf68 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentIntermediate.java @@ -0,0 +1,15 @@ +// "Replace with findFirst()" "true" + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +public class Main { + public void testMap(Map> map) throws Exception { + List firstList = new ArrayList<>(); + firstList.add("none"); + firstList = map.values().stream().filter(Objects::nonNull).findFirst().orElse(firstList); + System.out.println(firstList); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentWithLabel.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentWithLabel.java new file mode 100644 index 000000000000..e901baff38ee --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/afterFindFirstAssignmentWithLabel.java @@ -0,0 +1,13 @@ +// "Replace with findFirst()" "true" + +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +public class Main { + public void testMap(Map> map) throws Exception { + String firstStr = map.values().stream().filter(Objects::nonNull).flatMap(Collection::stream).filter(str -> !str.isEmpty()).findFirst().orElse(""); + System.out.println(firstStr); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignment.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignment.java new file mode 100644 index 000000000000..f502b379e643 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignment.java @@ -0,0 +1,17 @@ +// "Replace with findFirst()" "true" + +import java.util.List; +import java.util.Map; + +public class Main { + public void testMap(Map> map) throws Exception { + int firstSize = 0; + for(List list : map.values()) { + if(list != null) { + firstSize = list.size(); + break; + } + } + System.out.println(firstSize); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentIntermediate.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentIntermediate.java new file mode 100644 index 000000000000..2b05c0c36845 --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentIntermediate.java @@ -0,0 +1,19 @@ +// "Replace with findFirst()" "true" + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +public class Main { + public void testMap(Map> map) throws Exception { + List firstList = new ArrayList<>(); + firstList.add("none"); + for(List list : map.values()) { + if(list != null) { + firstList = list; + break; + } + } + System.out.println(firstList); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentNested.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentNested.java new file mode 100644 index 000000000000..3bc2741e77bb --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentNested.java @@ -0,0 +1,22 @@ +// "Replace with findFirst()" "false" + +import java.util.List; +import java.util.Map; + +public class Main { + public void testMap(Map> map) throws Exception { + String firstStr = ""; + OUTER: + for(List list : map.values()) { + if(list != null) { + for(String str : list) { + if(!str.isEmpty()) { + firstStr = str; + break; + } + } + } + } + System.out.println(firstStr); + } +} \ No newline at end of file diff --git a/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentWithLabel.java b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentWithLabel.java new file mode 100644 index 000000000000..287fbc9d68dd --- /dev/null +++ b/java/java-tests/testData/codeInsight/daemonCodeAnalyzer/quickFix/streamApiMigration/beforeFindFirstAssignmentWithLabel.java @@ -0,0 +1,22 @@ +// "Replace with findFirst()" "true" + +import java.util.List; +import java.util.Map; + +public class Main { + public void testMap(Map> map) throws Exception { + String firstStr = ""; + OUTER: + for(List list : map.values()) { + if(list != null) { + for(String str : list) { + if(!str.isEmpty()) { + firstStr = str; + break OUTER; + } + } + } + } + System.out.println(firstStr); + } +} \ No newline at end of file diff --git a/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java b/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java index 950c8e6119c3..ed6f57a71a5a 100644 --- a/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java +++ b/plugins/InspectionGadgets/InspectionGadgetsAnalysis/src/com/siyeh/ig/psiutils/ExpressionUtils.java @@ -601,7 +601,7 @@ public class ExpressionUtils { @Nullable public static PsiVariable getVariableFromNullComparison(PsiExpression expression, boolean equals) { final PsiReferenceExpression referenceExpression = getReferenceExpressionFromNullComparison(expression, equals); - final PsiElement target = referenceExpression.resolve(); + final PsiElement target = referenceExpression != null ? referenceExpression.resolve() : null; return target instanceof PsiVariable ? (PsiVariable)target : null; } @@ -700,4 +700,22 @@ public class ExpressionUtils { } return false; } + + /** + * Returns assignment expression if supplied element is a statement which contains assignment expression + * or it's an assignment expression itself + * + * @param element element to get assignment expression from + * @return extracted assignment or null if assignment is not found + */ + @Contract("null -> null") + public static PsiAssignmentExpression getAssignment(PsiElement element) { + if(element instanceof PsiExpressionStatement) { + element = ((PsiExpressionStatement)element).getExpression(); + } + if (element instanceof PsiAssignmentExpression) { + return (PsiAssignmentExpression)element; + } + return null; + } } \ No newline at end of file