[clang] Support -clear-ast-before-backend without -disable-free

[aeubanks]

Summary

Previously without -disable-free, -clear-ast-before-backend would crash in ~ASTContext() due to various reasons.
This works around that by doing a lot of the cleanup ahead of the destructor so that the destructor doesn't actually do any manual cleanup if we've already cleaned up beforehand.

This actually does save a measurable amount of memory with -clear-ast-before-backend, although at an almost unnoticeable runtime cost:
https://llvm-compile-time-tracker.com/compare.php?from=5d755b32f2775b9219f6d6e2feda5e1417dc993b&to=58ef1c7ad7e2ad45f9c97597905a8cf05a26258c&stat=max-rss

Previously we weren't doing any cleanup with -disable-free, so I tried measuring the impact of always doing the cleanup and didn't measure anything noticeable on llvm-compile-time-tracker.

[dblaikie]

clang/lib/CodeGen/CodeGenAction.cpp

355–356

Any chance of refactoring such that the ASTContext is scoped such that it is destroyed here, rather than rendered unusable - to reduce the chance that it'd be accidentally used after this point?

Like what happens if CompilerInstance::setASTContext(nullptr) is called? It wouldn't null out everyone's ASTContext pointers, but would mean there's no object there - perhaps crash harder/faster than a "cleaned up" ASTContext being used?

[aeubanks]

clang/lib/CodeGen/CodeGenAction.cpp

355–356

I had tried this in the past and gave up for some reason. So I tried again (ASan was very helpful here) and ran into a couple of issues.

First, Sema depends on ASTContext, so we have to setSema(nullptr) before setASTContext(nullptr), but Sema has to finalize some things after codegen, so we can't CI->setSema(nullptr) until much later.

Another unrelated thing is that the diagnostics end up calling into ASTContext. This might be fixable by moving some data structures out of ASTContext?

Another unrelated thing is a weird crash in DiagStorageAllocator::~DiagStorageAllocator() when calling setASTContext(nullptr) early.

These individually might be solvable, but it seems like a lot of unrelated issues.
Basically, we depend a lot on various parts of ASTContext that aren't the actual AST nodes themselves.

[dblaikie]

Sounds OK, thanks!

clang/lib/CodeGen/CodeGenAction.cpp

355–356

Fair enough - not sure, but might be worth noting this info somewhere in the source in case someone else comes along to do this refactoring in the future if they have another reason to weigh into the choice to address it.

@rsmith @aaron.ballman - wouldn't mind at least a quick thought on this stuff if you've both got a moment - for long term direction/broad design of this stuff.

[aeubanks]

add more comments about the cleanup

clang/include/clang/AST/ASTContext.h

public:
 
   void setPrintingPolicy(const clang::PrintingPolicy &Policy) {
     PrintingPolicy = Policy;
   }
 
   SourceManager& getSourceManager() { return SourceMgr; }
   const SourceManager& getSourceManager() const { return SourceMgr; }
 
+  // Cleans up some of the data structures. This allows us to do cleanup
+  // normally done in the destructor earlier. Renders much of the ASTContext
+  // unusable, mostly the actual AST nodes, so should be called when we no
+  // longer need access to the AST.
+  void cleanup();
+
   llvm::BumpPtrAllocator &getAllocator() const {
     return BumpAlloc;
   }
 
   void *Allocate(size_t Size, unsigned Align = 8) const {
     return BumpAlloc.Allocate(Size, Align);
   }
   template <typename T> T *Allocate(size_t Num = 1) const {

clang/lib/AST/ASTContext.cpp

: ConstantArrayTypes(this_()), FunctionProtoTypes(this_()),
       ProfList(new ProfileList(LangOpts.ProfileListFiles, SM)),
       PrintingPolicy(LOpts), Idents(idents), Selectors(sels),
       BuiltinInfo(builtins), TUKind(TUKind), DeclarationNames(*this),
       Comments(SM), CommentCommandTraits(BumpAlloc, LOpts.CommentOpts),
       CompCategories(this_()), LastSDM(nullptr, 0) {
   addTranslationUnitDecl();
 }
 
-ASTContext::~ASTContext() {
+void ASTContext::cleanup() {
   // Release the DenseMaps associated with DeclContext objects.
   // FIXME: Is this the ideal solution?
   ReleaseDeclContextMaps();
 
   // Call all of the deallocation functions on all of their targets.
   for (auto &Pair : Deallocations)
     (Pair.first)(Pair.second);
+  Deallocations.clear();
 
   // ASTRecordLayout objects in ASTRecordLayouts must always be destroyed
   // because they can contain DenseMaps.
   for (llvm::DenseMap<const ObjCContainerDecl*,
        const ASTRecordLayout*>::iterator
        I = ObjCLayouts.begin(), E = ObjCLayouts.end(); I != E; )
     // Increment in loop to prevent using deallocated memory.
     if (auto *R = const_cast<ASTRecordLayout *>((I++)->second))
       R->Destroy(*this);
+  ObjCLayouts.clear();
 
   for (llvm::DenseMap<const RecordDecl*, const ASTRecordLayout*>::iterator
        I = ASTRecordLayouts.begin(), E = ASTRecordLayouts.end(); I != E; ) {
     // Increment in loop to prevent using deallocated memory.
     if (auto *R = const_cast<ASTRecordLayout *>((I++)->second))
       R->Destroy(*this);
   }
+  ASTRecordLayouts.clear();
 
   for (llvm::DenseMap<const Decl*, AttrVec*>::iterator A = DeclAttrs.begin(),
                                                     AEnd = DeclAttrs.end();
        A != AEnd; ++A)
     A->second->~AttrVec();
+  DeclAttrs.clear();
 
   for (const auto &Value : ModuleInitializers)
     Value.second->~PerModuleInitializers();
+  ModuleInitializers.clear();
 }
 
+ASTContext::~ASTContext() { cleanup(); }
+
 void ASTContext::setTraversalScope(const std::vector<Decl *> &TopLevelDecls) {
   TraversalScope = TopLevelDecls;
   getParentMapContext().clear();
 }
 
 void ASTContext::AddDeallocation(void (*Callback)(void *), void *Data) const {
   Deallocations.push_back({Callback, Data});
 }

clang/lib/AST/DeclBase.cpp

StoredDeclsMap *DeclContext::CreateStoredDeclsMap(ASTContext &C) const {
   return M;
 }
 
 void ASTContext::ReleaseDeclContextMaps() {
   // It's okay to delete DependentStoredDeclsMaps via a StoredDeclsMap
   // pointer because the subclass doesn't add anything that needs to
   // be deleted.
   StoredDeclsMap::DestroyAll(LastSDM.getPointer(), LastSDM.getInt());
+  LastSDM.setPointer(nullptr);
 }
 
 void StoredDeclsMap::DestroyAll(StoredDeclsMap *Map, bool Dependent) {
   while (Map) {
     // Advance the iteration before we invalidate memory.
     llvm::PointerIntPair<StoredDeclsMap*,1> Next = Map->Previous;
 
     if (Dependent)

clang/lib/CodeGen/CodeGenAction.cpp

void HandleTranslationUnit(ASTContext &C) override {
         if (const Decl *FD = Gen->GetDeclForMangledName(F.getName())) {
           auto Loc = FD->getASTContext().getFullLoc(FD->getLocation());
           // TODO: use a fast content hash when available.
           auto NameHash = llvm::hash_value(F.getName());
           ManglingFullSourceLocs.push_back(std::make_pair(NameHash, Loc));
         }
       }
 
-      // FIXME: Fix cleanup issues with clearing the AST when we properly free
-      // things.
-      if (CodeGenOpts.DisableFree && CodeGenOpts.ClearASTBeforeBackend)
+      if (CodeGenOpts.ClearASTBeforeBackend) {
+        // Access to the AST is no longer available after this.
+        // Other things that the ASTContext manages are still available, e.g.
+        // the SourceManager. It'd be nice if we could separate out all the
+        // things in ASTContext used after this point and null out the
+        // ASTContext, but too many various parts of the ASTContext are still
+        // used in various parts.
+        C.cleanup();
         C.getAllocator().Reset();
+      }
 
       EmbedBitcode(getModule(), CodeGenOpts, llvm::MemoryBufferRef());
 
       EmitBackendOutput(Diags, HeaderSearchOpts, CodeGenOpts, TargetOpts,
                         LangOpts, C.getTargetInfo().getDataLayoutString(),
                         getModule(), Action, std::move(AsmOutStream));
 
       Ctx.setDiagnosticHandler(std::move(OldDiagnosticHandler));