/*
* Copyright 2000-2017 JetBrains s.r.o.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.intellij.codeInspection.streamMigration;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.psi.*;
import com.intellij.psi.controlFlow.ControlFlow;
import com.intellij.psi.controlFlow.ControlFlowUtil;
import com.intellij.psi.search.LocalSearchScope;
import com.intellij.psi.search.searches.ReferencesSearch;
import com.intellij.psi.util.PsiUtil;
import com.intellij.util.ArrayUtil;
import com.intellij.util.containers.IntArrayList;
import com.siyeh.ig.psiutils.*;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.Arrays;
import java.util.Collection;
import java.util.NoSuchElementException;
import java.util.Objects;
import static com.intellij.codeInspection.streamMigration.StreamApiMigrationInspection.*;
import static com.intellij.util.ObjectUtils.tryCast;
/**
* This immutable class represents the code which should be performed
* as a part of forEach operation of resulting stream possibly with
* some intermediate operations extracted.
*/
class TerminalBlock {
private static final Logger LOG = Logger.getInstance(TerminalBlock.class);
private final @NotNull PsiVariable myVariable;
private final @NotNull Operation[] myOperations;
private final @NotNull PsiStatement[] myStatements;
private TerminalBlock(@NotNull Operation[] operations, @NotNull PsiVariable variable, @NotNull PsiStatement... statements) {
// At least one operation is present (stream source)
LOG.assertTrue(operations.length > 0);
for(Operation operation : operations) Objects.requireNonNull(operation);
for(PsiStatement statement : statements) Objects.requireNonNull(statement);
myVariable = variable;
while(true) {
if(statements.length == 1 && statements[0] instanceof PsiBlockStatement) {
statements = ((PsiBlockStatement)statements[0]).getCodeBlock().getStatements();
} else if(statements.length == 1 && statements[0] instanceof PsiLabeledStatement) {
PsiStatement statement = ((PsiLabeledStatement)statements[0]).getStatement();
statements = statement == null ? PsiStatement.EMPTY_ARRAY : new PsiStatement[] {statement};
} else break;
}
myStatements = statements;
myOperations = operations;
}
private TerminalBlock(@Nullable TerminalBlock previousBlock,
@NotNull Operation operation,
@NotNull PsiVariable variable,
@NotNull PsiStatement... statements) {
this(previousBlock == null ? new Operation[]{operation} : ArrayUtil.append(previousBlock.myOperations, operation), variable,
statements);
}
Collection<PsiStatement> findExitPoints(ControlFlow controlFlow) {
int startOffset = controlFlow.getStartOffset(myStatements[0]);
int endOffset = controlFlow.getEndOffset(myStatements[myStatements.length - 1]);
if (startOffset < 0 || endOffset < 0) return null;
return ControlFlowUtil
.findExitPointsAndStatements(controlFlow, startOffset, endOffset, new IntArrayList(), PsiContinueStatement.class,
PsiBreakStatement.class, PsiReturnStatement.class, PsiThrowStatement.class);
}
PsiStatement getSingleStatement() {
return myStatements.length == 1 ? myStatements[0] : null;
}
@NotNull
PsiStatement[] getStatements() {
return myStatements;
}
@Nullable
<T extends PsiExpression> T getSingleExpression(Class<T> wantedType) {
PsiStatement statement = getSingleStatement();
if(statement instanceof PsiExpressionStatement) {
PsiExpression expression = ((PsiExpressionStatement)statement).getExpression();
return tryCast(expression, wantedType);
}
return null;
}
/**
* @return PsiMethodCallExpression if this TerminalBlock contains single method call, null otherwise
*/
@Nullable
PsiMethodCallExpression getSingleMethodCall() {
return getSingleExpression(PsiMethodCallExpression.class);
}
@Nullable
private TerminalBlock extractFilter() {
if(getSingleStatement() instanceof PsiIfStatement) {
PsiIfStatement ifStatement = (PsiIfStatement)getSingleStatement();
if(ifStatement.getElseBranch() == null && ifStatement.getCondition() != null) {
PsiStatement thenBranch = ifStatement.getThenBranch();
if(thenBranch != null) {
return new TerminalBlock(this, new FilterOp(ifStatement.getCondition(), myVariable, false), myVariable, thenBranch);
}
}
}
if(myStatements.length >= 1) {
PsiStatement first = myStatements[0];
// extract filter with negation
if(first instanceof PsiIfStatement) {
PsiIfStatement ifStatement = (PsiIfStatement)first;
if(ifStatement.getCondition() == null) return null;
PsiStatement branch = ifStatement.getThenBranch();
if(branch instanceof PsiBlockStatement) {
PsiStatement[] statements = ((PsiBlockStatement)branch).getCodeBlock().getStatements();
if(statements.length == 1)
branch = statements[0];
}
if(!(branch instanceof PsiContinueStatement) || ((PsiContinueStatement)branch).getLabelIdentifier() != null) return null;
PsiStatement[] statements;
if(ifStatement.getElseBranch() != null) {
statements = myStatements.clone();
statements[0] = ifStatement.getElseBranch();
} else {
statements = Arrays.copyOfRange(myStatements, 1, myStatements.length);
}
return new TerminalBlock(this, new FilterOp(ifStatement.getCondition(), myVariable, true), myVariable, statements);
}
}
return null;
}
/**
* Returns an equivalent {@code TerminalBlock} with one more intermediate operation extracted
* or null if extraction is not possible.
*
* @return extracted operation or null if extraction is not possible
*/
@Nullable
private TerminalBlock extractOperation() {
TerminalBlock withFilter = extractFilter();
if(withFilter != null) return withFilter;
// extract flatMap
if(getSingleStatement() instanceof PsiLoopStatement) {
PsiLoopStatement loopStatement = (PsiLoopStatement)getSingleStatement();
StreamSource source = StreamSource.tryCreate(loopStatement);
final PsiStatement body = loopStatement.getBody();
if(source == null || body == null) return null;
FlatMapOp op = new FlatMapOp(source, myVariable);
TerminalBlock withFlatMap = new TerminalBlock(this, op, source.getVariable(), body);
if(!VariableAccessUtils.variableIsUsed(myVariable, body)) {
return withFlatMap;
} else {
// Try extract nested filter like this:
// for(List subList : list) for(T t : subList) if(condition.test(t)) { ...; break; }
// if t is not used in "...", then this could be converted to
// list.stream().filter(subList -> subList.stream().anyMatch(condition)).forEach(subList -> ...)
TerminalBlock withFlatMapFilter = withFlatMap.extractFilter();
if(withFlatMapFilter != null && !withFlatMapFilter.isEmpty()) {
PsiStatement[] statements = withFlatMapFilter.getStatements();
PsiStatement lastStatement = statements[statements.length-1];
if (lastStatement instanceof PsiBreakStatement && op.breaksMe((PsiBreakStatement)lastStatement) &&
ReferencesSearch.search(withFlatMapFilter.getVariable(), new LocalSearchScope(statements)).findFirst() == null) {
return new TerminalBlock(this, new CompoundFilterOp((FilterOp)withFlatMapFilter.getLastOperation(), op),
myVariable, Arrays.copyOfRange(statements, 0, statements.length-1));
}
}
}
}
if(myStatements.length >= 1) {
PsiStatement first = myStatements[0];
if(PsiUtil.isLanguageLevel9OrHigher(first.getContainingFile()) && first instanceof PsiIfStatement) {
PsiIfStatement ifStatement = (PsiIfStatement)first;
PsiExpression condition = ifStatement.getCondition();
if(ifStatement.getElseBranch() == null && condition != null) {
PsiStatement thenStatement = ControlFlowUtils.stripBraces(ifStatement.getThenBranch());
if(ControlFlowUtils.statementBreaksLoop(thenStatement, getMainLoop())) {
TakeWhileOp op = new TakeWhileOp(condition, myVariable, true);
PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
return new TerminalBlock(this, op, myVariable, leftOver);
}
}
}
// extract map
if(first instanceof PsiDeclarationStatement) {
PsiDeclarationStatement decl = (PsiDeclarationStatement)first;
PsiElement[] elements = decl.getDeclaredElements();
if(elements.length == 1) {
PsiLocalVariable declaredVar = tryCast(elements[0], PsiLocalVariable.class);
if (declaredVar != null && StreamApiUtil.isSupportedStreamElement(declaredVar.getType())) {
PsiExpression initializer = declaredVar.getInitializer();
PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
if (initializer != null && ReferencesSearch.search(myVariable, new LocalSearchScope(leftOver)).findFirst() == null) {
MapOp op = new MapOp(initializer, myVariable, declaredVar.getType());
return new TerminalBlock(this, op, declaredVar, leftOver);
}
}
}
}
PsiExpression rValue = ExpressionUtils.getAssignmentTo(first, myVariable);
if(rValue != null && operations().allMatch(op -> op.canReassignVariable(myVariable))) {
PsiStatement[] leftOver = Arrays.copyOfRange(myStatements, 1, myStatements.length);
MapOp op = new MapOp(rValue, myVariable, myVariable.getType());
return new TerminalBlock(this, op, myVariable, leftOver);
}
}
TerminalBlock withLimit = tryPeelLimit(true);
return withLimit == this ? null : withLimit;
}
/**
* Try to peel off the condition like if(count > limit) break; from the end of current terminal block.
*
* <p>It's not guaranteed that the peeled condition actually could be translated to the limit operation:
* additional checks will be necessary</p>
*
* @param dedicatedCounter whether peeled counter must be an increment statement for dedicated local variable
* @return new terminal block with additional limit operation or self if peeling is failed.
*/
private TerminalBlock tryPeelLimit(boolean dedicatedCounter) {
if(myStatements.length == 0) return this;
TerminalBlock tb = this;
PsiStatement[] statements = {};
if(myStatements.length > 1) {
int count = myStatements.length - 1;
if (myStatements[count] instanceof PsiBreakStatement || myStatements[count] instanceof PsiReturnStatement) {
// to support conditions like if(...) continue; break; or if(...) continue; return ...;
count--;
}
statements = Arrays.copyOfRange(myStatements, 0, count);
tb = new TerminalBlock(myOperations, myVariable, Arrays.copyOfRange(myStatements, count, myStatements.length)).extractFilter();
}
if (tb == null || !ControlFlowUtils.statementBreaksLoop(tb.getSingleStatement(), getMainLoop())) return this;
FilterOp filter = tb.getLastOperation(FilterOp.class);
if (filter == null) return this;
PsiBinaryExpression binOp = tryCast(PsiUtil.skipParenthesizedExprDown(filter.getExpression()), PsiBinaryExpression.class);
if (binOp == null || !ComparisonUtils.isComparison(binOp)) return this;
String comparison = filter.isNegated() ? ComparisonUtils.getNegatedComparison(binOp.getOperationTokenType())
: binOp.getOperationSign().getText();
boolean flipped = false;
int delta = 0;
switch (comparison) {
case "==":
case ">=":
break;
case ">":
delta = 1;
break;
case "<":
delta = 1;
flipped = true;
break;
case "<=":
flipped = true;
break;
default:
return this;
}
PsiExpression countExpression = PsiUtil.skipParenthesizedExprDown(flipped ? binOp.getROperand() : binOp.getLOperand());
if(countExpression == null || VariableAccessUtils.variableIsUsed(myVariable, countExpression)) return this;
PsiExpression incrementedValue = extractIncrementedLValue(countExpression);
PsiLocalVariable var = null;
if (dedicatedCounter) {
if (!(incrementedValue instanceof PsiReferenceExpression)) return this;
var = tryCast(((PsiReferenceExpression)incrementedValue).resolve(), PsiLocalVariable.class);
if (var == null || !ExpressionUtils.isZero(var.getInitializer()) || ReferencesSearch.search(var).findAll().size() != 1) return this;
}
PsiExpression limit = flipped ? binOp.getLOperand() : binOp.getROperand();
if(!ExpressionUtils.isSimpleExpression(limit) || VariableAccessUtils.variableIsUsed(myVariable, limit)) return this;
PsiType type = limit.getType();
if(!PsiType.INT.equals(type) && !PsiType.LONG.equals(type)) return this;
if(countExpression instanceof PsiPostfixExpression) {
delta++;
}
TerminalBlock block = new TerminalBlock(Arrays.copyOf(tb.myOperations, tb.myOperations.length-1), myVariable, statements);
if (incrementedValue == null) {
// when countExpression does not change the counter, we may try to continue extracting ops from the remaining statement
// this is helpful to cover cases like for(...) { if(...) list.add(x); if(list.size == limit) break; }
while (true) {
TerminalBlock newBlock = block.extractOperation();
if (newBlock == null || newBlock.getLastOperation() instanceof FlatMapOp) break;
block = newBlock;
}
}
return block.add(new LimitOp(block.getVariable(), countExpression, limit, var, delta));
}
@Nullable
private PsiLocalVariable extractCollectionAdditionVariable() {
PsiMethodCallExpression call = getSingleMethodCall();
if (!isCallOf(call, CommonClassNames.JAVA_UTIL_COLLECTION, "add")) return null;
PsiExpression[] args = call.getArgumentList().getExpressions();
if (args.length != 1 || !ExpressionUtils.isReferenceTo(args[0], myVariable)) return null;
PsiReferenceExpression qualifier = tryCast(call.getMethodExpression().getQualifierExpression(), PsiReferenceExpression.class);
if (qualifier == null) return null;
PsiLocalVariable var = tryCast(qualifier.resolve(), PsiLocalVariable.class);
if (var == null) return null;
PsiNewExpression initializer = tryCast(var.getInitializer(), PsiNewExpression.class);
if (initializer == null) return null;
PsiExpressionList argumentList = initializer.getArgumentList();
if (argumentList == null ||
argumentList.getExpressions().length != 0 ||
ControlFlowUtils.getInitializerUsageStatus(var, getMainLoop()) == ControlFlowUtils.InitializerUsageStatus.UNKNOWN) {
return null;
}
return var;
}
@NotNull
private TerminalBlock replace(Operation orig, Operation replacement) {
int idx = ArrayUtil.indexOf(myOperations, orig);
if(idx == -1) {
throw new NoSuchElementException(orig.toString());
}
Operation[] ops = myOperations.clone();
ops[idx] = replacement;
return new TerminalBlock(ops, myVariable, myStatements);
}
TerminalBlock add(Operation op) {
return new TerminalBlock(this, op, myVariable, myStatements);
}
private TerminalBlock tryExtractDistinct() {
PsiLocalVariable collectionVariable = extractCollectionAdditionVariable();
if(collectionVariable == null) return this;
for(int idx = myOperations.length-1; idx > 0; idx--) {
Operation op = myOperations[idx];
if (op instanceof FilterOp) {
FilterOp filter = (FilterOp)op;
PsiExpression condition = filter.getExpression();
if (BoolUtils.isNegation(condition)) {
if (filter.isNegated()) continue;
condition = BoolUtils.getNegated(condition);
}
else if (!filter.isNegated()) continue;
if (!(condition instanceof PsiMethodCallExpression)) continue;
PsiMethodCallExpression conditionCall = (PsiMethodCallExpression)condition;
if (!ExpressionUtils.isReferenceTo(conditionCall.getMethodExpression().getQualifierExpression(), collectionVariable) ||
!isCallOf(conditionCall, CommonClassNames.JAVA_UTIL_COLLECTION, "contains")) {
continue;
}
PsiExpression[] conditionArgs = conditionCall.getArgumentList().getExpressions();
if (conditionArgs.length == 1 && ExpressionUtils.isReferenceTo(conditionArgs[0], myVariable)) {
return replace(op, new DistinctOp(myVariable));
}
} else if (!(op instanceof LimitOp)) {
break;
}
}
return this;
}
@NotNull
Operation getLastOperation() {
return myOperations[myOperations.length-1];
}
@Nullable
<T extends Operation> T getLastOperation(Class<T> clazz) {
return tryCast(getLastOperation(), clazz);
}
@Nullable
PsiExpression getCountExpression() {
LimitOp limitOp = getLastOperation(LimitOp.class);
if (limitOp != null && limitOp.getCounterVariable() == null) {
return limitOp.getCountExpression();
}
return null;
}
/**
* Extract all possible intermediate operations
* @return the terminal block with all possible terminal operations extracted (may return this if no operations could be extracted)
*/
@NotNull
TerminalBlock extractOperations() {
return StreamEx.iterate(this, Objects::nonNull, TerminalBlock::extractOperation).reduce((a, b) -> b).orElse(this);
}
@NotNull
PsiVariable getVariable() {
return myVariable;
}
boolean hasOperations() {
return myOperations.length > 1;
}
boolean isEmpty() {
return myStatements.length == 0;
}
@NotNull
StreamEx<Operation> operations() {
return StreamEx.ofReversed(myOperations);
}
PsiLoopStatement getMainLoop() {
return ((StreamSource)myOperations[0]).getLoop();
}
/**
* @return stream of physical expressions used in intermediate operations in arbitrary order
*/
StreamEx<PsiExpression> intermediateExpressions() {
return StreamEx.of(myOperations, 1, myOperations.length).flatMap(Operation::expressions);
}
/**
* @return stream of physical expressions used in stream source and intermediate operations in arbitrary order
*/
StreamEx<PsiExpression> intermediateAndSourceExpressions() {
return operations().flatMap(Operation::expressions);
}
/**
* Converts this TerminalBlock to PsiElement (either PsiStatement or PsiCodeBlock)
*
* @param factory factory to use to create new element if necessary
* @return the PsiElement
*/
PsiElement convertToElement(PsiElementFactory factory) {
if (myStatements.length == 1) {
return myStatements[0];
}
PsiCodeBlock block = factory.createCodeBlock();
for (PsiStatement statement : myStatements) {
block.add(statement);
}
return block;
}
String generate() {
return generate(false);
}
String generate(boolean noStreamForEmpty) {
if(noStreamForEmpty && myOperations.length == 1 && myOperations[0] instanceof CollectionStream) {
return ParenthesesUtils.getText(myOperations[0].getExpression(), ParenthesesUtils.POSTFIX_PRECEDENCE);
}
return StreamEx.of(myOperations).map(Operation::createReplacement).joining();
}
@NotNull
static TerminalBlock from(StreamSource source, @NotNull PsiStatement body) {
return new TerminalBlock(null, source, source.myVariable, body).extractOperations().tryPeelLimit(false).tryExtractDistinct();
}
boolean dependsOn(PsiExpression qualifier) {
return intermediateExpressions().anyMatch(expression -> isExpressionDependsOnUpdatedCollections(expression, qualifier));
}
boolean isReferencedInOperations(PsiVariable variable) {
return intermediateAndSourceExpressions().anyMatch(expr -> VariableAccessUtils.variableIsUsed(variable, expr));
}
}