Store ASTUnit in PrettyPrint instead of SourceManager

This allows us to check for more stuff in PrettyPrint.

Signed-off-by: Giuliano Belinassi <[email protected]>

libcextract/Passes.cpp

@@ -86,7 +86,7 @@ static bool Build_ASTUnit(PassManager::Context *ctx, IntrusiveRefCntPtr<vfs::Fil
     return false;
   }
 
-  PrettyPrint::Set_Source_Manager(&AU->getSourceManager());
+  PrettyPrint::Set_AST(AU.get());
   ctx->AST = std::move(AU);
 
   return true;
@@ -171,16 +171,11 @@ class BuildASTPass : public Pass
     std::error_code ec;
     llvm::raw_fd_ostream out(Get_Dump_Name_From_Input(ctx), ec);
 
-    PrettyPrint::Set_Source_Manager(nullptr);
-    PrettyPrint::Set_Output_Ostream(&out);
-
     for (it = ctx->AST->top_level_begin(); it != ctx->AST->top_level_end(); ++it) {
       Decl *decl = *it;
       PrettyPrint::Print_Decl(decl);
     }
 
-    PrettyPrint::Set_Source_Manager(&ctx->AST->getSourceManager());
-    PrettyPrint::Set_Output_Ostream(nullptr);
     out.close();
   }
 
@@ -458,7 +453,7 @@ class FunctionExternalizerPass : public Pass
          and set its new SourceManager to the PrettyPrint class.  */
       ctx->AST->Reparse(std::make_shared<PCHContainerOperations>(),
                         ClangCompat_None, ctx->OFS);
-      PrettyPrint::Set_Source_Manager(&ctx->AST->getSourceManager());
+      PrettyPrint::Set_AST(ctx->AST.get());
 
       if (ctx->Ibt && ctx->AST) {
         /* Do a sanity check on IBT macros.  Some kernel branches can't use it,

libcextract/PrettyPrint.cpp

@@ -68,7 +68,7 @@ void PrettyPrint::Print_Decl_Raw(Decl *decl)
 
   /* If a SourceManager was specified and the decl source range seems valid,
      then output based on the original source code.  */
-  if (SM && Is_Range_Valid(decl_range)) {
+  if (Is_Range_Valid(decl_range)) {
 
     /* The range given by getSourceRange ignores attributes. Example :
 
@@ -215,28 +215,26 @@ void PrettyPrint::Print_RawComment(SourceManager &sm, RawComment *comment)
 
 StringRef PrettyPrint::Get_Source_Text(const SourceRange &range)
 {
-    /* Calling this function supposes that a SourceManager was given to this class.  */
-    assert(SM && "A SourceManager wasn't passed to PrettyPrint.");
-
     // NOTE: sm.getSpellingLoc() used in case the range corresponds to a macro/preprocessed source.
     // NOTE2: getSpellingLoc() breaks in the case where a macro was asigned to be expanded to typedef.
+    SourceManager &SM = AST->getSourceManager();
     auto start_loc = range.getBegin();//SM->getSpellingLoc(range.getBegin());
     auto last_token_loc = range.getEnd();//SM->getSpellingLoc(range.getEnd());
-    auto end_loc = clang::Lexer::getLocForEndOfToken(last_token_loc, 0, *SM, LangOpts);
+    auto end_loc = clang::Lexer::getLocForEndOfToken(last_token_loc, 0, SM, LangOpts);
     auto printable_range = clang::SourceRange{start_loc, end_loc};
     return Get_Source_Text_Raw(printable_range);
 }
 
 StringRef PrettyPrint::Get_Source_Text_Raw(const SourceRange &range)
 {
-    return clang::Lexer::getSourceText(CharSourceRange::getCharRange(range), *SM, LangOpts);
+    SourceManager &SM = AST->getSourceManager();
+    return clang::Lexer::getSourceText(CharSourceRange::getCharRange(range), SM, LangOpts);
 }
 
 /** Compare if SourceLocation a is before SourceLocation b in the source code.  */
 bool PrettyPrint::Is_Before(const SourceLocation &a, const SourceLocation &b)
 {
-  assert(SM && "No SourceManager were given");
-  BeforeThanCompare<SourceLocation> is_before(*SM);
+  BeforeThanCompare<SourceLocation> is_before(AST->getSourceManager());
 
   assert(a.isValid());
   assert(b.isValid());
@@ -246,15 +244,16 @@ bool PrettyPrint::Is_Before(const SourceLocation &a, const SourceLocation &b)
 
 void PrettyPrint::Debug_SourceLoc(const SourceLocation &loc)
 {
-  loc.dump(*SM);
+  loc.dump(AST->getSourceManager());
 }
 
 bool PrettyPrint::Contains_From_LineCol(const SourceRange &a, const SourceRange &b)
 {
-  PresumedLoc a_begin = SM->getPresumedLoc(a.getBegin());
-  PresumedLoc a_end   = SM->getPresumedLoc(a.getEnd());
-  PresumedLoc b_begin = SM->getPresumedLoc(b.getBegin());
-  PresumedLoc b_end   = SM->getPresumedLoc(b.getEnd());
+  SourceManager &SM = AST->getSourceManager();
+  PresumedLoc a_begin = SM.getPresumedLoc(a.getBegin());
+  PresumedLoc a_end   = SM.getPresumedLoc(a.getEnd());
+  PresumedLoc b_begin = SM.getPresumedLoc(b.getBegin());
+  PresumedLoc b_end   = SM.getPresumedLoc(b.getEnd());
 
   assert(a_begin.getFileID() == a_end.getFileID());
   assert(b_begin.getFileID() == b_end.getFileID());
@@ -292,8 +291,7 @@ bool PrettyPrint::Contains(const SourceRange &a, const SourceRange &b)
 /** Compare if SourceLocation a is after SourceLocation b in the source code.  */
 bool PrettyPrint::Is_After(const SourceLocation &a, const SourceLocation &b)
 {
-  assert(SM && "No SourceManager were given");
-  BeforeThanCompare<SourceLocation> is_before(*SM);
+  BeforeThanCompare<SourceLocation> is_before(AST->getSourceManager());
   return is_before(b, a);
 }
 
@@ -304,7 +302,7 @@ SourceLocation PrettyPrint::Get_Expanded_Loc(Decl *decl)
 
   /* If a SourceManager was specified and the decl source range seems valid,
      then output based on the original source code.  */
-  if (SM && Is_Range_Valid(decl_range)) {
+  if (Is_Range_Valid(decl_range)) {
 
     /* The range given by getSourceRange ignores attributes. Example :
 
@@ -319,11 +317,12 @@ SourceLocation PrettyPrint::Get_Expanded_Loc(Decl *decl)
 
     AttrVec &attrvec = decl->getAttrs();
     bool has_attr = false;
+    SourceManager &SM = AST->getSourceManager();
 
     for (size_t i = 0; i < attrvec.size(); i++) {
       const Attr *attr = attrvec[i];
       SourceLocation loc = attr->getRange().getEnd();
-      loc = SM->getExpansionLoc(loc);
+      loc = SM.getExpansionLoc(loc);
 
       if (loc.isValid() && Is_Before(furthest, loc)) {
         furthest = loc;
@@ -341,7 +340,7 @@ SourceLocation PrettyPrint::Get_Expanded_Loc(Decl *decl)
 
       /* Keep fetching tokens.  */
       while (true) {
-        auto maybe_next_tok = Lexer::findNextToken(head, *SM, LangOpts);
+        auto maybe_next_tok = Lexer::findNextToken(head, SM, LangOpts);
         Token *tok = ClangCompat_GetTokenPtr(maybe_next_tok);
 
         if (tok == nullptr) {
@@ -375,19 +374,20 @@ void PrettyPrint::Set_Output_To(const std::string &path)
 
 StringRef PrettyPrint::Get_Filename_From_Loc(const SourceLocation &loc)
 {
-  return SM->getFilename(loc);
+  return AST->getSourceManager().getFilename(loc);
 }
 
 OptionalFileEntryRef PrettyPrint::Get_FileEntry(const SourceLocation &loc)
 {
-  return SM->getFileEntryRefForID(SM->getFileID(loc));
+  SourceManager &SM = AST->getSourceManager();
+  return SM.getFileEntryRefForID(SM.getFileID(loc));
 }
 
 /* See PrettyPrint.hh for what they do.  */
 raw_ostream *PrettyPrint::Out = &llvm::outs();
 LangOptions PrettyPrint::LangOpts;
 PrintingPolicy PrettyPrint::PPolicy(LangOpts);
-SourceManager *PrettyPrint::SM;
+ASTUnit *PrettyPrint::AST;
 
 
 

libcextract/PrettyPrint.hh

@@ -80,14 +80,14 @@ class PrettyPrint
 
   static bool Contains(const SourceRange &a, const SourceRange &b);
 
-  static inline void Set_Source_Manager(SourceManager *sm)
+  static inline void Set_AST(ASTUnit *ast)
   {
-    SM = sm;
+    AST = ast;
   }
 
   static inline SourceManager *Get_Source_Manager(void)
   {
-    return SM;
+    return &AST->getSourceManager();
   }
 
   static inline LangOptions &Get_Lang_Options(void)
@@ -153,9 +153,10 @@ class PrettyPrint
   /** Policy for printing.  We use the default for now.  */
   static PrintingPolicy PPolicy;
 
-  /** SourceManager built when parsing the AST.  Must be set after constructing
-      the ast by calling Set_Source_Manager.  */
-  static SourceManager *SM;
+  /** ASTUnit object.  Must be set after constructing the ast by
+      calling Set_Source_Manager.  FIXME: This should not be a statc
+      class variable.  Refactor this class entirely.  */
+  static ASTUnit *AST;
 
   friend class RecursivePrint;
 };