Merge pull request #19609 from jrose-apple/your-code-is-programmed-not-to-love-me-but-you-cant-pretend

[Stats] Track NumASTBytesAllocated continuously

Author: jrose-apple

include/swift/AST/ASTContext.h

@@ -208,6 +208,12 @@ class ASTContext final {
                          DiagnosticEngine &Diags);
   ~ASTContext();
 
+  /// Optional table of counters to report, nullptr when not collecting.
+  ///
+  /// This must be initialized early so that Allocate() doesn't try to access
+  /// it before being set to null.
+  UnifiedStatsReporter *Stats = nullptr;
+
   /// \brief The language options used for translation.
   LangOptions &LangOpts;
 
@@ -261,12 +267,6 @@ class ASTContext final {
   /// Cache of remapped types (useful for diagnostics).
   llvm::StringMap<Type> RemappedTypes;
 
-  /// Optional table of counters to report, nullptr when not collecting.
-  UnifiedStatsReporter *Stats = nullptr;
-
-  /// Set a new stats reporter.
-  void setStatsReporter(UnifiedStatsReporter *stats);
-
 private:
   /// \brief The current generation number, which reflects the number of
   /// times that external modules have been loaded.
@@ -287,11 +287,11 @@ class ASTContext final {
   /// Cache of module names that fail the 'canImport' test in this context.
   llvm::SmallPtrSet<Identifier, 8> FailedModuleImportNames;
 
-public:
   /// \brief Retrieve the allocator for the given arena.
   llvm::BumpPtrAllocator &
   getAllocator(AllocationArena arena = AllocationArena::Permanent) const;
 
+public:
   /// Allocate - Allocate memory from the ASTContext bump pointer.
   void *Allocate(unsigned long bytes, unsigned alignment,
                  AllocationArena arena = AllocationArena::Permanent) const {
@@ -300,7 +300,9 @@ class ASTContext final {
 
     if (LangOpts.UseMalloc)
       return AlignedAlloc(bytes, alignment);
-    
+
+    if (arena == AllocationArena::Permanent && Stats)
+      Stats->getFrontendCounters().NumASTBytesAllocated += bytes;
     return getAllocator(arena).Allocate(bytes, alignment);
   }
 
@@ -395,6 +397,9 @@ class ASTContext final {
   /// Retrive the syntax node memory manager for this context.
   llvm::IntrusiveRefCntPtr<syntax::SyntaxArena> getSyntaxArena() const;
 
+  /// Set a new stats reporter.
+  void setStatsReporter(UnifiedStatsReporter *stats);
+
   /// Retrieve the lazy resolver for this context.
   LazyResolver *getLazyResolver() const;
 

lib/AST/ASTContext.cpp

@@ -562,6 +562,11 @@ llvm::BumpPtrAllocator &ASTContext::getAllocator(AllocationArena arena) const {
 void ASTContext::setStatsReporter(UnifiedStatsReporter *stats) {
   Stats = stats;
   evaluator.setStatsReporter(stats);
+
+  if (stats) {
+    stats->getFrontendCounters().NumASTBytesAllocated =
+        getAllocator().getBytesAllocated();
+  }
 }
 
 RC<syntax::SyntaxArena> ASTContext::getSyntaxArena() const {

lib/FrontendTool/FrontendTool.cpp

@@ -499,7 +499,6 @@ static void countStatsPostSema(UnifiedStatsReporter &Stats,
   auto const &AST = Instance.getASTContext();
   C.NumLoadedModules = AST.LoadedModules.size();
   C.NumImportedExternalDefinitions = AST.ExternalDefinitions.size();
-  C.NumASTBytesAllocated = AST.getAllocator().getBytesAllocated();
 
   if (auto *D = Instance.getDependencyTracker()) {
     C.NumDependencies = D->getDependencies().size();