[Sema] Add fix-it to import RegexBuilder

For code such as the following:

```
let r = Regex {
  /abc/
}
```

If RegexBuilder has not been imported, emit a
specialized diagnostic and fix-it to add
`import RegexBuilder` to the file.

Unfortunately we're currently prevented from
emitting the specialized diagnostic in cases where
the builder contains references to RegexBuilder
types, such as:

```
let r = Regex {
  Capture {
    /abc/
  }
}
```

This is due to the fact that we bail from CSGen
due to the reference to `Capture` being turned
into an `ErrorExpr`. We ought to be able to
handle solving in the presence of such errors, but
for now I'm leaving it as future work.

rdar://93176036

Author: hamishknight

include/swift/AST/DiagnosticsSema.def

@@ -4890,6 +4890,9 @@ ERROR(regex_capture_types_failed_to_decode,none,
       "failed to decode capture types for regular expression literal; this may "
       "be a compiler bug", ())
 
+ERROR(must_import_regex_builder_module,none,
+      "regex builder requires the 'RegexBuilder' module be imported'", ())
+
 //------------------------------------------------------------------------------
 // MARK: Type Check Types
 //------------------------------------------------------------------------------

include/swift/AST/KnownIdentifiers.def

@@ -258,6 +258,7 @@ IDENTIFIER(Regex)
 IDENTIFIER_(regexString)
 IDENTIFIER(version)
 IDENTIFIER_(StringProcessing)
+IDENTIFIER(RegexBuilder)
 
 // Distributed actors
 IDENTIFIER(ActorID)

lib/Sema/CSDiagnostics.cpp

@@ -26,6 +26,7 @@
 #include "swift/AST/Expr.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/Initializer.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Pattern.h"
 #include "swift/AST/ProtocolConformance.h"
@@ -6460,6 +6461,9 @@ bool ArgumentMismatchFailure::diagnoseAsError() {
   if (diagnosePropertyWrapperMismatch())
     return true;
 
+  if (diagnoseAttemptedRegexBuilder())
+    return true;
+
   if (diagnoseTrailingClosureMismatch())
     return true;
 
@@ -6704,6 +6708,91 @@ bool ArgumentMismatchFailure::diagnosePropertyWrapperMismatch() const {
   return true;
 }
 
+/// Add a fix-it to insert an import of a module.
+static void fixItImport(InFlightDiagnostic &diag, Identifier moduleName,
+                        DeclContext *dc) {
+  auto *SF = dc->getParentSourceFile();
+  if (!SF)
+    return;
+
+  SourceLoc insertLoc;
+  bool isTrailing = true;
+
+  // Check if we can insert as the last import statement.
+  auto decls = SF->getTopLevelDecls();
+  for (auto *decl : decls) {
+    auto *importDecl = dyn_cast<ImportDecl>(decl);
+    if (!importDecl) {
+      if (insertLoc.isValid())
+        break;
+      continue;
+    }
+    insertLoc = importDecl->getEndLoc();
+  }
+
+  // If not, insert it as the first decl with a valid source location.
+  if (insertLoc.isInvalid()) {
+    for (auto *decl : decls) {
+      if (auto loc = decl->getStartLoc()) {
+        insertLoc = loc;
+        isTrailing = false;
+        break;
+      }
+    }
+  }
+
+  // If we didn't resolve to a valid location, give up.
+  if (insertLoc.isInvalid())
+    return;
+
+  SmallString<32> insertText;
+  if (isTrailing) {
+    insertText.append("\n");
+  }
+  insertText.append("import ");
+  insertText.append(moduleName.str());
+  if (isTrailing) {
+    diag.fixItInsertAfter(insertLoc, insertText);
+  } else {
+    insertText.append("\n\n");
+    diag.fixItInsert(insertLoc, insertText);
+  }
+}
+
+bool ArgumentMismatchFailure::diagnoseAttemptedRegexBuilder() const {
+  auto &ctx = getASTContext();
+
+  // Should be a lone trailing closure argument.
+  if (!Info.isTrailingClosure() || !Info.getArgList()->isUnary())
+    return false;
+
+  // Check if this an application of a Regex initializer, and the user has not
+  // imported RegexBuilder.
+  auto *ctor = dyn_cast_or_null<ConstructorDecl>(getCallee());
+  if (!ctor)
+    return false;
+
+  auto *ctorDC = ctor->getInnermostTypeContext();
+  if (!ctorDC || ctorDC->getSelfNominalTypeDecl() != ctx.getRegexDecl())
+    return false;
+
+  // If the RegexBuilder module is loaded, make sure it hasn't been imported.
+  // Note this will cause us to diagnose even if another SourceFile has
+  // imported RegexBuilder, and its extensions have leaked into this file. This
+  // is a longstanding lookup bug, and it's probably a good idea to suggest
+  // explicitly importing RegexBuilder regardless in that case.
+  if (auto *regexBuilderModule = ctx.getLoadedModule(ctx.Id_RegexBuilder)) {
+    auto &importCache = getASTContext().getImportCache();
+    if (importCache.isImportedBy(regexBuilderModule, getDC()))
+      return false;
+  }
+
+  // Suggest importing RegexBuilder.
+  auto diag = emitDiagnostic(diag::must_import_regex_builder_module);
+  fixItImport(diag, ctx.Id_RegexBuilder, getDC());
+  return true;
+}
+
 bool ArgumentMismatchFailure::diagnoseTrailingClosureMismatch() const {
   if (!Info.isTrailingClosure())
     return false;

lib/Sema/CSDiagnostics.h

@@ -1997,6 +1997,10 @@ class ArgumentMismatchFailure : public ContextualFailure {
   /// so we have to attend to that in diagnostics.
   bool diagnoseMisplacedMissingArgument() const;
 
+  /// Diagnose an attempt to pass a trailing closure to a Regex initializer
+  /// without importing RegexBuilder.
+  bool diagnoseAttemptedRegexBuilder() const;
+
 protected:
   /// \returns The position of the argument being diagnosed, starting at 1.
   unsigned getArgPosition() const { return Info.getArgPosition(); }

test/StringProcessing/Sema/Inputs/dummy.swift

@@ -0,0 +1 @@
+

test/StringProcessing/Sema/regex_builder_already_imported.swift

@@ -0,0 +1,16 @@
+// RUN: %target-typecheck-verify-swift -enable-bare-slash-regex -disable-availability-checking
+
+import RegexBuilder
+
+extension Regex where Output == Substring {
+  init(_ x: String) {} // expected-note {{'init(_:)' declared here}}
+}
+
+func foo() {
+  // FIXME: This diagnostic could probably be better, it's not clear we should
+  // be resolving to init(_ x: String) vs the result builder API and diagnosing
+  // the fact that Int isn't a RegexComponent.
+  _ = Regex { // expected-error {{trailing closure passed to parameter of type 'String' that does not accept a closure}}
+    0
+  }
+}

test/StringProcessing/Sema/regex_builder_fix_import_after_imports.swift

@@ -0,0 +1,20 @@
+// RUN: %empty-directory(%t)
+
+// Generate some dummy modules to import
+// RUN: %target-swift-frontend -emit-module -module-name A -o %t/A.swiftmodule %S/Inputs/dummy.swift
+// RUN: %target-swift-frontend -emit-module -module-name B -o %t/B.swiftmodule %S/Inputs/dummy.swift
+// RUN: %target-swift-frontend -emit-module -module-name C -o %t/C.swiftmodule %S/Inputs/dummy.swift
+
+// RUN: %target-typecheck-verify-swift -enable-bare-slash-regex -disable-availability-checking -I %t
+
+// REQUIRES: swift_in_compiler
+
+import A
+
+import B
+
+Regex { // expected-error {{regex builder requires the 'RegexBuilder' module be imported'}} {{14:9-9=\nimport RegexBuilder}}
+  /abc/
+}
+
+import C

test/StringProcessing/Sema/regex_builder_fix_import_decl.swift

@@ -0,0 +1,11 @@
+// RUN: %target-typecheck-verify-swift -enable-bare-slash-regex -disable-availability-checking
+
+// REQUIRES: swift_in_compiler
+
+struct S {
+  func foo() {
+    Regex { // expected-error {{regex builder requires the 'RegexBuilder' module be imported'}} {{5:1-1=import RegexBuilder\n\n}}
+      /abc/
+    }
+  }
+}

test/StringProcessing/Sema/regex_builder_fix_import_top_level.swift

@@ -0,0 +1,37 @@
+// RUN: %target-typecheck-verify-swift -enable-bare-slash-regex -disable-availability-checking
+
+// REQUIRES: swift_in_compiler
+
+Regex {} // expected-error {{regex builder requires the 'RegexBuilder' module be imported'}} {{5:1-1=import RegexBuilder\n\n}}
+
+Regex { // expected-error {{regex builder requires the 'RegexBuilder' module be imported'}} {{5:1-1=import RegexBuilder\n\n}}
+  /abc/
+}
+
+Regex { // expected-error {{regex builder requires the 'RegexBuilder' module be imported'}} {{5:1-1=import RegexBuilder\n\n}}
+  /abc/
+  /def/
+}
+
+Regex { // expected-error {{regex builder requires the 'RegexBuilder' module be imported'}} {{5:1-1=import RegexBuilder\n\n}}
+  /abc/
+  "def"
+  /g(h)(i)/
+}
+
+// FIXME: Unfortunately we bail from CSGen if we end up with an ErrorExpr, so
+// don't get a chance to diagnose. We ought to try solving with holes.
+// For now at least, this error should at least hopefully nudge users into
+// realizing they have a missing import.
+Regex {
+  Capture { // expected-error {{cannot find 'Capture' in scope}}
+    /abc/
+  }
+}
+
+// It's not clear what the user is doing here so fall back to regular diagnostics.
+Regex(a: 0) {
+  // expected-error@-1 {{cannot convert value of type 'Int' to expected argument type 'String'}}
+  // expected-error@-2 {{extra trailing closure passed in call}}
+  /abc/
+}