From 8cb6e587fd40b97983e445ee3f17f4c6d7190df7 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Sat, 22 Jun 2024 12:54:19 -0700
Subject: [MC] Allocate MCFragment with a bump allocator

#95197 and 75006466296ed4b0f845cbbec4bf77c21de43b40 eliminated all raw
`new MCXXXFragment`. We can now place fragments in a bump allocator.
In addition, remove the dead `Kind == FragmentType(~0)` condition.

~CodeViewContext may call `StrTabFragment->destroy()` and need to be
reset before `FragmentAllocator.Reset()`.
Tested by llvm/test/MC/COFF/cv-compiler-info.ll using asan.

Pull Request: https://github.com/llvm/llvm-project/pull/96402
---
 llvm/include/llvm/MC/MCContext.h |  6 +++++-
 llvm/lib/MC/MCCodeView.cpp       |  4 ++--
 llvm/lib/MC/MCContext.cpp        |  7 +++++--
 llvm/lib/MC/MCFragment.cpp       | 38 ++++++++++++++++----------------------
 4 files changed, 28 insertions(+), 27 deletions(-)

diff --git a/llvm/include/llvm/MC/MCContext.h b/llvm/include/llvm/MC/MCContext.h
index 6c977a5..0f18ce5 100644
--- a/llvm/include/llvm/MC/MCContext.h
+++ b/llvm/include/llvm/MC/MCContext.h
@@ -136,6 +136,9 @@ private:
   /// objects.
   BumpPtrAllocator Allocator;
 
+  /// For MCFragment instances.
+  BumpPtrAllocator FragmentAllocator;
+
   SpecificBumpPtrAllocator<MCSectionCOFF> COFFAllocator;
   SpecificBumpPtrAllocator<MCSectionDXContainer> DXCAllocator;
   SpecificBumpPtrAllocator<MCSectionELF> ELFAllocator;
@@ -432,7 +435,8 @@ public:
   MCInst *createMCInst();
 
   template <typename F, typename... Args> F *allocFragment(Args &&...args) {
-    return new F(std::forward<Args>(args)...);
+    return new (FragmentAllocator.Allocate(sizeof(F), alignof(F)))
+        F(std::forward<Args>(args)...);
   }
 
   /// \name Symbol Management
diff --git a/llvm/lib/MC/MCCodeView.cpp b/llvm/lib/MC/MCCodeView.cpp
index 713ae07..89b28b4 100644
--- a/llvm/lib/MC/MCCodeView.cpp
+++ b/llvm/lib/MC/MCCodeView.cpp
@@ -29,8 +29,8 @@ using namespace llvm::codeview;
 CodeViewContext::~CodeViewContext() {
   // If someone inserted strings into the string table but never actually
   // emitted them somewhere, clean up the fragment.
-  if (!InsertedStrTabFragment)
-    delete StrTabFragment;
+  if (!InsertedStrTabFragment && StrTabFragment)
+    StrTabFragment->destroy();
 }
 
 /// This is a valid number for use with .cv_loc if we've already seen a .cv_file
diff --git a/llvm/lib/MC/MCContext.cpp b/llvm/lib/MC/MCContext.cpp
index bb87682..e7fc883 100644
--- a/llvm/lib/MC/MCContext.cpp
+++ b/llvm/lib/MC/MCContext.cpp
@@ -151,10 +151,15 @@ void MCContext::reset() {
   SPIRVAllocator.DestroyAll();
   WasmSignatureAllocator.DestroyAll();
 
+  // ~CodeViewContext may destroy a MCFragment outside of sections and need to
+  // be reset before FragmentAllocator.
+  CVContext.reset();
+
   MCSubtargetAllocator.DestroyAll();
   InlineAsmUsedLabelNames.clear();
   Symbols.clear();
   Allocator.Reset();
+  FragmentAllocator.Reset();
   Instances.clear();
   CompilationDir.clear();
   MainFileName.clear();
@@ -165,8 +170,6 @@ void MCContext::reset() {
   DwarfCompileUnitID = 0;
   CurrentDwarfLoc = MCDwarfLoc(0, 0, 0, DWARF2_FLAG_IS_STMT, 0, 0);
 
-  CVContext.reset();
-
   MachOUniquingMap.clear();
   ELFUniquingMap.clear();
   GOFFUniquingMap.clear();
diff --git a/llvm/lib/MC/MCFragment.cpp b/llvm/lib/MC/MCFragment.cpp
index d90fd34..b1ed594 100644
--- a/llvm/lib/MC/MCFragment.cpp
+++ b/llvm/lib/MC/MCFragment.cpp
@@ -202,60 +202,54 @@ MCFragment::MCFragment(FragmentType Kind, bool HasInstructions)
     : Kind(Kind), HasInstructions(HasInstructions), LinkerRelaxable(false) {}
 
 void MCFragment::destroy() {
-  // First check if we are the sentinel.
-  if (Kind == FragmentType(~0)) {
-    delete this;
-    return;
-  }
-
   switch (Kind) {
     case FT_Align:
-      delete cast<MCAlignFragment>(this);
+      cast<MCAlignFragment>(this)->~MCAlignFragment();
       return;
     case FT_Data:
-      delete cast<MCDataFragment>(this);
+      cast<MCDataFragment>(this)->~MCDataFragment();
       return;
     case FT_CompactEncodedInst:
-      delete cast<MCCompactEncodedInstFragment>(this);
+      cast<MCCompactEncodedInstFragment>(this)->~MCCompactEncodedInstFragment();
       return;
     case FT_Fill:
-      delete cast<MCFillFragment>(this);
+      cast<MCFillFragment>(this)->~MCFillFragment();
       return;
     case FT_Nops:
-      delete cast<MCNopsFragment>(this);
+      cast<MCNopsFragment>(this)->~MCNopsFragment();
       return;
     case FT_Relaxable:
-      delete cast<MCRelaxableFragment>(this);
+      cast<MCRelaxableFragment>(this)->~MCRelaxableFragment();
       return;
     case FT_Org:
-      delete cast<MCOrgFragment>(this);
+      cast<MCOrgFragment>(this)->~MCOrgFragment();
       return;
     case FT_Dwarf:
-      delete cast<MCDwarfLineAddrFragment>(this);
+      cast<MCDwarfLineAddrFragment>(this)->~MCDwarfLineAddrFragment();
       return;
     case FT_DwarfFrame:
-      delete cast<MCDwarfCallFrameFragment>(this);
+      cast<MCDwarfCallFrameFragment>(this)->~MCDwarfCallFrameFragment();
       return;
     case FT_LEB:
-      delete cast<MCLEBFragment>(this);
+      cast<MCLEBFragment>(this)->~MCLEBFragment();
       return;
     case FT_BoundaryAlign:
-      delete cast<MCBoundaryAlignFragment>(this);
+      cast<MCBoundaryAlignFragment>(this)->~MCBoundaryAlignFragment();
       return;
     case FT_SymbolId:
-      delete cast<MCSymbolIdFragment>(this);
+      cast<MCSymbolIdFragment>(this)->~MCSymbolIdFragment();
       return;
     case FT_CVInlineLines:
-      delete cast<MCCVInlineLineTableFragment>(this);
+      cast<MCCVInlineLineTableFragment>(this)->~MCCVInlineLineTableFragment();
       return;
     case FT_CVDefRange:
-      delete cast<MCCVDefRangeFragment>(this);
+      cast<MCCVDefRangeFragment>(this)->~MCCVDefRangeFragment();
       return;
     case FT_PseudoProbe:
-      delete cast<MCPseudoProbeAddrFragment>(this);
+      cast<MCPseudoProbeAddrFragment>(this)->~MCPseudoProbeAddrFragment();
       return;
     case FT_Dummy:
-      delete cast<MCDummyFragment>(this);
+      cast<MCDummyFragment>(this)->~MCDummyFragment();
       return;
   }
 }
-- 
cgit v1.1