2 files changed, 50 insertions, 5 deletions

diff --git a/tools/diagtool/ShowEnabledWarnings.cpp b/tools/diagtool/ShowEnabledWarnings.cpp
index 7162451b4e73..abd69fd0af61 100644
--- a/tools/diagtool/ShowEnabledWarnings.cpp
+++ b/tools/diagtool/ShowEnabledWarnings.cpp
@@ -60,7 +60,7 @@ createDiagnostics(unsigned int argc, char **argv) {
   // well formed diagnostic object.
   TextDiagnosticBuffer *DiagsBuffer = new TextDiagnosticBuffer;
   IntrusiveRefCntPtr<DiagnosticsEngine> InterimDiags(
-    new DiagnosticsEngine(DiagIDs, DiagsBuffer));
+    new DiagnosticsEngine(DiagIDs, new DiagnosticOptions(), DiagsBuffer));
 
   // Try to build a CompilerInvocation.
   OwningPtr<CompilerInvocation> Invocation(
@@ -71,7 +71,7 @@ createDiagnostics(unsigned int argc, char **argv) {
 
   // Build the diagnostics parser
   IntrusiveRefCntPtr<DiagnosticsEngine> FinalDiags =
-    CompilerInstance::createDiagnostics(Invocation->getDiagnosticOpts(),
+    CompilerInstance::createDiagnostics(&Invocation->getDiagnosticOpts(),
                                         argc, argv);
   if (!FinalDiags)
     return NULL;
diff --git a/tools/diagtool/TreeView.cpp b/tools/diagtool/TreeView.cpp
index 9db2c26dd92d..bf9f766c309a 100644
--- a/tools/diagtool/TreeView.cpp
+++ b/tools/diagtool/TreeView.cpp
@@ -14,11 +14,13 @@
 #include "DiagTool.h"
 #include "DiagnosticNames.h"
 #include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/DiagnosticOptions.h"
 #include "llvm/Support/Format.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "clang/AST/ASTDiagnostic.h"
 #include "clang/Basic/AllDiagnostics.h"
+#include "llvm/Support/Process.h"
 
 DEF_DIAGTOOL("tree",
              "Show warning flags in a tree view",
@@ -31,11 +33,39 @@ static void printUsage() {
   llvm::errs() << "Usage: diagtool tree [--flags-only] [<diagnostic-group>]\n";
 }
 
+static bool showColors(llvm::raw_ostream &out) {
+  if (&out != &llvm::errs() && &out != &llvm::outs())
+    return false;
+  return llvm::errs().is_displayed() && llvm::outs().is_displayed();
+}
+
+static void setColor(bool ShowColors, llvm::raw_ostream &out,
+                     llvm::raw_ostream::Colors Color) {
+  if (ShowColors)
+    out << llvm::sys::Process::OutputColor(Color, false, false);
+}
+
+static void resetColor(bool ShowColors, llvm::raw_ostream &out) {
+  if (ShowColors)
+    out << llvm::sys::Process::ResetColor();
+}
+
+static clang::DiagnosticsEngine::Level getLevel(unsigned DiagID) {
+  // FIXME: This feels like a hack.
+  static clang::DiagnosticsEngine Diags(new DiagnosticIDs,
+                                        new DiagnosticOptions);
+  return Diags.getDiagnosticLevel(DiagID, SourceLocation());
+}
+
 static void printGroup(llvm::raw_ostream &out, const GroupRecord &Group,
                        bool FlagsOnly, unsigned Indent = 0) {
   out.indent(Indent * 2);
+  
+  bool ShowColors = showColors(out);
+  setColor(ShowColors, out, llvm::raw_ostream::YELLOW);
   out << "-W" << Group.getName() << "\n";
-
+  resetColor(ShowColors, out);
+  
   ++Indent;
   for (GroupRecord::subgroup_iterator I = Group.subgroup_begin(),
                                       E = Group.subgroup_end();
@@ -47,8 +77,15 @@ static void printGroup(llvm::raw_ostream &out, const GroupRecord &Group,
     for (GroupRecord::diagnostics_iterator I = Group.diagnostics_begin(),
                                            E = Group.diagnostics_end();
          I != E; ++I) {
+      if (ShowColors) {
+        if (getLevel(I->DiagID) != DiagnosticsEngine::Ignored) {
+          setColor(ShowColors, out, llvm::raw_ostream::GREEN);
+        }
+      }
       out.indent(Indent * 2);
-      out << I->getName() << "\n";
+      out << I->getName();
+      resetColor(ShowColors, out);
+      out << "\n";
     }
   }
 }
@@ -106,7 +143,7 @@ int TreeView::run(unsigned int argc, char **argv, llvm::raw_ostream &out) {
       ++argv;
     }
   }
-
+  
   bool ShowAll = false;
   StringRef RootGroup;
 
@@ -127,6 +164,14 @@ int TreeView::run(unsigned int argc, char **argv, llvm::raw_ostream &out) {
     return -1;
   }
 
+  if (showColors(out)) {
+    out << '\n';
+    setColor(true, out, llvm::raw_ostream::GREEN);
+    out << "GREEN";
+    resetColor(true, out);
+    out << " = enabled by default\n\n";
+  }
+  
   if (ShowAll)
     return showAll(out, FlagsOnly);