Add Rewrite support.

Implements the rewrite flag so that it will rename the headers it finds.
loarabia · Jun 14, 2011 · 5d5a3d4 · 5d5a3d4
1 parent beb2220
commit 5d5a3d4
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 6 deletions.
diff --git a/makefile b/makefile
@@ -3,7 +3,8 @@ LLVMLDFLAGS := $(shell llvm-config --ldflags --libs)
 SOURCES = src/uh.cpp \
     src/IncludeHandler.cpp
 OBJECTS = $(SOURCES:.cpp=.o)
-CLANGLIBS = -lclangParse \
+CLANGLIBS = -lclangRewrite \
+    -lclangParse \
     -lclangSema \
     -lclangAnalysis \
     -lclangAST \

diff --git a/src/IncludeHandler.cpp b/src/IncludeHandler.cpp
@@ -1,6 +1,7 @@
 #include <iostream>
 
 #include "IncludeHandler.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace uh;
 
@@ -35,5 +36,28 @@ void IncludeHandler::InclusionDirective(SourceLocation hashLoc,
              std::cout << loc.getLine();
              std::cout << " ";
              std::cout << filename.str() << std::endl;
+             if ( rewrite )
+             {
+                 bool result = rewriter.ReplaceText(endLoc.getFileLocWithOffset(1),
+                    filename.size(),
+                    headerFilename);
+
+                 for (Rewriter::buffer_iterator I = rewriter.buffer_begin(),
+                      E = rewriter.buffer_end();
+                      I != E;
+                      ++I)
+                 {
+                      const FileEntry *fe = rewriter.getSourceMgr().getFileEntryForID(I->first);
+                      std::string Filename = fe->getName();
+                      std::string Err;
+                      llvm::raw_fd_ostream OS(Filename.c_str(), Err,
+                                              llvm::raw_fd_ostream::F_Binary);
+                      if (!Err.empty()) {
+                      }
+                      RewriteBuffer &RewriteBuf = I->second;
+                      RewriteBuf.write(OS);
+                      OS.flush();
+                 }
+             }
         }
 }
diff --git a/src/IncludeHandler.h b/src/IncludeHandler.h
@@ -4,6 +4,7 @@
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/SourceLocation.h"
+#include "clang/Rewrite/Rewriter.h"
 #include "clang/Lex/PPCallbacks.h"
 
 #include "llvm/ADT/StringRef.h"
@@ -21,14 +22,22 @@ namespace uh
         int includesFoundInFile;
         llvm::Regex headerRegex;
         llvm::StringRef headerFilename;
+        clang::Rewriter rewriter;
+        bool rewrite;
 
         public:
-            IncludeHandler( SourceManager &srcMgr, llvm::sys::Path file, std::string &header) :
+            IncludeHandler( SourceManager &srcMgr,
+                            llvm::sys::Path file, 
+                            std::string &header,
+                            clang::Rewriter &rw,
+                            bool rewrite) :
                 sm(srcMgr),
                 fileBeingParsed(file),
                 includesFoundInFile(0),
                 headerRegex(header, llvm::Regex::IgnoreCase),
-                headerFilename(header)
+                headerFilename(header),
+                rewriter(rw),
+                rewrite(rewrite)
                 {
                 }
 

diff --git a/src/uh.cpp b/src/uh.cpp
@@ -22,13 +22,14 @@
 #include "clang/AST/ASTConsumer.h"
 #include "clang/Basic/TargetOptions.h"
 #include "clang/Basic/TargetInfo.h"
+#include "clang/Rewrite/Rewriter.h"
 
 
 using namespace llvm;
 
 /*
- * If this flag is specified, the tool will actually rename and headers that
- * match the specified header. The fefault value is false, which just causes
+ * If this flag is specified, the tool will actually rename any headers that
+ * match the specified header. The default value is false, which just causes
  * the tool to print where a header was found.
  */ 
 static cl::opt<bool>        Rename(
@@ -109,6 +110,7 @@ std::set<sys::Path> FindFilesContainingHeaders()
     {
         ci.createSourceManager(ci.getFileManager());
         clang::SourceManager &sm = ci.getSourceManager();
+        clang::Rewriter rw(sm, lOpts);
 
         ci.createPreprocessor();
         clang::Preprocessor &pp = ci.getPreprocessor();
@@ -118,7 +120,9 @@ std::set<sys::Path> FindFilesContainingHeaders()
         uh::IncludeHandler *includeHandler = new uh::IncludeHandler(
             ci.getSourceManager(),
             *childIterator,
-            Filename);
+            Filename,
+            rw,
+            Rename);
 
         pp.addPPCallbacks(includeHandler);
         ci.setPreprocessor(&pp);