[ASTScope] Eliminate ASTScopeImpl::getCatchNodeBody() and centralize

The implementation of "lookup the catch node" is scattered throughout
various `ASTScopeImpl` subclasses. Remove that, and instead centralize
the code to find the innermost catch scope.

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -190,9 +190,8 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
     return parentAndWasExpanded.getPointer();
   }
 
-  const Children &getChildren() const { return storedChildren; }
-
 public:
+  const Children &getChildren() const { return storedChildren; }
   void addChild(ASTScopeImpl *child, ASTContext &);
 
 public:
@@ -235,7 +234,6 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
   ScopeKind getKind() const { return kind; }
 
   virtual NullablePtr<AbstractClosureExpr> getClosureIfClosureScope() const;
-  virtual NullablePtr<const BraceStmtScope> getAsBraceStmtScope() const;
   ASTContext &getASTContext() const;
   NullablePtr<Decl> getDeclIfAny() const;
   NullablePtr<Stmt> getStmtIfAny() const;
@@ -312,12 +310,6 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
   /// compact scope tree in the debug info.
   virtual bool ignoreInDebugInfo() const { return false; }
 
-  /// If this scope node represents a potential catch node, return body the
-  /// AST node describing the catch (a function, closure, or do...catch) and
-  /// the node of it's "body", i.e., the brace statement from which errors
-  /// thrown will be caught by that node.
-  virtual std::pair<CatchNode, const BraceStmtScope *> getCatchNodeBody() const;
-
 #pragma mark - - lookup- starting point
 private:
   static const ASTScopeImpl *findStartingScopeForLookup(SourceFile *,
@@ -896,8 +888,6 @@ class FunctionBodyScope : public ASTScopeImpl {
   Decl *getDecl() const { return decl; }
   bool ignoreInDebugInfo() const override { return true; }
 
-  std::pair<CatchNode, const BraceStmtScope *> getCatchNodeBody() const override;
-
 protected:
   bool lookupLocalsOrMembers(DeclConsumer) const override;
 
@@ -1174,8 +1164,6 @@ class ClosureParametersScope final : public ASTScopeImpl {
   NullablePtr<AbstractClosureExpr> getClosureIfClosureScope() const override {
     return closureExpr;
   }
-  std::pair<CatchNode, const BraceStmtScope *> getCatchNodeBody() const override;
-
   Expr *getExpr() const { return closureExpr; }
   bool ignoreInDebugInfo() const override { return true; }
 
@@ -1610,8 +1598,6 @@ class DoCatchStmtScope final : public AbstractStmtScope {
   void expandAScopeThatDoesNotCreateANewInsertionPoint(ScopeCreator &);
 
 public:
-  std::pair<CatchNode, const BraceStmtScope *> getCatchNodeBody() const override;
-
   Stmt *getStmt() const override { return stmt; }
 
   static bool classof(const ASTScopeImpl *scope) {
@@ -1848,9 +1834,7 @@ class BraceStmtScope final : public AbstractStmtScope {
   getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
 
   NullablePtr<AbstractClosureExpr> parentClosureIfAny() const; // public??
-  Stmt *getStmt() const override { return stmt; }
-
-  NullablePtr<const BraceStmtScope> getAsBraceStmtScope() const override;
+  BraceStmt *getStmt() const override { return stmt; }
 
 protected:
   bool lookupLocalsOrMembers(DeclConsumer) const override;

lib/AST/ASTScope.cpp

@@ -101,65 +101,10 @@ NullablePtr<AbstractClosureExpr> BraceStmtScope::parentClosureIfAny() const {
   return !getParent() ? nullptr : getParent().get()->getClosureIfClosureScope();
 }
 
-NullablePtr<const BraceStmtScope> BraceStmtScope::getAsBraceStmtScope() const {
-  return this;
-}
-
 NullablePtr<AbstractClosureExpr> ASTScopeImpl::getClosureIfClosureScope() const {
   return nullptr;
 }
 
-NullablePtr<const BraceStmtScope> ASTScopeImpl::getAsBraceStmtScope() const {
-  return nullptr;
-}
-
-std::pair<CatchNode, const BraceStmtScope *>
-ASTScopeImpl::getCatchNodeBody() const {
-  return { nullptr, nullptr };
-}
-
-std::pair<CatchNode, const BraceStmtScope *>
-ClosureParametersScope::getCatchNodeBody() const {
-  const BraceStmtScope *body = nullptr;
-  const auto &children = getChildren();
-  if (!children.empty()) {
-    body = children[0]->getAsBraceStmtScope().getPtrOrNull();
-    assert(body && "Not a brace statement?");
-  }
-  return { const_cast<AbstractClosureExpr *>(closureExpr), body };
-}
-
-std::pair<CatchNode, const BraceStmtScope *>
-FunctionBodyScope::getCatchNodeBody() const {
-  const BraceStmtScope *body = nullptr;
-  const auto &children = getChildren();
-  if (!children.empty()) {
-    body = children[0]->getAsBraceStmtScope().getPtrOrNull();
-    assert(body && "Not a brace statement?");
-  }
-  return { const_cast<AbstractFunctionDecl *>(decl), body };
-}
-
-/// Determine whether this is an empty brace statement, which doesn't have a
-/// node associated with it.
-static bool isEmptyBraceStmt(Stmt *stmt) {
-  if (auto braceStmt = dyn_cast_or_null<BraceStmt>(stmt))
-    return braceStmt->empty();
-
-  return false;
-}
-
-std::pair<CatchNode, const BraceStmtScope *>
-DoCatchStmtScope::getCatchNodeBody() const {
-  const BraceStmtScope *body = nullptr;
-  const auto &children = getChildren();
-  if (!children.empty() && !isEmptyBraceStmt(stmt->getBody())) {
-    body = children[0]->getAsBraceStmtScope().getPtrOrNull();
-    assert(body && "Not a brace statement?");
-  }
-  return { const_cast<DoCatchStmt *>(stmt), body };
-}
-
 std::string ASTScopeImpl::getClassName() const {
   // GenericTypeOrExtensionScope provides a custom implementation that deals
   // with declaration names and "portions".

lib/AST/ASTScopeLookup.cpp

@@ -666,6 +666,43 @@ void ASTScopeImpl::lookupEnclosingMacroScope(
   } while ((scope = scope->getParent().getPtrOrNull()));
 }
 
+static std::pair<CatchNode, const BraceStmtScope *>
+getCatchNodeBody(const ASTScopeImpl *scope, CatchNode node) {
+  const auto &children = scope->getChildren();
+  if (children.empty())
+    return { CatchNode(), nullptr };
+
+  auto stmt = dyn_cast<BraceStmtScope>(children[0]);
+  if (!stmt || stmt->getStmt()->empty())
+    return { CatchNode(), nullptr };
+
+  return std::make_pair(node, stmt);
+}
+
+/// Retrieve the catch node associated with this scope, if any.
+static std::pair<CatchNode, const BraceStmtScope *>
+getCatchNode(const ASTScopeImpl *scope) {
+  // Closures introduce a catch scope for errors initiated in their body.
+  if (auto closureParams = dyn_cast<ClosureParametersScope>(scope)) {
+    return getCatchNodeBody(
+        scope, const_cast<AbstractClosureExpr *>(closureParams->closureExpr));
+  }
+
+  // Functions introduce a catch scope for errors initiated in their body.
+  if (auto function = dyn_cast<FunctionBodyScope>(scope)) {
+    return getCatchNodeBody(
+        scope, const_cast<AbstractFunctionDecl *>(function->decl));
+  }
+
+  // Do..catch blocks introduce a catch scope for errors initiated in the `do`
+  // body.
+  if (auto doCatch = dyn_cast<DoCatchStmtScope>(scope)) {
+    return getCatchNodeBody(scope, const_cast<DoCatchStmt *>(doCatch->stmt));
+  }
+
+  return { CatchNode(), nullptr };
+}
+
 CatchNode ASTScopeImpl::lookupCatchNode(ModuleDecl *module, SourceLoc loc) {
   auto sourceFile = module->getSourceFileContainingLocation(loc);
   if (!sourceFile)
@@ -682,12 +719,12 @@ CatchNode ASTScopeImpl::lookupCatchNode(ModuleDecl *module, SourceLoc loc) {
   for (auto scope = innermost; scope; scope = scope->getParent().getPtrOrNull()) {
     // If we are at a catch node and in the body of the region from which that
     // node catches thrown errors, we have our result.
-    auto caught = scope->getCatchNodeBody();
+    auto caught = getCatchNode(scope);
     if (caught.first && caught.second == innerBodyScope) {
       return caught.first;
     }
 
-    innerBodyScope = scope->getAsBraceStmtScope().getPtrOrNull();
+    innerBodyScope = dyn_cast<BraceStmtScope>(scope);
   }
 
   return nullptr;