public static void foldWhenExpressionWithReturns(JetWhenExpression whenExpression) {
Project project = whenExpression.getProject();
assert !whenExpression.getEntries().isEmpty() : FOLD_WITHOUT_CHECK;
JetReturnExpression newReturnExpression = JetPsiFactory.createReturn(project, whenExpression);
JetWhenExpression newWhenExpression =
(JetWhenExpression) newReturnExpression.getReturnedExpression();
assertNotNull(newWhenExpression);
//noinspection ConstantConditions
for (JetWhenEntry entry : newWhenExpression.getEntries()) {
JetReturnExpression currReturn = getFoldableBranchedReturn(entry.getExpression());
assertNotNull(currReturn);
JetExpression currExpr = currReturn.getReturnedExpression();
assertNotNull(currExpr);
//noinspection ConstantConditions
currReturn.replace(currExpr);
}
whenExpression.replace(newReturnExpression);
}
public static boolean checkWhenExpressionHasSingleElse(
@NotNull JetWhenExpression whenExpression) {
int elseCount = 0;
for (JetWhenEntry entry : whenExpression.getEntries()) {
if (entry.isElse()) {
elseCount++;
}
}
return (elseCount == 1);
}
@SuppressWarnings("ConstantConditions")
public static void foldWhenExpressionWithAssignments(JetWhenExpression whenExpression) {
Project project = whenExpression.getProject();
assert !whenExpression.getEntries().isEmpty() : FOLD_WITHOUT_CHECK;
JetBinaryExpression firstAssignment =
getFoldableBranchedAssignment(whenExpression.getEntries().get(0).getExpression());
assertNotNull(firstAssignment);
String op = firstAssignment.getOperationReference().getText();
JetSimpleNameExpression lhs = (JetSimpleNameExpression) firstAssignment.getLeft();
JetBinaryExpression assignment =
JetPsiFactory.createBinaryExpression(project, lhs, op, whenExpression);
JetWhenExpression newWhenExpression = (JetWhenExpression) assignment.getRight();
assertNotNull(newWhenExpression);
for (JetWhenEntry entry : newWhenExpression.getEntries()) {
JetBinaryExpression currAssignment = getFoldableBranchedAssignment(entry.getExpression());
assertNotNull(currAssignment);
JetExpression currRhs = currAssignment.getRight();
assertNotNull(currRhs);
currAssignment.replace(currRhs);
}
whenExpression.replace(assignment);
}
private static boolean checkFoldableWhenExpressionWithReturns(JetWhenExpression whenExpression) {
if (!JetPsiUtil.checkWhenExpressionHasSingleElse(whenExpression)) return false;
List<JetWhenEntry> entries = whenExpression.getEntries();
if (entries.isEmpty()) return false;
for (JetWhenEntry entry : entries) {
if (getFoldableBranchedReturn(entry.getExpression()) == null) return false;
}
return true;
}
@Nullable
public static <T extends PsiElement> T getDirectParentOfTypeForBlock(
@NotNull JetBlockExpression block, @NotNull Class<T> aClass) {
T parent = PsiTreeUtil.getParentOfType(block, aClass);
if (parent instanceof JetIfExpression) {
JetIfExpression ifExpression = (JetIfExpression) parent;
if (ifExpression.getElse() == block || ifExpression.getThen() == block) {
return parent;
}
}
if (parent instanceof JetWhenExpression) {
JetWhenExpression whenExpression = (JetWhenExpression) parent;
for (JetWhenEntry whenEntry : whenExpression.getEntries()) {
if (whenEntry.getExpression() == block) {
return parent;
}
}
}
if (parent instanceof JetFunctionLiteral) {
JetFunctionLiteral functionLiteral = (JetFunctionLiteral) parent;
if (functionLiteral.getBodyExpression() == block) {
return parent;
}
}
if (parent instanceof JetTryExpression) {
JetTryExpression tryExpression = (JetTryExpression) parent;
if (tryExpression.getTryBlock() == block) {
return parent;
}
for (JetCatchClause clause : tryExpression.getCatchClauses()) {
if (clause.getCatchBody() == block) {
return parent;
}
}
}
return null;
}
private static boolean checkFoldableWhenExpressionWithAssignments(
JetWhenExpression whenExpression) {
if (!JetPsiUtil.checkWhenExpressionHasSingleElse(whenExpression)) return false;
List<JetWhenEntry> entries = whenExpression.getEntries();
if (entries.isEmpty()) return false;
List<JetBinaryExpression> assignments = new ArrayList<JetBinaryExpression>();
for (JetWhenEntry entry : entries) {
JetBinaryExpression assignment = getFoldableBranchedAssignment(entry.getExpression());
if (assignment == null) return false;
assignments.add(assignment);
}
assert !assignments.isEmpty();
JetBinaryExpression firstAssignment = assignments.get(0);
for (JetBinaryExpression assignment : assignments) {
if (!checkAssignmentsMatch(assignment, firstAssignment)) return false;
}
return true;
}