private static PsiStatement wrapWithBlockStatementIfNeeded(@NotNull PsiStatement statement, @NotNull PsiElementFactory factory) {
    PsiElement parent = statement.getParent();
    if (parent instanceof PsiLoopStatement && ((PsiLoopStatement)parent).getBody() == statement ||
        parent instanceof PsiIfStatement &&
        (((PsiIfStatement)parent).getThenBranch() == statement || ((PsiIfStatement)parent).getElseBranch() == statement)) {
      PsiBlockStatement blockStatement = (PsiBlockStatement)factory.createStatementFromText("{}", statement);
      blockStatement.getCodeBlock().add(statement);
      blockStatement = (PsiBlockStatement)statement.replace(blockStatement);
      return blockStatement.getCodeBlock().getStatements()[0];
    }
    if (parent instanceof PsiForStatement) {
      if (((PsiForStatement)parent).getInitialization() == statement || ((PsiForStatement)parent).getUpdate() == statement) {
        return wrapWithBlockStatementIfNeeded((PsiForStatement)parent, factory);
      }
    }
    return statement;
  }
private static PsiStatement wrapWithBlockStatementIfNeeded(@NotNull PsiStatement statement, @NotNull PsiElementFactory factory)