Add removeConsumer to DiagnosticEngine

Author: DavidGoldman

include/swift/AST/DiagnosticEngine.h

@@ -596,6 +596,12 @@ namespace swift {
       Consumers.push_back(&Consumer);
     }
 
+    /// Remove a specific DiagnosticConsumer.
+    void removeConsumer(DiagnosticConsumer &Consumer) {
+      Consumers.erase(
+          std::remove(Consumers.begin(), Consumers.end(), &Consumer));
+    }
+
     /// Remove and return all \c DiagnosticConsumers.
     std::vector<DiagnosticConsumer *> takeConsumers() {
       auto Result = std::vector<DiagnosticConsumer*>(Consumers.begin(),

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -449,11 +449,7 @@ bool SwiftASTManager::initCompilerInvocation(CompilerInvocation &Invocation,
   });
 
   // Remove the StreamDiagConsumer as it's no longer needed.
-  std::vector<DiagnosticConsumer *> OldC = Diags.takeConsumers();
-  OldC.erase(std::remove(OldC.begin(), OldC.end(), &DiagConsumer));
-  for (DiagnosticConsumer *Consumer : OldC) {
-    Diags.addConsumer(*Consumer);
-  }
+  Diags.removeConsumer(DiagConsumer);
 
   if (HadError) {
     Error = ErrOS.str();