[Distributed] Offer fixit for import Distributed when it is required (#72948)

Author: ktoso

include/swift/AST/DiagnosticEngine.h

@@ -727,6 +727,9 @@ namespace swift {
 
     static const char *fixItStringFor(const FixItID id);
 
+    /// Get the best location where an 'import' fixit might be offered.
+    SourceLoc getBestAddImportFixItLoc(const Decl *Member) const;
+
     /// Add a token-based replacement fix-it to the currently-active
     /// diagnostic.
     template <typename... ArgTypes>
@@ -783,6 +786,9 @@ namespace swift {
       return fixItReplaceChars(L, L, "%0", {Str});
     }
 
+    /// Add a fix-it suggesting to 'import' some module.
+    InFlightDiagnostic &fixItAddImport(StringRef ModuleName);
+
     /// Add an insertion fix-it to the currently-active diagnostic.  The
     /// text is inserted immediately *after* the token specified.
     InFlightDiagnostic &fixItInsertAfter(SourceLoc L, StringRef Str) {
@@ -1391,6 +1397,11 @@ namespace swift {
     SourceLoc getDefaultDiagnosticLoc() const {
       return bufferIndirectlyCausingDiagnostic;
     }
+    SourceLoc getBestAddImportFixItLoc(const Decl *Member,
+                                       SourceFile *sourceFile) const;
+    SourceLoc getBestAddImportFixItLoc(const Decl *Member) const {
+      return getBestAddImportFixItLoc(Member, nullptr);
+    }
   };
 
   /// Remember details about the state of a diagnostic engine and restore them

include/swift/AST/DiagnosticsSema.def

@@ -5436,7 +5436,7 @@ ERROR(conflicting_default_argument_isolation,none,
       "default argument cannot be both %0 and %1",
       (ActorIsolation, ActorIsolation))
 ERROR(distributed_decl_needs_explicit_distributed_import,none,
-      "%kind0 requires explicit import of Distributed module",
+      "%kind0 declared without importing module 'Distributed'",
       (const ValueDecl *))
 NOTE(distributed_func_cannot_overload_on_async_only,none,
      "%0 previously declared here, cannot overload distributed methods on effect only", (const ValueDecl *))

lib/AST/DiagnosticEngine.cpp

@@ -302,6 +302,75 @@ InFlightDiagnostic::fixItReplaceChars(SourceLoc Start, SourceLoc End,
   return *this;
 }
 
+SourceLoc
+DiagnosticEngine::getBestAddImportFixItLoc(const Decl *Member,
+                                           SourceFile *sourceFile) const {
+  auto &SM = SourceMgr;
+
+  SourceLoc bestLoc;
+
+  auto SF =
+      sourceFile ? sourceFile : Member->getDeclContext()->getParentSourceFile();
+  if (!SF) {
+    return bestLoc;
+  }
+
+  for (auto item : SF->getTopLevelItems()) {
+    // If we found an import declaration, we want to insert after it.
+    if (auto importDecl =
+            dyn_cast_or_null<ImportDecl>(item.dyn_cast<Decl *>())) {
+      SourceLoc loc = importDecl->getEndLoc();
+      if (loc.isValid()) {
+        bestLoc = Lexer::getLocForEndOfLine(SM, loc);
+      }
+
+      // Keep looking for more import declarations.
+      continue;
+    }
+
+    // If we got a location based on import declarations, we're done.
+    if (bestLoc.isValid())
+      break;
+
+    // For any other item, we want to insert before it.
+    SourceLoc loc = item.getStartLoc();
+    if (loc.isValid()) {
+      bestLoc = Lexer::getLocForStartOfLine(SM, loc);
+      break;
+    }
+  }
+
+  return bestLoc;
+}
+
+InFlightDiagnostic &InFlightDiagnostic::fixItAddImport(StringRef ModuleName) {
+  assert(IsActive && "Cannot modify an inactive diagnostic");
+  auto Member = Engine->ActiveDiagnostic->getDecl();
+  SourceLoc bestLoc = Engine->getBestAddImportFixItLoc(Member);
+
+  if (bestLoc.isValid()) {
+    llvm::SmallString<64> importText;
+
+    // @_spi imports.
+    if (Member->isSPI()) {
+      auto spiGroups = Member->getSPIGroups();
+      if (!spiGroups.empty()) {
+        importText += "@_spi(";
+        importText += spiGroups[0].str();
+        importText += ") ";
+      }
+    }
+
+    importText += "import ";
+    importText += ModuleName;
+    importText += "\n";
+
+    return fixItInsert(bestLoc, importText);
+  }
+
+  return *this;
+}
+
 InFlightDiagnostic &InFlightDiagnostic::fixItExchange(SourceRange R1,
                                                       SourceRange R2) {
   assert(IsActive && "Cannot modify an inactive diagnostic");

lib/Basic/SourceLoc.cpp

@@ -10,14 +10,15 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "swift/Basic/Range.h"
 #include "swift/Basic/SourceLoc.h"
+#include "swift/AST/SourceFile.h"
+#include "swift/Basic/Range.h"
 #include "swift/Basic/SourceManager.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/PrettyStackTrace.h"
-#include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace swift;
 

lib/Sema/CSDiagnostics.cpp

@@ -6195,34 +6195,8 @@ bool InaccessibleMemberFailure::diagnoseAsError() {
                      Member->getDescriptiveKind(), Member->getName(),
                      definingModule->getName());
 
-    auto enclosingSF = getDC()->getParentSourceFile();
-    SourceLoc bestLoc;
-    SourceManager &srcMgr = Member->getASTContext().SourceMgr;
-    for (auto item : enclosingSF->getTopLevelItems()) {
-      // If we found an import declaration, we want to insert after it.
-      if (auto importDecl =
-              dyn_cast_or_null<ImportDecl>(item.dyn_cast<Decl *>())) {
-        SourceLoc loc = importDecl->getEndLoc();
-        if (loc.isValid()) {
-          bestLoc = Lexer::getLocForEndOfLine(srcMgr, loc);
-        }
-
-        // Keep looking for more import declarations.
-        continue;
-      }
-
-      // If we got a location based on import declarations, we're done.
-      if (bestLoc.isValid())
-        break;
-
-      // For any other item, we want to insert before it.
-      SourceLoc loc = item.getStartLoc();
-      if (loc.isValid()) {
-        bestLoc = Lexer::getLocForStartOfLine(srcMgr, loc);
-        break;
-      }
-    }
-
+    SourceLoc bestLoc =
+        getBestAddImportFixItLocation(Member, getDC()->getParentSourceFile());
     if (bestLoc.isValid()) {
       llvm::SmallString<64> importText;
 

lib/Sema/CSDiagnostics.h

@@ -90,6 +90,12 @@ class FailureDiagnostic {
 
   ConstraintLocator *getLocator() const { return Locator; }
 
+  SourceLoc getBestAddImportFixItLocation(const Decl *Member,
+                                          SourceFile *sourceFile) const {
+    auto &engine = Member->getASTContext().Diags;
+    return engine.getBestAddImportFixItLoc(Member, sourceFile);
+  }
+
   Type getType(ASTNode node, bool wantRValue = true) const;
 
   /// Get type associated with a given ASTNode without resolving it,

lib/Sema/TypeCheckDistributed.cpp

@@ -49,7 +49,8 @@ DistributedModuleIsAvailableRequest::evaluate(Evaluator &evaluator,
   auto DistributedModule = C.getLoadedModule(C.Id_Distributed);
   if (!DistributedModule) {
     decl->diagnose(diag::distributed_decl_needs_explicit_distributed_import,
-                   decl);
+                   decl)
+        .fixItAddImport("Distributed");
     return false;
   }
 

test/Distributed/distributed_actor_missing_distributed_import.swift

@@ -5,7 +5,7 @@
 actor SomeActor {}
 
 distributed actor DA {}
-// expected-error@-1{{distributed actor 'DA' requires explicit import of Distributed module}}
+// expected-error@-1{{distributed actor 'DA' declared without importing module 'Distributed'}}
 
 distributed actor class DAC {}
 // expected-error@-1{{distributed' can only be applied to 'actor' definitions, and distributed actor-isolated async functions}}
@@ -18,23 +18,23 @@ actor A {
 }
 
 actor B {
-  distributed var neverOk: String { // expected-error{{distributed property 'neverOk' requires explicit import of Distributed module}}
+  distributed var neverOk: String { // expected-error{{distributed property 'neverOk' declared without importing module 'Distributed'}}
     ""
   }
 }
 
-  // expected-error@+1{{distributed actor 'DA2' requires explicit import of Distributed module}}
+  // expected-error@+1{{distributed actor 'DA2' declared without importing module 'Distributed'}}
 distributed actor DA2 {
 
   func normal() async {}
 
-  // expected-error@+1{{distributed instance method 'dist()' requires explicit import of Distributed module}}
+  // expected-error@+1{{distributed instance method 'dist()' declared without importing module 'Distributed'}}
   distributed func dist() {}
 
-  // expected-error@+1{{distributed instance method 'distAsync()' requires explicit import of Distributed module}}
+  // expected-error@+1{{distributed instance method 'distAsync()' declared without importing module 'Distributed'}}
   distributed func distAsync() async {}
 
-  // expected-error@+1{{distributed property 'neverOk' requires explicit import of Distributed module}}
+  // expected-error@+1{{distributed property 'neverOk' declared without importing module 'Distributed'}}
   distributed var neverOk: String {
     ""
   }

test/Distributed/distributed_missing_import.swift

@@ -5,8 +5,7 @@
 actor SomeActor { }
 
 distributed actor MissingImportDistributedActor_0 { }
-// expected-error@-1{{distributed actor 'MissingImportDistributedActor_0' requires explicit import of Distributed module}}
+// expected-error@-1{{distributed actor 'MissingImportDistributedActor_0' declared without importing module 'Distributed'}}{{5:1-1=import Distributed\n}}
 
 let t: DistributedActorSystem // expected-error{{cannot find type 'DistributedActorSystem' in scope}}
-let a: ActorAddress // expected-error{{cannot find type 'ActorAddress' in scope}}
 

test/Distributed/distributed_missing_import_fixit.swift

@@ -0,0 +1,13 @@
+// RUN: %target-typecheck-verify-swift -disable-availability-checking
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+import Swift // just any import so we can anchor the fixit onto it
+
+actor SomeActor { }
+
+distributed actor MissingImportDistributedActor_0 { }
+// expected-error@-1{{distributed actor 'MissingImportDistributedActor_0' declared without importing module 'Distributed'}}{{6:1-1=import Distributed\n}}
+
+let t: DistributedActorSystem // expected-error{{cannot find type 'DistributedActorSystem' in scope}}
+