PsiElementVisitor (@NotNull final ProblemsHolder holder, boolean isOnTheFly) { if (!JavaFeature.ADVANCED_COLLECTIONS_API.isFeatureSupported(holder.getFile())) { return PsiElementVisitor.EMPTY_VISITOR; } return new JavaElementVisitor() { @Override public void visitConditionalExpression(@NotNull PsiConditionalExpression expression) { MapCheckCondition condition = fromConditional(expression, myTreatGetNullAsContainsKey); if(condition == null || condition.hasVariable()) return; PsiExpression existsBranch = condition.getExistsBranch(expression.getThenExpression(), expression.getElseExpression()); PsiExpression noneBranch = condition.getNoneBranch(expression.getThenExpression(), expression.getElseExpression()); processGetPut(condition, existsBranch, existsBranch, noneBranch); } @Override public void visitLocalVariable(@NotNull PsiLocalVariable variable) { PsiExpression expression = variable.getInitializer(); PsiMethodCallExpression getCall = extractMapMethodCall(expression, "get"); if (getCall == null) return; List<PsiReferenceExpression> references = VariableAccessUtils .getVariableReferences(variable, PsiTreeUtil.getParentOfType(variable, PsiCodeBlock.class)); if (references.isEmpty()) return; PsiMethodCallExpression commonPutCall = findPutMethodParent(references.get(0).getElement()); if (commonPutCall == null || !isCommonPutCallForAllReferences(references, commonPutCall)) return; PsiExpression getCallQualifierExpression = getCall.getMethodExpression().getQualifierExpression(); PsiExpression putCallQualifierExpression = commonPutCall.getMethodExpression().getQualifierExpression(); EquivalenceChecker equivalenceChecker = getCanonicalPsiEquivalence(); if (! equivalenceChecker.expressionsAreEquivalent(getCallQualifierExpression, putCallQualifierExpression)) return; PsiStatement variableDeclarationStatement = PsiTreeUtil.getParentOfType(variable, PsiDeclarationStatement.class); PsiElement nextSibling = PsiTreeUtil.skipWhitespacesAndCommentsForward(variableDeclarationStatement); if (! (nextSibling instanceof PsiStatement)) return; PsiExpressionStatement putCallStatement = ObjectUtils.tryCast(commonPutCall.getParent(), PsiExpressionStatement.class); if (nextSibling != putCallStatement) return; PsiExpression[] getArgs = getCall.getArgumentList().getExpressions(); PsiExpression[] putArgs = commonPutCall.getArgumentList().getExpressions(); if (getArgs.length != 1 || putArgs.length != 2 || ! equivalenceChecker.expressionsAreEquivalent(getArgs[0], putArgs[0])) return; PsiExpression value = putArgs[1]; if (LambdaGenerationUtil.canBeUncheckedLambda(value)) { GetPutToComputeFix fix = new GetPutToComputeFix(variable); holder.registerProblem(commonPutCall, QuickFixBundle.message("java.8.map.api.inspection.description", "compute"), fix); } } @Override public void visitIfStatement(@NotNull PsiIfStatement statement) { MapCheckCondition condition = fromConditional(statement, myTreatGetNullAsContainsKey); if(condition == null) return; PsiStatement existsBranch = ControlFlowUtils.stripBraces(condition.getExistsBranch(statement.getThenBranch(), statement.getElseBranch())); PsiStatement noneBranch = ControlFlowUtils.stripBraces(condition.getNoneBranch(statement.getThenBranch(), statement.getElseBranch())); if(existsBranch == null) { processSingleBranch(condition, noneBranch); } else { if(mySuggestMapMerge && condition.isGetNull()) { processMerge(condition, existsBranch, noneBranch); } if(condition.hasVariable()) return; EquivalenceChecker.Match match = getCanonicalPsiEquivalence().statementsMatch(noneBranch, existsBranch); processGetPut(condition, existsBranch, match.getRightDiff(), match.getLeftDiff()); } } @Override public void visitForeachStatement(@NotNull PsiForeachStatement statement) { if (!mySuggestMapReplaceAll) return; MapLoopCondition condition = MapLoopCondition.create(statement); if (condition == null) return; PsiMethodCallExpression putCall = condition.extractPut(statement); if (putCall == null) return; PsiExpression[] args = putCall.getArgumentList().getExpressions(); if (args.length != 2) return; PsiExpression key = args[0]; if (!condition.isKeyAccess(key)) return; PsiExpression value = args[1]; if (condition.isEntrySet() && isUsedAsReference(value, condition)) return; if (hasMapUsages(condition, value)) return; if (!LambdaGenerationUtil.canBeUncheckedLambda(value, variable -> condition.getMap().equals(variable))) return; ReplaceWithSingleMapOperation fix = ReplaceWithSingleMapOperation.create("replaceAll", putCall, value); holder.registerProblem(statement.getFirstChild(), QuickFixBundle.message("java.8.map.api.inspection.description", fix.myMethodName), fix); } private static PsiMethodCallExpression findPutMethodParent(PsiElement element) { while (element != null && !(element instanceof PsiMethod)) { if (element instanceof PsiExpression expression) { PsiMethodCallExpression putCall = extractMapMethodCall(expression, "put"); if (putCall != null) return putCall; } element = element.getParent(); } return null; } private static boolean isCommonPutCallForAllReferences(List<PsiReferenceExpression> references, PsiMethodCallExpression commonPutCall) { for (PsiReferenceExpression reference : references) { PsiMethodCallExpression putCall = findPutMethodParent(reference); if (putCall != commonPutCall) { return false; } } return true; } private static boolean hasMapUsages(@NotNull MapLoopCondition condition, @Nullable PsiExpression value) { return !VariableAccessUtils.getVariableReferences(condition.getMap(), value).stream() .map(ExpressionUtils::getCallForQualifier) .allMatch(call -> condition.isValueAccess(call)); } private static boolean isUsedAsReference(@NotNull PsiElement value, @NotNull MapLoopCondition condition) { return !VariableAccessUtils.getVariableReferences(condition.getIterParam(), value).stream() .map(ExpressionUtils::getCallForQualifier) .allMatch(KEY_VALUE_GET_METHODS); } private void processMerge(MapCheckCondition condition, PsiStatement existsBranch, PsiStatement noneBranch) { if(noneBranch instanceof PsiExpressionStatement && existsBranch instanceof PsiExpressionStatement) { PsiExpression absentValue = extractPutValue(condition, noneBranch); if (absentValue == null) return; PsiExpression presentValue = extractPutValue(condition, existsBranch); if (presentValue == null || !LambdaGenerationUtil.canBeUncheckedLambda(presentValue)) return; // absentValue should not refer map if (!PsiTreeUtil.processElements(absentValue, e -> !condition.isMap(e))) return; boolean hasVariable = condition.hasVariable(); if (hasVariable && PsiTreeUtil.collectElements(presentValue, condition::isValueReference).length == 0) return; PsiElement[] mapRefs = PsiTreeUtil.collectElements(presentValue, condition::isMap); if(hasVariable ^ mapRefs.length == 0) return; for(PsiElement mapRef : mapRefs) { PsiElement parent = mapRef.getParent(); if (!(parent instanceof PsiReferenceExpression) || condition.extractGetCall(parent.getParent()) == null) return; } if (PsiTreeUtil.collectElements(presentValue, e -> PsiEquivalenceUtil.areElementsEquivalent(e, absentValue)).length == 0) { return; } if (NullabilityUtil.getExpressionNullability(absentValue) == Nullability.NULLABLE || NullabilityUtil.getExpressionNullability(presentValue) == Nullability.NULLABLE) { return; } boolean informationLevel = !mySideEffects && SideEffectChecker.mayHaveSideEffects(presentValue, ex -> condition.extractGetCall(ex) != null); register(condition, holder, informationLevel, new ReplaceWithSingleMapOperation("merge", PsiTreeUtil .getParentOfType(absentValue, PsiMethodCallExpression.class), presentValue, noneBranch)); } } private void processGetPut(MapCheckCondition condition, PsiElement result, PsiElement exists, PsiElement none) { PsiMethodCallExpression getCall = condition.extractGetCall(exists); if(getCall == null) return; if(!(none instanceof PsiExpression noneExpression)) return; PsiMethodCallExpression putCall = extractMapMethodCall(noneExpression, "put"); if (mySuggestMapPutIfAbsent && putCall != null && condition.isGetNull() && condition.isMap(putCall.getMethodExpression().getQualifierExpression())) { PsiExpression[] putArgs = putCall.getArgumentList().getExpressions(); if (putArgs.length != 2 || !condition.isKey(putArgs[0]) || !ExpressionUtils.isSafelyRecomputableExpression(putArgs[1])) return; register(condition, holder, false, new ReplaceWithSingleMapOperation("putIfAbsent", getCall, putArgs[1], result)); } if (mySuggestMapGetOrDefault && condition.isContainsKey() && ExpressionUtils.isSafelyRecomputableExpression(noneExpression) && condition.isMapValueType(noneExpression.getType())) { register(condition, holder, false, new ReplaceWithSingleMapOperation("getOrDefault", getCall, noneExpression, result)); } } private void processSingleBranch(MapCheckCondition condition, PsiStatement noneBranch) { PsiAssignmentExpression assignment = ExpressionUtils.getAssignment(noneBranch); if(assignment != null && mySuggestMapGetOrDefault && condition.isContainsKey()) { /* value = map.get(key); if(value == null) { value = ... } */ PsiExpression rValue = assignment.getRExpression(); if (ExpressionUtils.isSafelyRecomputableExpression(rValue) && condition.isValueReference(assignment.getLExpression()) && !condition.isValueReference(rValue) && condition.isMapValueType(rValue.getType())) { register(condition, holder, false, ReplaceWithSingleMapOperation.fromIf("getOrDefault", condition, rValue)); } } else if (condition.isGetNull()) { /* value = map.get(key); if(value == null) { value = ... map.put(key, value); } */ PsiExpression lambdaCandidate = extractLambdaCandidate(condition, noneBranch); if (lambdaCandidate != null && mySuggestMapComputeIfAbsent) { if (NullabilityUtil.getExpressionNullability(lambdaCandidate) == Nullability.NULLABLE) return; boolean informationLevel = !mySideEffects && SideEffectChecker.mayHaveSideEffects(lambdaCandidate); register(condition, holder, informationLevel, ReplaceWithSingleMapOperation.fromIf("computeIfAbsent", condition, lambdaCandidate)); } if (lambdaCandidate == null) { PsiExpression expression = extractPutValue(condition, noneBranch); if(expression != null) { String replacement = null; boolean informationLevel = false; if (mySuggestMapPutIfAbsent && ExpressionUtils.isSafelyRecomputableExpression(expression) && !condition.isValueReference(expression)) { replacement = "putIfAbsent"; } else if (mySuggestMapComputeIfAbsent && !condition.hasVariable()) { informationLevel = !mySideEffects && SideEffectChecker.mayHaveSideEffects(expression); replacement = "computeIfAbsent"; } if(replacement != null) { if(condition.hasVariable()) { register(condition, holder, informationLevel, ReplaceWithSingleMapOperation.fromIf(replacement, condition, expression)); } else { PsiMethodCallExpression call = PsiTreeUtil.getParentOfType(expression, PsiMethodCallExpression.class); LOG.assertTrue(call != null); register(condition, holder, informationLevel, new ReplaceWithSingleMapOperation(replacement, call, expression, noneBranch)); } } } } } } }; }