Improve diagnostics for uses of unsafe declarations in functions

Instead of producing a warning for each use of an unsafe entity,
collect all of the uses of unsafe constructs within a given function
and batch them together in a single diagnostic at the function level
that tells you what you can do (add `@unsafe` or `@safe(unchecked)`,
depending on whether all unsafe uses were in the definition), plus
notes identifying every unsafe use within that declaration. The new
diagnostic renderer nicely collects together in a single snippet, so
it's easier to reason about.

Here's an example from the embedded runtime that previously would have
been 6 separate warnings, each with 1-2 notes:

```
swift/stdlib/public/core/EmbeddedRuntime.swift:397:13: warning: global function 'swift_retainCount' involves unsafe code; use '@safe(unchecked)' to assert that the code is memory-safe
395 |
396 | @_cdecl("swift_retainCount")
397 | public func swift_retainCount(object: Builtin.RawPointer) -> Int {
    |             `- warning: global function 'swift_retainCount' involves unsafe code; use '@safe(unchecked)' to assert that the code is memory-safe
398 |   if !isValidPointerForNativeRetain(object: object) { return 0 }
399 |   let o = UnsafeMutablePointer<HeapObject>(object)
    |           |                              `- note: call to unsafe initializer 'init(_:)'
    |           `- note: reference to unsafe generic struct 'UnsafeMutablePointer'
400 |   let refcount = refcountPointer(for: o)
    |                  |                    `- note: reference to let 'o' involves unsafe type 'UnsafeMutablePointer<HeapObject>'
    |                  `- note: call to global function 'refcountPointer(for:)' involves unsafe type 'UnsafeMutablePointer<Int>'
401 |   return loadAcquire(refcount) & HeapObject.refcountMask
    |          |           `- note: reference to let 'refcount' involves unsafe type 'UnsafeMutablePointer<Int>'
    |          `- note: call to global function 'loadAcquire' involves unsafe type 'UnsafeMutablePointer<Int>'
402 | }
403 |
```

Note that we have lost a little bit of information, because we no
longer produce "unsafe declaration was here" notes pointing back at
things like `UnsafeMutablePointer` or `recountPointer(for:)`. However,
strict memory safety tends to be noisy to turn on, so it's worth
losing a little bit of easily-recovered information to gain some
brevity.

Author: DougGregor

include/swift/AST/DiagnosticsSema.def

@@ -8072,6 +8072,19 @@ NOTE(sending_function_result_with_sending_param_note, none,
 //------------------------------------------------------------------------------
 ERROR(unsafe_attr_disabled,none,
       "attribute requires '-enable-experimental-feature AllowUnsafeAttribute'", ())
+
+GROUPED_WARNING(decl_involves_unsafe,Unsafe,none,
+      "%kindbase0 involves unsafe code; "
+      "use %select{'@unsafe' to indicate that its use is not memory-safe|"
+      "'@safe(unchecked)' to assert that the code is memory-safe}1",
+      (const Decl *, bool))
+NOTE(note_reference_to_unsafe_decl,none,
+      "%select{reference|call}0 to unsafe %kind1",
+      (bool, const ValueDecl *))
+NOTE(note_reference_to_unsafe_typed_decl,none,
+     "%select{reference|call}0 to %kind1 involves unsafe type %2",
+     (bool, const ValueDecl *, Type))
+
 GROUPED_WARNING(override_safe_withunsafe,Unsafe,none,
       "override of safe %0 with unsafe %0", (DescriptiveDeclKind))
 GROUPED_WARNING(witness_unsafe,Unsafe,none,

include/swift/AST/SourceFile.h

@@ -27,6 +27,7 @@
 namespace swift {
 
 class PersistentParserState;
+struct SourceFileExtras;
 
 /// Kind of import affecting how a decl can be reexported.
 ///
@@ -243,6 +244,9 @@ class SourceFile final : public FileUnit {
   /// Storage for \c HasImportsMatchingFlagRequest.
   ImportOptions cachedImportOptions;
 
+  /// Extra information for the source file, allocated as needed.
+  SourceFileExtras *extras = nullptr;
+
   friend ASTContext;
 
 public:
@@ -368,6 +372,11 @@ class SourceFile final : public FileUnit {
   /// \c #sourceLocation(file:) declarations.
   llvm::StringMap<SourceFilePathInfo> getInfoForUsedFilePaths() const;
 
+  /// Retrieve "extra" information associated with this source file, which is
+  /// lazily and separately constructed. Use this for scratch information
+  /// that isn't needed for all source files.
+  SourceFileExtras &getExtras() const;
+
   SourceFile(ModuleDecl &M, SourceFileKind K, unsigned bufferID,
              ParsingOptions parsingOpts = {}, bool isPrimary = false);
 

include/swift/AST/SourceFileExtras.h

@@ -0,0 +1,36 @@
+//===--- SourceFileExtras.h - Extra data for a source file ------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_SOURCEFILEEXTRAS_H
+#define SWIFT_AST_SOURCEFILEEXTRAS_H
+
+#include "swift/AST/UnsafeUse.h"
+#include "llvm/ADT/DenseMap.h"
+#include <vector>
+
+namespace swift {
+
+class Decl;
+
+/// Extra information associated with a source file that is lazily created and
+/// stored in a separately-allocated side structure.
+struct SourceFileExtras {
+  /// Captures all of the unsafe uses associated with a given declaration.
+  ///
+  /// The declaration is the entity that can be annotated (e.g., with @unsafe)
+  /// to suppress all of the unsafe-related diagnostics listed here.
+  llvm::DenseMap<const Decl *, std::vector<UnsafeUse>> unsafeUses;
+};
+  
+}
+
+#endif // SWIFT_AST_SOURCEFILEEXTRAS_H

include/swift/AST/UnsafeUse.h

@@ -67,6 +67,7 @@ class UnsafeUse {
 
     struct {
       NormalProtocolConformance *conformance;
+      DeclContext *declContext;
       const void *location;
     } conformance;
 
@@ -139,9 +140,11 @@ class UnsafeUse {
   }
 
   static UnsafeUse forConformance(NormalProtocolConformance *conformance,
-                                  SourceLoc location) {
+                                  SourceLoc location,
+                                  DeclContext *dc) {
     UnsafeUse result(UnsafeConformance);
     result.storage.conformance.conformance = conformance;
+    result.storage.conformance.declContext = dc;
     result.storage.conformance.location = location.getOpaquePointerValue();
     return result;
   }
@@ -230,10 +233,14 @@ class UnsafeUse {
       return storage.entity.declContext;
 
     case Override:
+      return getDecl()->getDeclContext();
+
     case Witness:
     case TypeWitness:
+      return getConformance()->getDeclContext();
+
     case UnsafeConformance:
-      return nullptr;
+      return storage.conformance.declContext;
     }
   }
 

lib/AST/Module.cpp

@@ -41,6 +41,7 @@
 #include "swift/AST/PrintOptions.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/SourceFileExtras.h"
 #include "swift/AST/SynthesizedFileUnit.h"
 #include "swift/AST/Type.h"
 #include "swift/AST/TypeCheckRequests.h"
@@ -1175,6 +1176,17 @@ ASTNode SourceFile::getNodeInEnclosingSourceFile() const {
   return ASTNode::getFromOpaqueValue(getGeneratedSourceFileInfo()->astNode);
 }
 
+SourceFileExtras &SourceFile::getExtras() const {
+  if (!extras) {
+    const_cast<SourceFile *>(this)->extras = new SourceFileExtras();
+    getASTContext().addCleanup([extras=this->extras] {
+      delete extras;
+    });
+  }
+
+  return *extras;
+}
+
 void ModuleDecl::lookupClassMember(ImportPath::Access accessPath,
                                    DeclName name,
                                    SmallVectorImpl<ValueDecl*> &results) const {

lib/Sema/TypeCheckAvailability.h

@@ -274,7 +274,12 @@ std::pair<const Decl *, bool /*inDefinition*/>
 enclosingContextForUnsafe(SourceLoc referenceLoc, const DeclContext *referenceDC);
 
 /// Diagnose the given unsafe use right now.
-void diagnoseUnsafeUse(const UnsafeUse &use);
+void diagnoseUnsafeUse(const UnsafeUse &use, bool asNote = false);
+
+/// Diagnose any unsafe uses within the signature or definition of the given
+/// declaration, if there are any.
+void diagnoseUnsafeUsesIn(const Decl *decl);
+
 
 } // namespace swift
 

lib/Sema/TypeCheckProtocol.cpp

@@ -2469,7 +2469,8 @@ checkIndividualConformance(NormalProtocolConformance *conformance) {
 
     if (!conformance->getDeclContext()->allowsUnsafe()) {
       diagnoseUnsafeUse(
-          UnsafeUse::forConformance(conformance, ComplainLoc));
+          UnsafeUse::forConformance(conformance, ComplainLoc,
+                                    conformance->getDeclContext()));
     }
   }
 

lib/Sema/TypeCheckStmt.cpp

@@ -2949,6 +2949,8 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
     TypeChecker::checkFunctionEffects(AFD);
   }
 
+  diagnoseUnsafeUsesIn(AFD);
+
   return hadError ? errorBody() : body;
 }
 

lib/Sema/TypeCheckUnsafe.cpp

@@ -17,11 +17,40 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/UnsafeUse.h"
 #include "swift/AST/DiagnosticsSema.h"
+#include "swift/AST/SourceFile.h"
+#include "swift/AST/SourceFileExtras.h"
 #include "TypeCheckAvailability.h"
 #include "TypeCheckType.h"
 
 using namespace swift;
 
+static std::pair<const Decl *, bool /*inDefinition*/>
+enclosingContextForUnsafe(const UnsafeUse &use) {
+  return swift::enclosingContextForUnsafe(use.getLocation(),
+                                          use.getDeclContext());
+}
+
+/// Whether this particular unsafe use occurs within the definition of an
+/// entity, but not in its signature, meaning that the unsafety could be
+/// encapsulated.
+static bool isUnsafeUseInDefinition(const UnsafeUse &use) {
+  switch (use.getKind()) {
+  case UnsafeUse::Override:
+  case UnsafeUse::Witness:
+  case UnsafeUse::TypeWitness:
+  case UnsafeUse::UnsafeConformance:
+  case UnsafeUse::UnownedUnsafe:
+  case UnsafeUse::NonisolatedUnsafe:
+    // Never part of the definition. These are always part of the interface.
+    return false;
+
+  case UnsafeUse::ReferenceToUnsafe:
+  case UnsafeUse::CallToUnsafe:
+    return enclosingContextForUnsafe(use).second;
+  }
+}
+
+
 static void suggestUnsafeOnEnclosingDecl(
     SourceLoc referenceLoc, const DeclContext *referenceDC) {
   const Decl *decl = nullptr;
@@ -58,7 +87,39 @@ static void suggestUnsafeMarkerOnConformance(
   ).fixItInsert(decl->getAttributeInsertionLoc(false), "@unsafe ");
 }
 
-void swift::diagnoseUnsafeUse(const UnsafeUse &use) {
+static bool shouldDeferUnsafeUseIn(const Decl *decl) {
+  if (auto func = dyn_cast_or_null<AbstractFunctionDecl>(decl)) {
+    if (func->hasBody()) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+/// Retrieve the extra information
+static SourceFileExtras &getSourceFileExtrasFor(const Decl *decl) {
+  auto dc = decl->getDeclContext();
+  return dc->getOutermostParentSourceFile()->getExtras();
+}
+
+void swift::diagnoseUnsafeUse(const UnsafeUse &use, bool asNote) {
+  if (!asNote) {
+    // If we can associate this unsafe use within a particular declaration, do so.
+    // It will be diagnosed later, along with all other unsafe uses within this
+    // same declaration.
+    if (use.getKind() == UnsafeUse::ReferenceToUnsafe ||
+        use.getKind() == UnsafeUse::CallToUnsafe) {
+      auto [enclosingDecl, _] = enclosingContextForUnsafe(
+          use.getLocation(), use.getDeclContext());
+      if (enclosingDecl && shouldDeferUnsafeUseIn(enclosingDecl)) {
+        getSourceFileExtrasFor(enclosingDecl).unsafeUses[enclosingDecl]
+          .push_back(use);
+        return;
+      }
+    }
+  }
+
   switch (use.getKind()) {
   case UnsafeUse::Override: {
     auto override = use.getDecl();
@@ -142,18 +203,52 @@ void swift::diagnoseUnsafeUse(const UnsafeUse &use) {
           use.getDeclContext(),
           [&](Type specificType) {
             ctx.Diags.diagnose(
-                loc, diag::reference_to_unsafe_typed_decl, isCall, decl,
-                specificType);
+                loc,
+                asNote ? diag::note_reference_to_unsafe_typed_decl
+                       : diag::reference_to_unsafe_typed_decl,
+                isCall, decl, specificType);
           });
-      suggestUnsafeOnEnclosingDecl(loc, use.getDeclContext());
-      decl->diagnose(diag::unsafe_decl_here, decl);
     } else {
-      ctx.Diags
-        .diagnose(loc, diag::reference_to_unsafe_decl, isCall, decl);
+      ctx.Diags.diagnose(
+          loc,
+          asNote ? diag::note_reference_to_unsafe_decl
+                 : diag::reference_to_unsafe_decl,
+          isCall, decl);
+    }
+
+    if (!asNote) {
       suggestUnsafeOnEnclosingDecl(loc, use.getDeclContext());
       decl->diagnose(diag::unsafe_decl_here, decl);
     }
+
     return;
   }
   }
 }
+
+void swift::diagnoseUnsafeUsesIn(const Decl *decl) {
+  auto &extras = getSourceFileExtrasFor(decl);
+  auto known = extras.unsafeUses.find(decl);
+  if (known == extras.unsafeUses.end())
+    return;
+
+  // Take the unsafe uses.
+  auto unsafeUses = std::move(known->second);
+  extras.unsafeUses.erase(known);
+  if (unsafeUses.empty())
+    return;
+
+  // Determine whether all of the uses are in the definition.
+  bool canEncapsulateUnsafety = std::all_of(
+      unsafeUses.begin(), unsafeUses.end(), isUnsafeUseInDefinition);
+  StringRef fixItStr = canEncapsulateUnsafety
+      ? "@safe(unchecked) "
+      : "@unsafe ";
+  decl->diagnose(diag::decl_involves_unsafe, decl, canEncapsulateUnsafety)
+    .fixItInsert(decl->getAttributeInsertionLoc(/*forModifier=*/false),
+                 fixItStr);
+
+  for (const auto &unsafeUse : unsafeUses) {
+    diagnoseUnsafeUse(unsafeUse, /*asNote=*/true);
+  }
+}

test/Interop/Cxx/class/safe-interop-mode.swift

@@ -60,22 +60,22 @@ using SpanOfIntAlias = SpanOfInt;
 import Test
 import CoreFoundation
 
-// expected-note@+1{{make global function 'useUnsafeParam' @unsafe to indicate that its use is not memory-safe}}{{1-1=@unsafe }}
-func useUnsafeParam(x: Unannotated) { // expected-warning{{reference to unsafe struct 'Unannotated'}}
+// expected-warning@+1{{global function 'useUnsafeParam' involves unsafe code; use '@unsafe' to indicate that its use is not memory-safe}}{{1-1=@unsafe }}
+func useUnsafeParam(x: Unannotated) { // expected-note{{reference to unsafe struct 'Unannotated'}}
 }
 
-// expected-note@+2{{make global function 'useUnsafeParam2' @unsafe to indicate that its use is not memory-safe}}{{10:1-1=@unsafe }}
+// expected-warning@+2{{global function 'useUnsafeParam2' involves unsafe code; use '@unsafe' to indicate that its use is not memory-safe}}{{10:1-1=@unsafe }}
 @available(SwiftStdlib 5.8, *)
-func useUnsafeParam2(x: UnsafeReference) { // expected-warning{{reference to unsafe class 'UnsafeReference'}}
+func useUnsafeParam2(x: UnsafeReference) { // expected-note{{reference to unsafe class 'UnsafeReference'}}
 }
 
-// expected-note@+1{{make global function 'useUnsafeParam3' @unsafe to indicate that its use is not memory-safe}}{{1-1=@unsafe }}
-func useUnsafeParam3(x: UnknownEscapabilityAggregate) { // expected-warning{{reference to unsafe struct 'UnknownEscapabilityAggregate'}}
+// expected-warning@+1{{global function 'useUnsafeParam3' involves unsafe code; use '@unsafe' to indicate that its use is not memory-safe}}{{1-1=@unsafe }}
+func useUnsafeParam3(x: UnknownEscapabilityAggregate) { // expected-note{{reference to unsafe struct 'UnknownEscapabilityAggregate'}}
 }
 
-// expected-note@+1{{make global function 'useSafeParams' @safe(unchecked) to allow it to use unsafe constructs in its definition}}{{1-1=@safe(unchecked) }}
+// expected-warning@+1{{global function 'useSafeParams' involves unsafe code; use '@safe(unchecked)' to assert that the code is memory-safe}}{{1-1=@safe(unchecked) }}
 func useSafeParams(x: Owner, y: View, z: SafeEscapableAggregate, c: MyContainer) {
-    let _ = c.__beginUnsafe() // expected-warning{{call to unsafe instance method '__beginUnsafe'}}
+    let _ = c.__beginUnsafe() // expected-note{{call to unsafe instance method '__beginUnsafe()'}}
 }
 
 func useCfType(x: CFArray) {