TerminalBlock (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.copyOf(myStatements, count); tb = new TerminalBlock(myOperations, myVariable, Arrays.copyOfRange(myStatements, count, myStatements.length)).extractFilter(); } PsiStatement sourceStatement = getStreamSourceStatement(); if (tb == null || (sourceStatement instanceof PsiLoopStatement && !ControlFlowUtils.statementBreaksLoop(tb.getSingleStatement(), (PsiLoopStatement)sourceStatement))) return this; FilterOp filter = tb.getLastOperation(FilterOp.class); if (filter == null) return this; PsiBinaryExpression binOp = tryCast(PsiUtil.skipParenthesizedExprDown(filter.getExpression()), PsiBinaryExpression.class); if (!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 ">" -> delta = 1; case "<" -> { delta = 1; flipped = true; } case "<=" -> flipped = true; 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.isSafelyRecomputableExpression(limit) || VariableAccessUtils.variableIsUsed(myVariable, limit)) return this; PsiType type = limit.getType(); if(!PsiTypes.intType().equals(type) && !PsiTypes.longType().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)); }