Track regex literal source locations in SourceManager

This is unfortunately needed to ensure we correctly
re-lex regex literal tokens correctly, which is
needed for diagnostic logic to correctly compute
source ranges.

rdar://92469692

Author: hamishknight

include/swift/Basic/SourceManager.h

@@ -16,6 +16,7 @@
 #include "swift/Basic/FileSystem.h"
 #include "swift/Basic/SourceLoc.h"
 #include "clang/Basic/FileManager.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/Support/SourceMgr.h"
 #include <map>
@@ -55,6 +56,10 @@ class SourceManager {
   /// reusing compilation.
   llvm::DenseMap<SourceRange, SourceRange> ReplacedRanges;
 
+  /// The starting source locations of regex literals written in source. This
+  /// is an unfortunate hack needed to allow for correct re-lexing.
+  llvm::DenseSet<SourceLoc> RegexLiteralStartLocs;
+
   std::map<const char *, VirtualFile> VirtualFiles;
   mutable std::pair<const char *, const VirtualFile*> CachedVFile = {nullptr, nullptr};
 
@@ -107,6 +112,17 @@ class SourceManager {
     ReplacedRanges[Orig] = New;
   }
 
+  /// Record the starting source location of a regex literal.
+  void recordRegexLiteralStartLoc(SourceLoc loc) {
+    RegexLiteralStartLocs.insert(loc);
+  }
+
+  /// Checks whether a given source location is for the start of a regex
+  /// literal.
+  bool isRegexLiteralStart(SourceLoc loc) const {
+    return RegexLiteralStartLocs.contains(loc);
+  }
+
   /// Returns true if \c LHS is before \c RHS in the source buffer.
   bool isBeforeInBuffer(SourceLoc LHS, SourceLoc RHS) const {
     return LHS.Value.getPointer() < RHS.Value.getPointer();

lib/Parse/Lexer.cpp

@@ -2677,6 +2677,16 @@ Token Lexer::getTokenAtLocation(const SourceManager &SM, SourceLoc Loc,
   // we need to lex just the comment token.
   Lexer L(FakeLangOpts, SM, BufferID, nullptr, LexerMode::Swift,
           HashbangMode::Allowed, CRM);
+
+  if (SM.isRegexLiteralStart(Loc)) {
+    // HACK: If this was previously lexed as a regex literal, make sure we
+    // re-lex with forward slash regex literals enabled to make sure we get an
+    // accurate length. We can force EnableExperimentalStringProcessing on, as
+    // we know it must have been enabled to parse the regex in the first place.
+    FakeLangOpts.EnableExperimentalStringProcessing = true;
+    L.ForwardSlashRegexMode = LexerForwardSlashRegexMode::Always;
+  }
+
   L.restoreState(State(Loc));
   return L.peekNextToken();
 }

lib/Parse/ParseRegex.cpp

@@ -54,6 +54,8 @@ ParserResult<Expr> Parser::parseExprRegexLiteral() {
                             /*diagBaseLoc*/ getBridgedSourceLoc(Tok.getLoc()),
                             getBridgedDiagnosticEngine(&Diags));
   auto loc = consumeToken();
+  SourceMgr.recordRegexLiteralStartLoc(loc);
+
   if (hadError) {
     return makeParserResult(new (Context) ErrorExpr(loc));
   }

test/StringProcessing/Sema/regex_literal_diagnostics.swift

@@ -0,0 +1,20 @@
+// RUN: %target-typecheck-verify-swift -enable-bare-slash-regex -disable-availability-checking
+
+// REQUIRES: swift_in_compiler
+
+postfix operator ^^
+postfix func ^^ <T> (_ x: T) -> T { x }
+
+prefix operator !!
+prefix func !! <T> (_ x: T) -> T { x }
+
+// rdar://92469692 - Make sure we get a correct fix-it location here.
+func foo<T>(_ x: T, y: Int) {} // expected-note 3{{'foo(_:y:)' declared here}}
+foo(/a/) // expected-error {{missing argument for parameter 'y' in call}} {{8-8=, y: <#Int#>}}
+foo(/, /) // expected-error {{missing argument for parameter 'y' in call}} {{9-9=, y: <#Int#>}}
+foo(/a/^^) // expected-error {{missing argument for parameter 'y' in call}} {{10-10=, y: <#Int#>}}
+
+func bar<T>(x: Int, _ y: T) {} // expected-note 3{{'bar(x:_:)' declared here}}
+bar(/a/) // expected-error {{missing argument for parameter 'x' in call}} {{5-5=x: <#Int#>, }}
+bar(/, /) // expected-error {{missing argument for parameter 'x' in call}} {{5-5=x: <#Int#>, }}
+bar(!!/a/) // expected-error {{missing argument for parameter 'x' in call}} {{5-5=x: <#Int#>, }}