final_returns.cpp

#include <torch/csrc/jit/ir.h>
#include <torch/csrc/jit/script/final_returns.h>

namespace torch {
namespace jit {
namespace script {

struct ReturnInfo {
  bool returns_; // true - all paths through stmts_ always return
                 // false - all paths through stmts_ do not return
  List<Stmt> stmts_;
};

void checkNoReturn(const TreeRef& ref) {
  if (ref->kind() == TK_RETURN)
    throw ErrorReport(ref) << "return is not allowed from a loop.";
  for (const TreeRef& child : ref->trees()) {
    checkNoReturn(child);
  }
}

// transform stmts so that its last action is to return or report that it
// never returns.
// return_none - if true, add an implicit `return None` to the end of the block
//   this handles the case where the return is implicit at the end of the
//   function.
ReturnInfo makeReturnsFinal(
    const SourceRange& range,
    at::ArrayRef<TreeRef> stmts,
    bool return_none);
ReturnInfo makeReturnsFinal(const List<Stmt>& stmts, bool return_none) {
  return makeReturnsFinal(stmts.range(), stmts.get()->trees(), return_none);
}
ReturnInfo makeReturnsFinal(
    const SourceRange& range,
    at::ArrayRef<TreeRef> stmts,
    bool return_none) {
  std::vector<TreeRef> changed;
  changed.reserve(stmts.size());
  for (size_t i = 0; i < stmts.size(); ++i) {
    const TreeRef& stmt = stmts[i];
    switch (stmt->kind()) {
      case TK_IF: {
        auto if_stmt = If(stmt);
        auto true_final = makeReturnsFinal(if_stmt.trueBranch(), false);
        // (3) early return an if statement without an else block:
        if (true_final.returns_ && if_stmt.falseBranch().size() == 0) {
          auto rest_final =
              makeReturnsFinal(range, stmts.slice(i + 1), return_none);
          if (!rest_final.returns_) {
            throw ErrorReport(if_stmt)
                << "This if statement performs an early return, but the block of code that follows it does not return."
                << " Early returns are only allowed when the block following them also returns.";
          }
          changed.emplace_back(
              if_stmt.withNewBranches(true_final.stmts_, rest_final.stmts_));
          return {true, List<Stmt>::unsafeCreate(range, std::move(changed))};
        }

        auto false_final = makeReturnsFinal(if_stmt.falseBranch(), false);
        // (1) neither branch returns just keep processing the block
        if (!true_final.returns_ && !false_final.returns_) {
          changed.emplace_back(if_stmt);
          break;
        }
        // (2) all branches return
        if (true_final.returns_ && false_final.returns_) {
          changed.emplace_back(
              if_stmt.withNewBranches(true_final.stmts_, false_final.stmts_));
          return {true, List<Stmt>::unsafeCreate(range, std::move(changed))};
        }
        throw ErrorReport(if_stmt)
            << "This if statement contains some paths that return and some paths that do not. "
            << "If statements must either entirely return or never return.";
      } break;
      case TK_WHILE:
      case TK_FOR:
        changed.emplace_back(stmt);
        checkNoReturn(stmt);
        break;
      case TK_RETURN:
        changed.emplace_back(stmt);
        // ignore the rest the the block, it is dead.
        return {true, List<Stmt>::unsafeCreate(range, std::move(changed))};
      default:
        changed.emplace_back(stmt);
        break;
    }
  }
  if (return_none) {
    // add an implicit return none node
    changed.emplace_back(
        Return::create(range, Expr(Compound::create(TK_NONE, range, {}))));
  }
  // we reach the end of the block, no returns have happened
  // unless we just inserted a return_none implicit return.
  return {return_none, List<Stmt>::unsafeCreate(range, std::move(changed))};
}

List<Stmt> moveAllReturnsToEnd(const List<Stmt>& stmts) {
  return makeReturnsFinal(stmts, true).stmts_;
}

} // namespace script
} // namespace jit
} // namespace torch