[Sema] Clean up diagnoseStmtAvailability a bit

The walker here is really doing 2 completely
separate things, split it into 2 walkers. For
ExprAvailabilityWalker, we just want to recursively
continue walking looking for any sub-expressions
to diagnose. For StmtAvailabilityWalker we only
want to walk the top-level statement, and are
only interested in any TypeReprs it has. This lets
us get rid of the special handling for BraceStmt.

Author: hamishknight

lib/Sema/TypeCheckAvailability.cpp

@@ -3790,13 +3790,22 @@ class ExprAvailabilityWalker : public ASTWalker {
   }
 
   PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
-
-    // We end up here when checking the output of the result builder transform,
-    // which includes closures that are not "separately typechecked" and yet
-    // contain statements and declarations. We need to walk them recursively,
-    // since these availability for these statements is not diagnosed from
-    // typeCheckStmt() as usual.
-    diagnoseStmtAvailability(S, Where.getDeclContext(), /*walkRecursively=*/true);
+    // We need to recursively call diagnoseExprAvailability for any
+    // sub-expressions in the statement since the availability context may
+    // differ, e.g for things like `guard #available(...)`.
+    class StmtRecurseWalker : public BaseDiagnosticWalker {
+      DeclContext *DC;
+
+    public:
+      StmtRecurseWalker(DeclContext *DC) : DC(DC) {}
+
+      PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
+        diagnoseExprAvailability(E, DC);
+        return Action::SkipNode(E);
+      }
+    };
+    StmtRecurseWalker W(Where.getDeclContext());
+    S->walk(W);
     return Action::SkipNode(S);
   }
 
@@ -4396,26 +4405,29 @@ void swift::diagnoseExprAvailability(const Expr *E, DeclContext *DC) {
 namespace {
 
 class StmtAvailabilityWalker : public BaseDiagnosticWalker {
+  const Stmt *TopLevelStmt;
   DeclContext *DC;
-  bool WalkRecursively;
 
 public:
-  explicit StmtAvailabilityWalker(DeclContext *dc, bool walkRecursively)
-    : DC(dc), WalkRecursively(walkRecursively) {}
+  explicit StmtAvailabilityWalker(const Stmt *S, DeclContext *dc)
+    : TopLevelStmt(S), DC(dc) {}
 
   PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
-    if (!WalkRecursively && isa<BraceStmt>(S))
-      return Action::SkipNode(S);
-
-    return Action::Continue(S);
+    // `diagnoseStmtAvailability` is called for every statement, so we don't
+    // want to walk into any nested statements.
+    return Action::VisitNodeIf(S == TopLevelStmt, S);
   }
 
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-    if (WalkRecursively)
-      diagnoseExprAvailability(E, DC);
+    // Handled by ExprAvailabilityWalker.
     return Action::SkipNode(E);
   }
 
+  PreWalkAction walkToDeclPre(Decl *D) override {
+    // Handled by DeclAvailabilityChecker.
+    return Action::SkipNode();
+  }
+
   PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
     auto where = ExportContext::forFunctionBody(DC, T->getStartLoc());
     diagnoseTypeReprAvailability(T, where);
@@ -4434,13 +4446,8 @@ class StmtAvailabilityWalker : public BaseDiagnosticWalker {
 };
 }
 
-void swift::diagnoseStmtAvailability(const Stmt *S, DeclContext *DC,
-                                     bool walkRecursively) {
-  // We'll visit the individual statements when we check them.
-  if (!walkRecursively && isa<BraceStmt>(S))
-    return;
-
-  StmtAvailabilityWalker walker(DC, walkRecursively);
+void swift::diagnoseStmtAvailability(const Stmt *S, DeclContext *DC) {
+  StmtAvailabilityWalker walker(S, DC);
   const_cast<Stmt*>(S)->walk(walker);
 }
 

lib/Sema/TypeCheckAvailability.h

@@ -211,12 +211,8 @@ bool isExported(const Decl *D);
 void diagnoseExprAvailability(const Expr *E, DeclContext *DC);
 
 /// Diagnose uses of unavailable declarations in statements (via patterns, etc)
-/// but not expressions, unless \p walkRecursively was specified.
-///
-/// \param walkRecursively Whether nested statements and expressions should
-/// be visited, too.
-void diagnoseStmtAvailability(const Stmt *S, DeclContext *DC,
-                              bool walkRecursively=false);
+/// but not expressions.
+void diagnoseStmtAvailability(const Stmt *S, DeclContext *DC);
 
 /// Checks both a TypeRepr and a Type, but avoids emitting duplicate
 /// diagnostics by only checking the Type if the TypeRepr succeeded.