2 files changed, 76 insertions, 31 deletions

diff --git a/lib/Frontend/FixItRewriter.cpp b/lib/Frontend/FixItRewriter.cpp
index 20d452e76a..b26401212f 100644
--- a/lib/Frontend/FixItRewriter.cpp
+++ b/lib/Frontend/FixItRewriter.cpp
@@ -14,6 +14,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Frontend/FixItRewriter.h"
+#include "clang/Basic/FileManager.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/FrontendDiagnostic.h"
 #include "llvm/Support/raw_ostream.h"
@@ -34,47 +36,43 @@ FixItRewriter::~FixItRewriter() {
   Diags.setClient(Client);
 }
 
-bool FixItRewriter::WriteFixedFile(const std::string &InFileName,
-                                   const std::string &OutFileName) {
+bool FixItRewriter::WriteFixedFile(FileID ID, llvm::raw_ostream &OS) {
+  const RewriteBuffer *RewriteBuf = Rewrite.getRewriteBufferFor(ID);
+  if (!RewriteBuf) return true;
+  OS << std::string(RewriteBuf->begin(), RewriteBuf->end());
+  OS.flush();
+  return false;
+}
+
+bool FixItRewriter::WriteFixedFiles() {
   if (NumFailures > 0) {
     Diag(FullSourceLoc(), diag::warn_fixit_no_changes);
     return true;
   }
 
-  llvm::OwningPtr<llvm::raw_ostream> OwnedStream;
-  llvm::raw_ostream *OutFile;
-  if (!OutFileName.empty()) {
-    std::string Err;
-    OutFile = new llvm::raw_fd_ostream(OutFileName.c_str(), Err,
-                                       llvm::raw_fd_ostream::F_Binary);
-    OwnedStream.reset(OutFile);
-  } else if (InFileName == "-") {
-    OutFile = &llvm::outs();
-  } else {
-    llvm::sys::Path Path(InFileName);
+  for (iterator I = buffer_begin(), E = buffer_end(); I != E; ++I) {
+    const FileEntry *Entry = Rewrite.getSourceMgr().getFileEntryForID(I->first);
+    llvm::sys::Path Path(Entry->getName());
     std::string Suffix = Path.getSuffix();
     Path.eraseSuffix();
     Path.appendSuffix("fixit." + Suffix);
     std::string Err;
-    OutFile = new llvm::raw_fd_ostream(Path.c_str(), Err,
-                                       llvm::raw_fd_ostream::F_Binary);
-    OwnedStream.reset(OutFile);
-  }
-
-  FileID MainFileID = Rewrite.getSourceMgr().getMainFileID();
-  if (const RewriteBuffer *RewriteBuf =
-        Rewrite.getRewriteBufferFor(MainFileID)) {
-    *OutFile << std::string(RewriteBuf->begin(), RewriteBuf->end());
-  } else {
-    Diag(FullSourceLoc(), diag::note_fixit_main_file_unchanged);
+    llvm::raw_fd_ostream OS(Path.c_str(), Err, llvm::raw_fd_ostream::F_Binary);
+    if (!Err.empty()) {
+      Diags.Report(clang::diag::err_fe_unable_to_open_output)
+          << Path.c_str() << Err;
+      continue;
+    }
+    RewriteBuffer &RewriteBuf = I->second;
+    OS << std::string(RewriteBuf.begin(), RewriteBuf.end());
+    OS.flush();
   }
-  OutFile->flush();
 
   return false;
 }
 
 bool FixItRewriter::IncludeInDiagnosticCounts() const {
-  return Client? Client->IncludeInDiagnosticCounts() : true;
+  return Client ? Client->IncludeInDiagnosticCounts() : true;
 }
 
 void FixItRewriter::HandleDiagnostic(Diagnostic::Level DiagLevel,
@@ -85,6 +83,7 @@ void FixItRewriter::HandleDiagnostic(Diagnostic::Level DiagLevel,
   if (DiagLevel == Diagnostic::Ignored)
     return;
 
+  const SourceManager &SM = Rewrite.getSourceMgr();
   if (!FixItLocations.empty()) {
     // The user has specified the locations where we should perform
     // the various fix-it modifications.
@@ -99,15 +98,17 @@ void FixItRewriter::HandleDiagnostic(Diagnostic::Level DiagLevel,
     // See if the location of the error is one that matches what the
     // user requested.
     bool AcceptableLocation = false;
-    const FileEntry *File
-      = Rewrite.getSourceMgr().getFileEntryForID(
-                                            Info.getLocation().getFileID());
+    const FileEntry *File = SM.getFileEntryForID(
+                                                Info.getLocation().getFileID());
     unsigned Line = Info.getLocation().getSpellingLineNumber();
     unsigned Column = Info.getLocation().getSpellingColumnNumber();
     for (llvm::SmallVector<RequestedSourceLocation, 4>::iterator
            Loc = FixItLocations.begin(), LocEnd = FixItLocations.end();
          Loc != LocEnd; ++Loc) {
-      if (Loc->File == File && Loc->Line == Line && Loc->Column == Column) {
+      if (Loc->File == File &&
+          ((Loc->Line == 0 && Loc->Column == 0 &&
+            DiagLevel > Diagnostic::Note) ||
+           (Loc->Line == Line && Loc->Column == Column))) {
         AcceptableLocation = true;
         break;
       }
diff --git a/lib/Frontend/FrontendActions.cpp b/lib/Frontend/FrontendActions.cpp
index 6d7c6a8460..ce3e841b90 100644
--- a/lib/Frontend/FrontendActions.cpp
+++ b/lib/Frontend/FrontendActions.cpp
@@ -134,6 +134,23 @@ static bool AddFixItLocations(CompilerInstance &CI,
     FixItRewrite.addFixItLocation(Requested);
   }
 
+  const std::string &OutputFile = CI.getFrontendOpts().OutputFile;
+  if (Locs.empty() && !OutputFile.empty()) {
+    // FIXME: we will issue "FIX-IT applied suggested code changes" for every
+    // input, but only the main file will actually be rewritten.
+    const std::vector<std::pair<FrontendOptions::InputKind, std::string> > &Inputs =
+      CI.getFrontendOpts().Inputs;
+    for (unsigned i = 0, e = Inputs.size(); i != e; ++i) {
+      const FileEntry *File = CI.getFileManager().getFile(Inputs[i].second);
+      assert(File && "Input file not found in FileManager");
+      RequestedSourceLocation Requested;
+      Requested.File = File;
+      Requested.Line = 0;
+      Requested.Column = 0;
+      FixItRewrite.addFixItLocation(Requested);
+    }
+  }
+
   return true;
 }
 
@@ -149,7 +166,34 @@ bool FixItAction::BeginSourceFileAction(CompilerInstance &CI,
 
 void FixItAction::EndSourceFileAction() {
   const FrontendOptions &FEOpts = getCompilerInstance().getFrontendOpts();
-  Rewriter->WriteFixedFile(getCurrentFile(), FEOpts.OutputFile);
+  if (!FEOpts.OutputFile.empty()) {
+    // When called with 'clang -fixit -o filename' output only the main file.
+
+    const SourceManager &SM = getCompilerInstance().getSourceManager();
+    FileID MainFileID = SM.getMainFileID();
+    if (!Rewriter->IsModified(MainFileID)) {
+      getCompilerInstance().getDiagnostics().Report(
+          diag::note_fixit_main_file_unchanged);
+      return;
+    }
+
+    llvm::OwningPtr<llvm::raw_ostream> OwnedStream;
+    llvm::raw_ostream *OutFile;
+    if (FEOpts.OutputFile == "-") {
+      OutFile = &llvm::outs();
+    } else {
+      std::string Err;
+      OutFile = new llvm::raw_fd_ostream(FEOpts.OutputFile.c_str(), Err,
+                                         llvm::raw_fd_ostream::F_Binary);
+      OwnedStream.reset(OutFile);
+    }
+
+    Rewriter->WriteFixedFile(MainFileID, *OutFile);
+    return;
+  }
+
+  // Otherwise rewrite all files.
+  Rewriter->WriteFixedFiles();
 }
 
 ASTConsumer *RewriteObjCAction::CreateASTConsumer(CompilerInstance &CI,