void (@NotNull Project project, @NotNull PsiElement element, @NotNull ModPsiUpdater updater) { PsiElement outerElement = PsiTreeUtil.getParentOfType(element, PsiIfStatement.class, PsiConditionalExpression.class, PsiForeachStatement.class); if (outerElement == null) return; MapCondition condition = outerElement instanceof PsiForeachStatement ? MapLoopCondition.create((PsiForeachStatement)outerElement) : fromConditional(outerElement, true); if(condition == null) return; PsiMethodCallExpression call = updater.getWritable(myCallPointer.getElement()); if (call == null) return; PsiExpressionList argsList = call.getArgumentList(); PsiExpression[] args = argsList.getExpressions(); if(args.length == 0) return; if ((myMethodName.equals("merge") || myMethodName.equals("replaceAll")) && args.length != 2) return; PsiExpression value = updater.getWritable(myValuePointer.getElement()); if (value == null) return; PsiElement result = updater.getWritable(myResultPointer.getElement()); if(result == null) return; PsiElementFactory factory = JavaPsiFacade.getElementFactory(project); CommentTracker ct = new CommentTracker(); ExpressionUtils.bindCallTo(call, myMethodName); PsiExpression replacement; switch (myMethodName) { case "computeIfAbsent" -> { PsiExpression key = args[0]; List<PsiReferenceExpression> refs = Collections.emptyList(); String nameCandidate = "k"; if (key instanceof PsiReferenceExpression && ((PsiReferenceExpression)key).getQualifier() == null) { // try to use lambda parameter if key is simple reference and has the same type as map keys PsiMethod method = call.resolveMethod(); if (method != null) { PsiType argType = method.getParameterList().getParameters()[0].getType(); PsiType mapKeyType = call.resolveMethodGenerics().getSubstitutor().substitute(argType); PsiType keyType = key.getType(); if (mapKeyType != null && keyType != null && keyType.isAssignableFrom(mapKeyType)) { PsiElement target = ((PsiReferenceExpression)key).resolve(); refs = target == null ? Collections.emptyList() : StreamEx.of(PsiTreeUtil.collectElementsOfType(value, PsiReferenceExpression.class)) .filter(ref -> ref.getQualifierExpression() == null && ref.isReferenceTo(target)).toList(); if (!refs.isEmpty()) { nameCandidate = getNameCandidate(((PsiReferenceExpression)key).getReferenceName()); } } } } String varName = JavaCodeStyleManager.getInstance(project).suggestUniqueVariableName(nameCandidate, value, true); for (PsiReferenceExpression ref : refs) { ExpressionUtils.bindReferenceTo(ref, varName); } replacement = factory.createExpressionFromText(varName + " -> " + ct.text(value), value); } case "merge" -> { MapCheckCondition checkCondition = ObjectUtils.tryCast(condition, MapCheckCondition.class); if (checkCondition == null) return; PsiExpression absentValue = args[1]; String aVar = JavaCodeStyleManager.getInstance(project).suggestUniqueVariableName("a", value, true); String bVar = JavaCodeStyleManager.getInstance(project).suggestUniqueVariableName("b", value, true); for (PsiElement e : PsiTreeUtil.collectElements(value, e -> PsiEquivalenceUtil.areElementsEquivalent(e, absentValue))) { ct.replace(e, factory.createIdentifier(bVar)); } for (PsiElement e : PsiTreeUtil .collectElements(value, e -> checkCondition.extractGetCall(e) != null || checkCondition.isValueReference(e))) { ct.replace(e, factory.createIdentifier(aVar)); } replacement = factory.createExpressionFromText("(" + aVar + "," + bVar + ") -> " + ct.text(value), value); } case "replaceAll" -> { MapLoopCondition loopCondition = ObjectUtils.tryCast(condition, MapLoopCondition.class); if (loopCondition == null) return; String kVar = suggestKeyName(loopCondition, value); String vVar = new VariableNameGenerator(value, VariableKind.PARAMETER).byName("v", "value").generate(true); replacement = createLambdaForLoopReplacement(factory, kVar, vVar, loopCondition, value, ct); ct.delete(args); } default -> replacement = ct.markUnchanged(value); } PsiElement newArg; if (args.length == 2 && !myMethodName.equals("merge") && !myMethodName.equals("replaceAll")) { newArg = ct.replace(args[1], replacement); } else { newArg = argsList.add(replacement); } if(newArg instanceof PsiLambdaExpression) { LambdaCanBeMethodReferenceInspection.replaceLambdaWithMethodReference((PsiLambdaExpression)newArg); } if (PsiTreeUtil.isAncestor(outerElement, result, true)) { result = ct.replaceAndRestoreComments(outerElement, result); } else { ct.deleteAndRestoreComments(outerElement); } PsiVariable variable = condition instanceof MapCheckCondition ? ((MapCheckCondition)condition).extractDeclaration() : null; if (variable != null && !PsiTreeUtil.isAncestor(result, variable, true) && ReferencesSearch.search(variable).findFirst() == null) { new CommentTracker().deleteAndRestoreComments(variable); } CodeStyleManager.getInstance(project).reformat(result); }