private static PsiStatement @Nullable [] replaceAllReturnWithYield(PsiStatement[] statements) {
    PsiStatement[] result = ArrayUtil.copyOf(statements);
    for (int i = 0; i < statements.length; i++) {
      PsiStatement statement = result[i];
      if (statement instanceof PsiReturnStatement returnStatement) {
        PsiExpression returnValue = returnStatement.getReturnValue();
        if (returnValue == null || PsiTreeUtil.hasErrorElements(returnValue) || returnValue.getType() == null) {
          return null;
        }
        result[i] = createYieldStatement(returnValue);
        continue;
      }
      PsiStatement copy = (PsiStatement)statement.copy();
      result[i] = copy;
      Collection<PsiReturnStatement> returnStatements = PsiTreeUtil.findChildrenOfType(copy, PsiReturnStatement.class);
      for (PsiReturnStatement returnStatement : returnStatements) {
        PsiExpression returnValue = returnStatement.getReturnValue();
        if (returnValue == null || PsiTreeUtil.hasErrorElements(returnValue) || returnValue.getType() == null) {
          return null;
        }
        returnStatement.replace(createYieldStatement(returnValue));
      }
    }
    return result;
  }
private static PsiStatement @Nullable [] replaceAllReturnWithYield(PsiStatement[] statements)