public static void foldIfExpressionWithReturns(JetIfExpression ifExpression) {
    Project project = ifExpression.getProject();

    JetReturnExpression newReturnExpression = JetPsiFactory.createReturn(project, ifExpression);
    JetIfExpression newIfExpression = (JetIfExpression) newReturnExpression.getReturnedExpression();

    assertNotNull(newIfExpression);

    //noinspection ConstantConditions
    JetReturnExpression thenReturn = getFoldableBranchedReturn(newIfExpression.getThen());
    JetReturnExpression elseReturn = getFoldableBranchedReturn(newIfExpression.getElse());

    assertNotNull(thenReturn);
    assertNotNull(elseReturn);

    JetExpression thenExpr = thenReturn.getReturnedExpression();
    JetExpression elseExpr = elseReturn.getReturnedExpression();

    assertNotNull(thenExpr);
    assertNotNull(elseExpr);

    //noinspection ConstantConditions
    thenReturn.replace(thenExpr);
    //noinspection ConstantConditions
    elseReturn.replace(elseExpr);

    ifExpression.replace(newReturnExpression);
  }
  public static void foldIfExpressionWithAsymmetricReturns(JetIfExpression ifExpression) {
    Project project = ifExpression.getProject();

    JetExpression condition = ifExpression.getCondition();
    JetExpression thenRoot = ifExpression.getThen();
    JetExpression elseRoot =
        (JetExpression) JetPsiUtil.skipTrailingWhitespacesAndComments(ifExpression);

    assertNotNull(condition);
    assertNotNull(thenRoot);
    assertNotNull(elseRoot);

    //noinspection ConstantConditions
    JetIfExpression newIfExpression =
        JetPsiFactory.createIf(project, condition, thenRoot, elseRoot);
    JetReturnExpression newReturnExpression = JetPsiFactory.createReturn(project, newIfExpression);

    newIfExpression = (JetIfExpression) newReturnExpression.getReturnedExpression();

    assertNotNull(newIfExpression);

    //noinspection ConstantConditions
    JetReturnExpression thenReturn = getFoldableBranchedReturn(newIfExpression.getThen());
    JetReturnExpression elseReturn = getFoldableBranchedReturn(newIfExpression.getElse());

    assertNotNull(thenReturn);
    assertNotNull(elseReturn);

    JetExpression thenExpr = thenReturn.getReturnedExpression();
    JetExpression elseExpr = elseReturn.getReturnedExpression();

    assertNotNull(thenExpr);
    assertNotNull(elseExpr);

    //noinspection ConstantConditions
    thenReturn.replace(thenExpr);
    //noinspection ConstantConditions
    elseReturn.replace(elseExpr);

    elseRoot.delete();
    ifExpression.replace(newReturnExpression);
  }
  public static void foldIfExpressionWithAssignments(JetIfExpression ifExpression) {
    Project project = ifExpression.getProject();

    JetBinaryExpression thenAssignment = getFoldableBranchedAssignment(ifExpression.getThen());

    assertNotNull(thenAssignment);

    String op = thenAssignment.getOperationReference().getText();
    JetSimpleNameExpression lhs = (JetSimpleNameExpression) thenAssignment.getLeft();

    JetBinaryExpression assignment =
        JetPsiFactory.createBinaryExpression(project, lhs, op, ifExpression);
    JetIfExpression newIfExpression = (JetIfExpression) assignment.getRight();

    assertNotNull(newIfExpression);

    //noinspection ConstantConditions
    thenAssignment = getFoldableBranchedAssignment(newIfExpression.getThen());
    JetBinaryExpression elseAssignment = getFoldableBranchedAssignment(newIfExpression.getElse());

    assertNotNull(thenAssignment);
    assertNotNull(elseAssignment);

    JetExpression thenRhs = thenAssignment.getRight();
    JetExpression elseRhs = elseAssignment.getRight();

    assertNotNull(thenRhs);
    assertNotNull(elseRhs);

    //noinspection ConstantConditions
    thenAssignment.replace(thenRhs);
    //noinspection ConstantConditions
    elseAssignment.replace(elseRhs);

    ifExpression.replace(assignment);
  }