PsiStatement (PsiStatement currentContext) { PsiElement contextParent = currentContext.getParent(); if (contextParent instanceof PsiLoopStatement) { Object mark = new Object(); PsiTreeUtil.mark(myReturnStatement, mark); currentContext = BlockUtils.expandSingleStatementToBlockStatement(currentContext); contextParent = currentContext.getParent(); myReturnStatement = (PsiReturnStatement)requireNonNull(PsiTreeUtil.releaseMark(currentContext, mark)); } if (contextParent instanceof PsiCodeBlock) { PsiElement[] tail = extractTail(currentContext, (PsiCodeBlock)contextParent); PsiStatement loopOrSwitch = PsiTreeUtil.getParentOfType(currentContext, PsiLoopStatement.class, PsiSwitchStatement.class); if (loopOrSwitch != null && PsiTreeUtil.isAncestor(myBlock, loopOrSwitch, true)) { myExitContext.register(myReplacements); String exitStatement = "if(" + myExitContext.generateExitCondition() + ") break;"; contextParent.addAfter(myFactory.createStatementFromText(exitStatement, currentContext), currentContext); currentContext = loopOrSwitch; return currentContext; } List<PsiStatement> statements = ContainerUtil.filterIsInstance(tail, PsiStatement.class); if (!statements.isEmpty()) { PsiStatement statement = statements.get(0); if (statements.size() == 1 && myExitContext.isDefaultReturn(statement)) { new CommentTracker().deleteAndRestoreComments(statement); } else { myExitContext.register(myReplacements); if (!myExitContext.isFinishCondition(statement)) { String conditionalBlock = "if(" + myExitContext.getNonExitCondition() + ") {}"; PsiIfStatement ifStatement = (PsiIfStatement)myFactory.createStatementFromText(conditionalBlock, currentContext); PsiCodeBlock ifBlock = requireNonNull(((PsiBlockStatement)requireNonNull(ifStatement.getThenBranch())).getCodeBlock()); PsiJavaToken lBrace = requireNonNull(ifBlock.getLBrace()); PsiElement tailStart = ArrayUtil.getFirstElement(tail); PsiElement tailEnd = ArrayUtil.getLastElement(tail); ifBlock.addRangeAfter(tailStart, tailEnd, lBrace); contextParent.deleteChildRange(tailStart, tailEnd); PsiElement insertedIf = contextParent.addAfter(ifStatement, currentContext); fixNonInitializedVars(insertedIf); } } } if (contextParent == myBlock) { currentContext = null; } else if (contextParent.getParent() instanceof PsiStatement) { currentContext = (PsiStatement)contextParent.getParent(); } else if (contextParent.getParent() instanceof PsiCatchSection) { currentContext = (PsiStatement)contextParent.getParent().getParent(); } else { throw new RuntimeExceptionWithAttachments("Unexpected structure: " + contextParent.getParent().getClass(), new Attachment("body.txt", myBlock.getText()), new Attachment("context.txt", contextParent.getText())); } } else if (contextParent instanceof PsiIfStatement || contextParent instanceof PsiLabeledStatement) { currentContext = (PsiStatement)contextParent; } else { throw new RuntimeExceptionWithAttachments("Unexpected structure: " + contextParent.getClass(), new Attachment("body.txt", myBlock.getText()), new Attachment("context.txt", contextParent.getText())); } return currentContext; }