private static PsiMethodCallExpression extractAddCall( PsiStatement body, PsiIfStatement ifStatement) { if (ifStatement != null) { final PsiStatement thenBranch = ifStatement.getThenBranch(); return extractAddCall(thenBranch, null); } PsiExpressionStatement stmt = null; if (body instanceof PsiBlockStatement) { final PsiStatement[] statements = ((PsiBlockStatement) body).getCodeBlock().getStatements(); if (statements.length == 1 && statements[0] instanceof PsiExpressionStatement) { stmt = (PsiExpressionStatement) statements[0]; } } else if (body instanceof PsiExpressionStatement) { stmt = (PsiExpressionStatement) body; } if (stmt != null) { final PsiExpression expression = stmt.getExpression(); if (expression instanceof PsiMethodCallExpression) { return (PsiMethodCallExpression) expression; } } return null; }
@Override public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) { final PsiForeachStatement foreachStatement = PsiTreeUtil.getParentOfType(descriptor.getPsiElement(), PsiForeachStatement.class); if (foreachStatement != null) { if (!FileModificationService.getInstance().preparePsiElementForWrite(foreachStatement)) return; PsiStatement body = foreachStatement.getBody(); final PsiExpression iteratedValue = foreachStatement.getIteratedValue(); if (body != null && iteratedValue != null) { restoreComments(foreachStatement, body); final PsiParameter parameter = foreachStatement.getIterationParameter(); final PsiIfStatement ifStmt = extractIfStatement(body); StringBuilder buffer = new StringBuilder(getIteratedValueText(iteratedValue)); if (ifStmt != null) { final PsiStatement thenBranch = ifStmt.getThenBranch(); LOG.assertTrue(thenBranch != null); buffer.append(".stream()"); buffer.append(createFiltersChainText(body, parameter, ifStmt)); body = thenBranch; } buffer.append(".").append(getForEachMethodName()).append("("); final String functionalExpressionText = createForEachFunctionalExpressionText(project, body, parameter); final PsiElementFactory elementFactory = JavaPsiFacade.getElementFactory(project); PsiExpressionStatement callStatement = (PsiExpressionStatement) elementFactory.createStatementFromText( buffer.toString() + functionalExpressionText + ");", foreachStatement); callStatement = (PsiExpressionStatement) foreachStatement.replace(callStatement); final PsiExpressionList argumentList = ((PsiCallExpression) callStatement.getExpression()).getArgumentList(); LOG.assertTrue(argumentList != null, callStatement.getText()); final PsiExpression[] expressions = argumentList.getExpressions(); LOG.assertTrue(expressions.length == 1); if (expressions[0] instanceof PsiFunctionalExpression && ((PsiFunctionalExpression) expressions[0]).getFunctionalInterfaceType() == null) { callStatement = (PsiExpressionStatement) callStatement.replace( elementFactory.createStatementFromText( buffer.toString() + "(" + parameter.getText() + ") -> " + wrapInBlock(body) + ");", callStatement)); } simplifyRedundantCast(callStatement); CodeStyleManager.getInstance(project) .reformat( JavaCodeStyleManager.getInstance(project).shortenClassReferences(callStatement)); } } }