[Parse] Unify recovery for invalid tokens following a #if body

Previously we would only diagnose and recover for
invalid tokens following a `#if` body for the decl
and postfix expression case. Sink this logic into
`parseIfConfigRaw`, ensuring that we do this for
all `#if` cases. This requires propagating the
context we're parsing in to customize the
diagnostic.

Author: hamishknight

include/swift/AST/DiagnosticsParse.def

@@ -51,6 +51,8 @@ ERROR(extra_tokens_conditional_compilation_directive,none,
       "extra tokens following conditional compilation directive", ())
 ERROR(unexpected_rbrace_in_conditional_compilation_block,none,
       "unexpected '}' in conditional compilation block", ())
+ERROR(ifconfig_unexpectedtoken,none,
+      "unexpected tokens in '#if' body", ())
 ERROR(unexpected_if_following_else_compilation_directive,none,
       "unexpected 'if' keyword following '#else' conditional compilation "
       "directive; did you mean '#elseif'?",

include/swift/Parse/Parser.h

@@ -105,6 +105,15 @@ enum class IfConfigElementsRole {
   Skipped
 };
 
+/// Describes the context in which the '#if' is being parsed.
+enum class IfConfigContext {
+  BraceItems,
+  DeclItems,
+  SwitchStmt,
+  PostfixExpr,
+  DeclAttrs
+};
+
 /// The main class used for parsing a source file (.swift or .sil).
 ///
 /// Rather than instantiating a Parser yourself, use one of the parsing APIs
@@ -993,8 +1002,9 @@ class Parser {
   /// them however they wish. The parsing function will be provided with the
   /// location of the clause token (`#if`, `#else`, etc.), the condition,
   /// whether this is the active clause, and the role of the elements.
-  template<typename Result>
+  template <typename Result>
   Result parseIfConfigRaw(
+      IfConfigContext ifConfigContext,
       llvm::function_ref<void(SourceLoc clauseLoc, Expr *condition,
                               bool isActive, IfConfigElementsRole role)>
           parseElements,
@@ -1003,7 +1013,8 @@ class Parser {
   /// Parse a #if ... #endif directive.
   /// Delegate callback function to parse elements in the blocks.
   ParserResult<IfConfigDecl> parseIfConfig(
-    llvm::function_ref<void(SmallVectorImpl<ASTNode> &, bool)> parseElements);
+      IfConfigContext ifConfigContext,
+      llvm::function_ref<void(SmallVectorImpl<ASTNode> &, bool)> parseElements);
 
   /// Parse an #if ... #endif containing only attributes.
   ParserStatus parseIfConfigDeclAttributes(

lib/Parse/ParseDecl.cpp

@@ -6191,23 +6191,16 @@ ParserResult<Decl> Parser::parseDecl(bool IsAtStartOfLineOrPreviousHadSemi,
 
   if (Tok.is(tok::pound_if) && !ifConfigContainsOnlyAttributes()) {
     auto IfConfigResult = parseIfConfig(
-      [&](SmallVectorImpl<ASTNode> &Decls, bool IsActive) {
-        ParserStatus Status;
-        bool PreviousHadSemi = true;
-        while (Tok.isNot(tok::pound_else, tok::pound_endif, tok::pound_elseif,
-                         tok::eof)) {
-          if (Tok.is(tok::r_brace)) {
-            diagnose(Tok.getLoc(),
-                      diag::unexpected_rbrace_in_conditional_compilation_block);
-            // If we see '}', following declarations don't look like belong to
-            // the current decl context; skip them.
-            skipUntilConditionalBlockClose();
-            break;
+        IfConfigContext::DeclItems,
+        [&](SmallVectorImpl<ASTNode> &Decls, bool IsActive) {
+          ParserStatus Status;
+          bool PreviousHadSemi = true;
+          while (Tok.isNot(tok::pound_else, tok::pound_endif, tok::pound_elseif,
+                           tok::r_brace, tok::eof)) {
+            Status |= parseDeclItem(PreviousHadSemi,
+                                    [&](Decl *D) { Decls.emplace_back(D); });
           }
-          Status |= parseDeclItem(PreviousHadSemi,
-                                  [&](Decl *D) { Decls.emplace_back(D); });
-        }
-      });
+        });
     if (IfConfigResult.hasCodeCompletion() && isIDEInspectionFirstPass()) {
       consumeDecl(BeginParserPosition, CurDeclContext->isModuleScopeContext());
       return makeParserError();

lib/Parse/ParseExpr.cpp

@@ -1461,8 +1461,9 @@ Parser::parseExprPostfixSuffix(ParserResult<Expr> Result, bool isExprBasic,
       }
 
       llvm::SmallPtrSet<Expr *, 4> exprsWithBindOptional;
-      auto ICD =
-          parseIfConfig([&](SmallVectorImpl<ASTNode> &elements, bool isActive) {
+      auto ICD = parseIfConfig(
+          IfConfigContext::PostfixExpr,
+          [&](SmallVectorImpl<ASTNode> &elements, bool isActive) {
             // Although we know the '#if' body starts with period,
             // '#elseif'/'#else' bodies might start with invalid tokens.
             if (isAtStartOfPostfixExprSuffix() || Tok.is(tok::pound_if)) {
@@ -1474,13 +1475,6 @@ Parser::parseExprPostfixSuffix(ParserResult<Expr> Result, bool isExprBasic,
                 exprsWithBindOptional.insert(expr.get());
               elements.push_back(expr.get());
             }
-
-            // Don't allow any character other than the postfix expression.
-            if (!Tok.isAny(tok::pound_elseif, tok::pound_else, tok::pound_endif,
-                           tok::eof)) {
-              diagnose(Tok, diag::expr_postfix_ifconfig_unexpectedtoken);
-              skipUntilConditionalBlockClose();
-            }
           });
       if (ICD.isNull())
         break;

lib/Parse/ParseIfConfig.cpp

@@ -768,11 +768,12 @@ static Expr *findAnyLikelySimulatorEnvironmentTest(Expr *Condition) {
 
 /// Parse and populate a #if ... #endif directive.
 /// Delegate callback function to parse elements in the blocks.
-template<typename Result>
+template <typename Result>
 Result Parser::parseIfConfigRaw(
-    llvm::function_ref<void(SourceLoc clauseLoc, Expr *condition,
-                            bool isActive, IfConfigElementsRole role)>
-      parseElements,
+    IfConfigContext ifConfigContext,
+    llvm::function_ref<void(SourceLoc clauseLoc, Expr *condition, bool isActive,
+                            IfConfigElementsRole role)>
+        parseElements,
     llvm::function_ref<Result(SourceLoc endLoc, bool hadMissingEnd)> finish) {
   assert(Tok.is(tok::pound_if));
 
@@ -896,6 +897,24 @@ Result Parser::parseIfConfigRaw(
           ClauseLoc, Condition, isActive, IfConfigElementsRole::Skipped);
     }
 
+    // We ought to be at the end of the clause, diagnose if not and skip to
+    // the closing token. `#if` + `#endif` are considered stronger delimiters
+    // than `{` + `}`, so we can skip over those too.
+    if (Tok.isNot(tok::pound_elseif, tok::pound_else, tok::pound_endif,
+                  tok::eof)) {
+      if (Tok.is(tok::r_brace)) {
+        diagnose(Tok, diag::unexpected_rbrace_in_conditional_compilation_block);
+      } else if (ifConfigContext == IfConfigContext::PostfixExpr) {
+        diagnose(Tok, diag::expr_postfix_ifconfig_unexpectedtoken);
+      } else {
+        // We ought to never hit this case in practice, but fall back to a
+        // generic 'unexpected tokens' diagnostic if we weren't able to produce
+        // a better diagnostic during the parsing of the clause.
+        diagnose(Tok, diag::ifconfig_unexpectedtoken);
+      }
+      skipUntilConditionalBlockClose();
+    }
+
     // Record the clause range info in SourceFile.
     if (shouldEvaluate) {
       auto kind = isActive ? IfConfigClauseRangeInfo::ActiveClause
@@ -926,9 +945,11 @@ Result Parser::parseIfConfigRaw(
 /// Parse and populate a #if ... #endif directive.
 /// Delegate callback function to parse elements in the blocks.
 ParserResult<IfConfigDecl> Parser::parseIfConfig(
+    IfConfigContext ifConfigContext,
     llvm::function_ref<void(SmallVectorImpl<ASTNode> &, bool)> parseElements) {
   SmallVector<IfConfigClause, 4> clauses;
   return parseIfConfigRaw<ParserResult<IfConfigDecl>>(
+      ifConfigContext,
       [&](SourceLoc clauseLoc, Expr *condition, bool isActive,
           IfConfigElementsRole role) {
         SmallVector<ASTNode, 16> elements;
@@ -939,7 +960,8 @@ ParserResult<IfConfigDecl> Parser::parseIfConfig(
 
         clauses.emplace_back(
             clauseLoc, condition, Context.AllocateCopy(elements), isActive);
-      }, [&](SourceLoc endLoc, bool hadMissingEnd) {
+      },
+      [&](SourceLoc endLoc, bool hadMissingEnd) {
         auto *ICD = new (Context) IfConfigDecl(CurDeclContext,
                                                Context.AllocateCopy(clauses),
                                                endLoc, hadMissingEnd);
@@ -952,6 +974,7 @@ ParserStatus Parser::parseIfConfigDeclAttributes(
     PatternBindingInitializer *initContext) {
   ParserStatus status = makeParserSuccess();
   return parseIfConfigRaw<ParserStatus>(
+      IfConfigContext::DeclAttrs,
       [&](SourceLoc clauseLoc, Expr *condition, bool isActive,
           IfConfigElementsRole role) {
         if (isActive) {

lib/Parse/ParseStmt.cpp

@@ -369,12 +369,14 @@ ParserStatus Parser::parseBraceItems(SmallVectorImpl<ASTNode> &Entries,
     PreviousHadSemi = false;
     if (Tok.is(tok::pound_if) && !isStartOfSwiftDecl()) {
       auto IfConfigResult = parseIfConfig(
-        [&](SmallVectorImpl<ASTNode> &Elements, bool IsActive) {
-          parseBraceItems(Elements, Kind, IsActive
-                            ? BraceItemListKind::ActiveConditionalBlock
-                            : BraceItemListKind::InactiveConditionalBlock,
-                          IsFollowingGuard);
-        });
+          IfConfigContext::BraceItems,
+          [&](SmallVectorImpl<ASTNode> &Elements, bool IsActive) {
+            parseBraceItems(Elements, Kind,
+                            IsActive
+                                ? BraceItemListKind::ActiveConditionalBlock
+                                : BraceItemListKind::InactiveConditionalBlock,
+                            IsFollowingGuard);
+          });
       if (IfConfigResult.hasCodeCompletion() && isIDEInspectionFirstPass()) {
         consumeDecl(BeginParserPosition, IsTopLevel);
         return IfConfigResult;
@@ -2572,10 +2574,11 @@ Parser::parseStmtCases(SmallVectorImpl<ASTNode> &cases, bool IsActive) {
     } else if (Tok.is(tok::pound_if)) {
       // '#if' in 'case' position can enclose one or more 'case' or 'default'
       // clauses.
-      auto IfConfigResult = parseIfConfig(
-        [&](SmallVectorImpl<ASTNode> &Elements, bool IsActive) {
-          parseStmtCases(Elements, IsActive);
-        });
+      auto IfConfigResult =
+          parseIfConfig(IfConfigContext::SwitchStmt,
+                        [&](SmallVectorImpl<ASTNode> &Elements, bool IsActive) {
+                          parseStmtCases(Elements, IsActive);
+                        });
       Status |= IfConfigResult;
       if (auto ICD = IfConfigResult.getPtrOrNull()) {
         cases.emplace_back(ICD);

test/Parse/ConditionalCompilation/pound-if-inside-function-5.swift

@@ -0,0 +1,16 @@
+// RUN: %target-typecheck-verify-swift
+
+// Make sure we can recover by ignoring the '}' in the #if.
+struct S {
+  func foo() {
+#if true
+  } // expected-error {{unexpected '}' in conditional compilation block}}
+#endif
+  }
+  func bar() {
+#if false
+  } // expected-error {{unexpected '}' in conditional compilation block}}
+#endif
+  }
+  func baz() {}
+}