diff --git a/tools/wpa/clang-wpa.cpp b/tools/wpa/clang-wpa.cpp
index e7515ec..2cdc1c6 100644
--- a/tools/wpa/clang-wpa.cpp
+++ b/tools/wpa/clang-wpa.cpp
@@ -15,7 +15,6 @@
 #include "clang/Analysis/CallGraph.h"
 
 #include "clang/Basic/FileManager.h"
-#include "clang/Index/TranslationUnit.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace clang;
@@ -24,20 +23,10 @@
 static llvm::cl::list<std::string>
 InputFilenames(llvm::cl::Positional, llvm::cl::desc("<input AST files>"));
 
-// FIXME: this duplicates the one in index-test.cpp.
-class TUnit : public TranslationUnit {
-public:
-  TUnit(ASTUnit *ast, const std::string &filename)
-    : AST(ast), Filename(filename) {}
-  ASTContext &getASTContext() { return AST->getASTContext(); }
-  llvm::OwningPtr<ASTUnit> AST;
-  std::string Filename;
-};
-
 int main(int argc, char **argv) {
   llvm::cl::ParseCommandLineOptions(argc, argv, "clang-wpa");
   FileManager FileMgr;
-  std::vector<TUnit*> TUnits;
+  std::vector<ASTUnit*> ASTUnits;
 
   if (InputFilenames.empty())
     return 0;
@@ -55,15 +44,14 @@
       return 1;
     }
 
-    TUnit *TU = new TUnit(AST.take(), InFile);
-    TUnits.push_back(TU);
+    ASTUnits.push_back(AST.take());
   }
 
   llvm::OwningPtr<CallGraph> CG;
   CG.reset(new CallGraph());
 
-  for (unsigned i = 0, e = TUnits.size(); i != e; ++i)
-    CG->addTU(*(TUnits[i]->AST));
+  for (unsigned i = 0, e = ASTUnits.size(); i != e; ++i)
+    CG->addTU(*ASTUnits[i]);
 
   CG->ViewCallGraph();
 }