[ASTGen] Fix end location of TopLevelCodeDecl

TopLevelCodeDecl::getEndLoc() should return the token location of the
last token. If the last token is a string literal or a regex literal,
it should be the location of the literal token.

Author: rintaro

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -110,6 +110,21 @@ struct ASTGenVisitor {
       split: Self.splitCodeBlockItemIfConfig
     ) { element in
       let loc = self.generateSourceLoc(element)
+
+      func endLoc() -> BridgedSourceLoc {
+        if let endTok = element.lastToken(viewMode: .sourceAccurate) {
+          switch endTok.parent?.kind {
+          case .stringLiteralExpr, .regexLiteralExpr:
+            // string/regex literal are single token in AST.
+            return self.generateSourceLoc(endTok.parent)
+          default:
+            return self.generateSourceLoc(endTok)
+          }
+        } else {
+          return loc
+        }
+      }
+
       let swiftASTNodes = generate(codeBlockItem: element)
       switch swiftASTNodes {
       case .decl(let d):
@@ -120,7 +135,7 @@ struct ASTGenVisitor {
           declContext: self.declContext,
           startLoc: loc,
           stmt: s,
-          endLoc: loc
+          endLoc: endLoc()
         )
         out.append(topLevelDecl.asDecl)
       case .expr(let e):
@@ -129,7 +144,7 @@ struct ASTGenVisitor {
           declContext: self.declContext,
           startLoc: loc,
           expr: e,
-          endLoc: loc
+          endLoc: endLoc()
         )
         out.append(topLevelDecl.asDecl)
       }

test/ASTGen/top_level.swift

@@ -0,0 +1,26 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend %s -dump-parse -disable-availability-checking -enable-experimental-move-only -enable-experimental-feature ParserASTGen > %t/astgen.ast.raw
+// RUN: %target-swift-frontend %s -dump-parse -disable-availability-checking -enable-experimental-move-only > %t/cpp-parser.ast.raw
+
+// Filter out any addresses in the dump, since they can differ.
+// RUN: sed -E 's#0x[0-9a-fA-F]+##g' %t/cpp-parser.ast.raw > %t/cpp-parser.ast
+// RUN: sed -E 's#0x[0-9a-fA-F]+##g' %t/astgen.ast.raw > %t/astgen.ast
+
+// RUN: %diff -u %t/astgen.ast %t/cpp-parser.ast
+
+// -enable-experimental-feature requires an asserts build
+// REQUIRES: asserts
+// rdar://116686158
+// UNSUPPORTED: asan
+
+_ = ##"foo bar"##
+_ = ##/foo bar/##
+
+if let first = [1,2,3].first {
+  foo(x: first)
+}
+
+func foo(x: Int) {}
+
+// FIXME: Top-level pattern binding decl must be enclosed with TopLevelCodeDecl
+// let a = 42