Merge pull request #74559 from hamishknight/pound-if-skip

[Parse] Handle `#if` in brace skipping logic

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

@@ -5659,10 +5659,11 @@ static unsigned skipUntilMatchingRBrace(Parser &P,
   HasPotentialRegexLiteral = false;
 
   unsigned OpenBraces = 1;
+  unsigned OpenPoundIf = 0;
 
   bool LastTokenWasFunc = false;
 
-  while (OpenBraces != 0 && P.Tok.isNot(tok::eof)) {
+  while ((OpenBraces != 0 || OpenPoundIf != 0) && P.Tok.isNot(tok::eof)) {
     // Detect 'func' followed by an operator identifier.
     if (LastTokenWasFunc) {
       LastTokenWasFunc = false;
@@ -5693,6 +5694,24 @@ static unsigned skipUntilMatchingRBrace(Parser &P,
       return OpenBraces;
     }
 
+    // Match opening `#if` with closing `#endif` to match what the parser does,
+    // `#if` + `#endif` are considered stronger delimiters than `{` + `}`.
+    if (P.consumeIf(tok::pound_if)) {
+      OpenPoundIf += 1;
+      continue;
+    }
+    if (OpenPoundIf != 0) {
+      // We're in a `#if`, check to see if we've reached the end.
+      if (P.consumeIf(tok::pound_endif)) {
+        OpenPoundIf -= 1;
+        continue;
+      }
+      // Consume the next token and continue iterating. We can swallow any
+      // amount of '{' and '}' while in the `#if`.
+      P.consumeToken();
+      continue;
+    }
+
     if (P.consumeIf(tok::l_brace)) {
       ++OpenBraces;
       continue;
@@ -6192,23 +6211,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

@@ -1479,8 +1479,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)) {
@@ -1492,13 +1493,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

@@ -769,11 +769,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));
 
@@ -897,6 +898,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
@@ -927,9 +946,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;
@@ -940,7 +961,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);
@@ -953,6 +975,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

@@ -370,12 +370,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;
@@ -2573,10 +2575,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() {}
+}

test/Parse/rdar129195380.swift

@@ -0,0 +1,36 @@
+// RUN: not %target-swift-frontend -experimental-skip-all-function-bodies -dump-parse %s | %FileCheck %s
+
+// rdar://129195380 - Make sure the skipping logic can handle #if.
+struct S {
+  // CHECK: func_decl{{.*}}:[[@LINE+1]]:3 - line:[[@LINE+11]]:3{{.*}}"foo()"
+  func foo() {
+#if true
+  }
+#if true
+  func bar() {
+#else
+  }
+#endif
+  }
+#endif
+  }
+  // CHECK: func_decl{{.*}}:[[@LINE+1]]:3 - line:[[@LINE+1]]:15{{.*}}"baz()"
+  func baz() {}
+}
+
+// The '#if' is unterminated here, so swallows the rest of the file.
+// CHECK: struct_decl{{.*}}:[[@LINE+1]]:1 - line:[[@LINE+14]]:14{{.*}}"R"
+struct R {
+#if false
+  }
+#if true
+  }
+#endif
+  }
+#else
+  }
+  // CHECK-NOT: qux
+  func qux() {}
+}
+
+func flim() {}

test/SourceKit/CursorInfo/rdar129195380.swift

@@ -0,0 +1,13 @@
+// rdar://129195380 - Make sure we correctly handle '#if' when skipping function
+// bodies.
+class C {
+  func test1() {
+  #if FOOBAR
+// RUN: %sourcekitd-test -req=cursor -pos=%(line + 2):5 %s -- %s -DFOOBAR
+// RUN: %sourcekitd-test -req=cursor -pos=%(line + 1):5 %s -- %s
+    abc
+  }
+
+  func test2() {
+  }
+}