[Parse] Avoid skipping bodies with #sourceLocation

`#sourceLocation` requires setting and restoring
state in the parser, so we cannot skip any bodies
containing `#sourceLocation` tokens. Member skipping
was already doing this, extend it to function
body skipping too.

Author: hamishknight

lib/Parse/ParseDecl.cpp

@@ -5620,11 +5620,13 @@ static void diagnoseOperatorFixityAttributes(Parser &P,
 
 static unsigned skipUntilMatchingRBrace(Parser &P,
                                         bool &HasPoundDirective,
+                                        bool &HasPoundSourceLocation,
                                         bool &HasOperatorDeclarations,
                                         bool &HasNestedClassDeclarations,
                                         bool &HasNestedTypeDeclarations,
                                         bool &HasPotentialRegexLiteral) {
   HasPoundDirective = false;
+  HasPoundSourceLocation = false;
   HasOperatorDeclarations = false;
   HasNestedClassDeclarations = false;
   HasNestedTypeDeclarations = false;
@@ -5646,8 +5648,11 @@ static unsigned skipUntilMatchingRBrace(Parser &P,
 
     HasNestedClassDeclarations |= P.Tok.is(tok::kw_class);
 
-    HasPoundDirective |= P.Tok.isAny(tok::pound_sourceLocation, tok::pound_line,
-      tok::pound_if, tok::pound_else, tok::pound_endif, tok::pound_elseif);
+    HasPoundSourceLocation |= P.Tok.isAny(tok::pound_sourceLocation,
+                                          tok::pound_line);
+    HasPoundDirective |= HasPoundSourceLocation;
+    HasPoundDirective |= P.Tok.isAny(tok::pound_if, tok::pound_else,
+                                     tok::pound_endif, tok::pound_elseif);
 
     HasNestedTypeDeclarations |= P.Tok.isAny(tok::kw_class, tok::kw_struct,
                                              tok::kw_enum, tok::kw_typealias,
@@ -7085,16 +7090,22 @@ bool Parser::canDelayMemberDeclParsing(bool &HasOperatorDeclarations,
   // we can't lazily parse.
   CancellableBacktrackingScope BackTrack(*this);
   bool HasPoundDirective;
+  bool HasPoundSourceLocation;
   bool HasNestedTypeDeclarations;
   bool HasPotentialRegexLiteral;
   skipUntilMatchingRBrace(*this,
                           HasPoundDirective,
+                          HasPoundSourceLocation,
                           HasOperatorDeclarations,
                           HasNestedClassDeclarations,
                           HasNestedTypeDeclarations,
                           HasPotentialRegexLiteral);
-  if (!HasPoundDirective && !HasPotentialRegexLiteral)
+  if (!HasPoundDirective && !HasPotentialRegexLiteral) {
+    // If we didn't see any pound directive, we must not have seen
+    // #sourceLocation either.
+    ASSERT(!HasPoundSourceLocation);
     BackTrack.cancelBacktrack();
+  }
   return !BackTrack.willBacktrack();
 }
 
@@ -7690,17 +7701,20 @@ bool Parser::canDelayFunctionBodyParsing(bool &HasNestedTypeDeclarations) {
     return false;
 
   // Skip until the matching right curly bracket; If it has a potential regex
-  // literal, we can't skip. We don't care others, so just ignore them;
+  // literal, we can't skip. If there's a `#sourceLocation`, we also can't skip
+  // since we rely on setting and restoring state in the parser. Other cases we
+  // can handle.
   CancellableBacktrackingScope BackTrack(*this);
   consumeToken(tok::l_brace);
   bool HasPoundDirectives;
+  bool HasPoundSourceLocation;
   bool HasOperatorDeclarations;
   bool HasNestedClassDeclarations;
   bool HasPotentialRegexLiteral;
-  skipUntilMatchingRBrace(*this, HasPoundDirectives, HasOperatorDeclarations,
-                          HasNestedClassDeclarations, HasNestedTypeDeclarations,
-                          HasPotentialRegexLiteral);
-  if (HasPotentialRegexLiteral)
+  skipUntilMatchingRBrace(*this, HasPoundDirectives, HasPoundSourceLocation,
+                          HasOperatorDeclarations, HasNestedClassDeclarations,
+                          HasNestedTypeDeclarations, HasPotentialRegexLiteral);
+  if (HasPoundSourceLocation || HasPotentialRegexLiteral)
     return false;
 
   BackTrack.cancelBacktrack();

test/Parse/function_skipping_source_location.swift

@@ -0,0 +1,53 @@
+// Note we use '--implicit-check-not' to ensure we exhaustively match errors.
+// RUN: not %target-swift-frontend -parse -experimental-skip-all-function-bodies -diagnostic-style=llvm %s 2>&1 | %FileCheck --implicit-check-not='error:' %s
+
+// We cannot skip function bodies with #sourceLocation.
+func foo() {
+  // CHECK: [[@LINE+1]]:4: error: expected expression
+  [;
+  #sourceLocation(file: "A", line: 1)
+  [;
+  // CHECK: A:1:4: error: expected expression
+}
+
+func bar() {
+  // CHECK: A:7:4: error: expected expression
+  [;
+#sourceLocation()
+  [;
+  // CHECK: [[@LINE-1]]:4: error: expected expression
+}
+
+// This function body is skipped.
+func baz() {
+  [;
+}
+
+// This member list and function are not skipped.
+struct S {
+  func qux() {
+    // CHECK: [[@LINE+1]]:6: error: expected expression
+    [;
+    #sourceLocation(file: "B", line: 1)
+    [;
+    // CHECK: B:1:6: error: expected expression
+  }
+  func ;
+  // CHECK: B:4:8: error: expected identifier
+}
+
+// This member list is also not skipped.
+struct R {
+  // CHECK: B:11:8: error: expected identifier
+  func ;
+
+#sourceLocation()
+
+  func ;
+  // CHECK: [[@LINE-1]]:8: error: expected identifier
+}
+
+// This member list is skipped.
+struct Q {
+    [;
+}

test/Parse/issue-74561.swift

@@ -2,7 +2,7 @@
 
 // Make sure we can parse with and without skipping.
 // RUN: %target-typecheck-verify-swift
-// RUN: not %target-swift-frontend -emit-module -experimental-skip-non-inlinable-function-bodies -module-name Mod -emit-module-path %t/Mod.swiftmodule -diagnostic-style=llvm %s 2>&1 | %FileCheck %s
+// RUN: %target-swift-frontend -verify -emit-module -experimental-skip-non-inlinable-function-bodies -module-name Mod -emit-module-path %t/Mod.swiftmodule %s
 
 // https://github.com/swiftlang/swift/issues/74561
 // Make sure we can parse this.
@@ -11,10 +11,8 @@ public func foo(_ param: Int) {
 #sourceLocation()
 }
 
-// FIXME: This should parse correctly.
 #sourceLocation(file: "B", line: 3)
 @inlinable
 public func bar(_ param: Int) {
 #sourceLocation()
 }
-// CHECK: B:6:1: error: parameterless closing #sourceLocation() directive without prior opening #sourceLocation(file:,line:) directive