Merge pull request #26342 from xedin/rdar-53404683

[AST] DiagnosticEngine: retain strings used in tentative diagnostic arguments

Author: xedin

include/swift/AST/DiagnosticEngine.h

@@ -22,6 +22,8 @@
 #include "swift/AST/TypeLoc.h"
 #include "swift/AST/DeclNameLoc.h"
 #include "swift/AST/DiagnosticConsumer.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Allocator.h"
 
 namespace swift {
   class Decl;
@@ -323,6 +325,8 @@ namespace swift {
     SourceLoc Loc;
     const Decl *Decl = nullptr;
 
+    friend DiagnosticEngine;
+
   public:
     // All constructors are intentionally implicit.
     template<typename ...ArgTypes>
@@ -564,6 +568,12 @@ namespace swift {
     /// results that we can point to on the command line.
     llvm::DenseMap<const Decl *, SourceLoc> PrettyPrintedDeclarations;
 
+    llvm::BumpPtrAllocator TransactionAllocator;
+    /// A set of all strings involved in current transactional chain.
+    /// This is required because diagnostics are not directly emitted
+    /// but rather stored until all transactions complete.
+    llvm::StringMap<char, llvm::BumpPtrAllocator &> TransactionStrings;
+
     /// The number of open diagnostic transactions. Diagnostics are only
     /// emitted once all transactions have closed.
     unsigned TransactionCount = 0;
@@ -579,8 +589,8 @@ namespace swift {
 
   public:
     explicit DiagnosticEngine(SourceManager &SourceMgr)
-      : SourceMgr(SourceMgr), ActiveDiagnostic() {
-    }
+        : SourceMgr(SourceMgr), ActiveDiagnostic(),
+          TransactionStrings(TransactionAllocator) {}
 
     /// hadAnyError - return true if any *error* diagnostics have been emitted.
     bool hadAnyError() const { return state.hadAnyError(); }
@@ -797,6 +807,11 @@ namespace swift {
         DiagnosticFormatOptions FormatOpts = DiagnosticFormatOptions());
 
   private:
+    /// Called when tentative diagnostic is about to be flushed,
+    /// to apply any required transformations e.g. copy string arguments
+    /// to extend their lifetime.
+    void onTentativeDiagnosticFlush(Diagnostic &diagnostic);
+
     /// Flush the active diagnostic.
     void flushActiveDiagnostic();
 
@@ -856,6 +871,9 @@ namespace swift {
     bool IsOpen = true;
 
   public:
+    DiagnosticTransaction(const DiagnosticTransaction &) = delete;
+    DiagnosticTransaction &operator=(const DiagnosticTransaction &) = delete;
+
     explicit DiagnosticTransaction(DiagnosticEngine &engine)
       : Engine(engine),
         PrevDiagnostics(Engine.TentativeDiagnostics.size()),
@@ -870,6 +888,11 @@ namespace swift {
       if (IsOpen) {
         commit();
       }
+
+      if (Depth == 0) {
+        Engine.TransactionStrings.clear();
+        Engine.TransactionAllocator.Reset();
+      }
     }
 
     /// Abort and close this transaction and erase all diagnostics

lib/AST/DiagnosticEngine.cpp

@@ -754,6 +754,7 @@ void DiagnosticEngine::flushActiveDiagnostic() {
   if (TransactionCount == 0) {
     emitDiagnostic(*ActiveDiagnostic);
   } else {
+    onTentativeDiagnosticFlush(*ActiveDiagnostic);
     TentativeDiagnostics.emplace_back(std::move(*ActiveDiagnostic));
   }
   ActiveDiagnostic.reset();
@@ -940,3 +941,17 @@ BufferIndirectlyCausingDiagnosticRAII::BufferIndirectlyCausingDiagnosticRAII(
   if (loc.isValid())
     Diags.setBufferIndirectlyCausingDiagnosticToInput(loc);
 }
+
+void DiagnosticEngine::onTentativeDiagnosticFlush(Diagnostic &diagnostic) {
+  for (auto &argument : diagnostic.Args) {
+    if (argument.getKind() != DiagnosticArgumentKind::String)
+      continue;
+
+    auto content = argument.getAsString();
+    if (content.empty())
+      continue;
+
+    auto I = TransactionStrings.insert(std::make_pair(content, char())).first;
+    argument = DiagnosticArgument(StringRef(I->getKeyData()));
+  }
+}

lib/Sema/DerivedConformanceCodable.cpp

@@ -1180,7 +1180,7 @@ ValueDecl *DerivedConformance::deriveEncodable(ValueDecl *requirement) {
   // diagnostics, then potentially collect notes. If we succeed in
   // synthesizing Encodable, we can cancel the transaction and get rid of the
   // fake failures.
-  auto diagnosticTransaction = DiagnosticTransaction(TC.Context.Diags);
+  DiagnosticTransaction diagnosticTransaction(TC.Context.Diags);
   TC.diagnose(ConformanceDecl, diag::type_does_not_conform,
               Nominal->getDeclaredType(), getProtocolType());
   TC.diagnose(requirement, diag::no_witnesses, diag::RequirementKind::Func,
@@ -1216,7 +1216,7 @@ ValueDecl *DerivedConformance::deriveDecodable(ValueDecl *requirement) {
   // diagnostics produced by canSynthesize and deriveDecodable_init to produce
   // them in the right order -- see the comment in deriveEncodable for
   // background on this transaction.
-  auto diagnosticTransaction = DiagnosticTransaction(TC.Context.Diags);
+  DiagnosticTransaction diagnosticTransaction(TC.Context.Diags);
   TC.diagnose(ConformanceDecl->getLoc(), diag::type_does_not_conform,
               Nominal->getDeclaredType(), getProtocolType());
   TC.diagnose(requirement, diag::no_witnesses,