[MC] Eliminate two symbol-related hash maps (#95464)

Previously, a symbol insertion requires (at least) three hash table
operations:

- Lookup/create entry in Symbols (main symbol table)
- Lookup NextUniqueID to deduplicate identical temporary labels
- Add entry to UsedNames, which is also used to serve as storage for the
symbol name in the MCSymbol.

All three lookups are done with the same name, so combining these into a
single table reduces the number of lookups to one. Thus, a pointer to a
symbol table entry can be passed to createSymbol to avoid a duplicate
lookup of the same name.

The new symbol table entry value is placed in a separate header to avoid
including MCContext in MCSymbol or vice versa.

4 files changed, 69 insertions, 68 deletions

diff --git a/llvm/lib/MC/MCContext.cpp b/llvm/lib/MC/MCContext.cpp
index 432ff5e7..9c7cb74 100644
--- a/llvm/lib/MC/MCContext.cpp
+++ b/llvm/lib/MC/MCContext.cpp
@@ -70,7 +70,7 @@ MCContext::MCContext(const Triple &TheTriple, const MCAsmInfo *mai,
                      bool DoAutoReset, StringRef Swift5ReflSegmentName)
     : Swift5ReflectionSegmentName(Swift5ReflSegmentName), TT(TheTriple),
       SrcMgr(mgr), InlineSrcMgr(nullptr), DiagHandler(defaultDiagHandler),
-      MAI(mai), MRI(mri), MSTI(msti), Symbols(Allocator), UsedNames(Allocator),
+      MAI(mai), MRI(mri), MSTI(msti), Symbols(Allocator),
       InlineAsmUsedLabelNames(Allocator),
       CurrentDwarfLoc(0, 0, 0, DWARF2_FLAG_IS_STMT, 0, 0),
       AutoReset(DoAutoReset), TargetOptions(TargetOpts) {
@@ -153,7 +153,6 @@ void MCContext::reset() {
 
   MCSubtargetAllocator.DestroyAll();
   InlineAsmUsedLabelNames.clear();
-  UsedNames.clear();
   Symbols.clear();
   Allocator.Reset();
   Instances.clear();
@@ -179,7 +178,6 @@ void MCContext::reset() {
   ELFEntrySizeMap.clear();
   ELFSeenGenericMergeableSections.clear();
 
-  NextID.clear();
   DwarfLocSeen = false;
   GenDwarfForAssembly = false;
   GenDwarfFileNumber = 0;
@@ -214,11 +212,21 @@ MCSymbol *MCContext::getOrCreateSymbol(const Twine &Name) {
 
   assert(!NameRef.empty() && "Normal symbols cannot be unnamed!");
 
-  MCSymbol *&Sym = Symbols[NameRef];
-  if (!Sym)
-    Sym = createSymbol(NameRef, false, false);
+  MCSymbolTableEntry &Entry = getSymbolTableEntry(NameRef);
+  if (!Entry.second.Symbol) {
+    bool IsRenamable = NameRef.starts_with(MAI->getPrivateGlobalPrefix());
+    bool IsTemporary = IsRenamable && !SaveTempLabels;
+    if (!Entry.second.Used) {
+      Entry.second.Used = true;
+      Entry.second.Symbol = createSymbolImpl(&Entry, IsTemporary);
+    } else {
+      assert(IsRenamable && "cannot rename non-private symbol");
+      // Slow path: we need to rename a temp symbol from the user.
+      Entry.second.Symbol = createRenamableSymbol(NameRef, false, IsTemporary);
+    }
+  }
 
-  return Sym;
+  return Entry.second.Symbol;
 }
 
 MCSymbol *MCContext::getOrCreateFrameAllocSymbol(const Twine &FuncName,
@@ -237,7 +245,11 @@ MCSymbol *MCContext::getOrCreateLSDASymbol(const Twine &FuncName) {
                            FuncName);
 }
 
-MCSymbol *MCContext::createSymbolImpl(const StringMapEntry<bool> *Name,
+MCSymbolTableEntry &MCContext::getSymbolTableEntry(StringRef Name) {
+  return *Symbols.try_emplace(Name, MCSymbolTableValue{}).first;
+}
+
+MCSymbol *MCContext::createSymbolImpl(const MCSymbolTableEntry *Name,
                                       bool IsTemporary) {
   static_assert(std::is_trivially_destructible<MCSymbolCOFF>(),
                 "MCSymbol classes must be trivially destructible");
@@ -273,49 +285,37 @@ MCSymbol *MCContext::createSymbolImpl(const StringMapEntry<bool> *Name,
       MCSymbol(MCSymbol::SymbolKindUnset, Name, IsTemporary);
 }
 
-MCSymbol *MCContext::createSymbol(StringRef Name, bool AlwaysAddSuffix,
-                                  bool IsTemporary) {
-  // Determine whether this is a user written assembler temporary or normal
-  // label, if used.
-  if (!SaveTempLabels && !IsTemporary)
-    IsTemporary = Name.starts_with(MAI->getPrivateGlobalPrefix());
-
-  SmallString<128> NewName = Name;
-  bool AddSuffix = AlwaysAddSuffix;
-  unsigned &NextUniqueID = NextID[Name];
-  while (true) {
-    if (AddSuffix) {
-      NewName.resize(Name.size());
-      raw_svector_ostream(NewName) << NextUniqueID++;
-    }
-    auto NameEntry = UsedNames.insert(std::make_pair(NewName.str(), true));
-    if (NameEntry.second || !NameEntry.first->second) {
-      // Ok, we found a name.
-      // Mark it as used for a non-section symbol.
-      NameEntry.first->second = true;
-      // Have the MCSymbol object itself refer to the copy of the string that is
-      // embedded in the UsedNames entry.
-      return createSymbolImpl(&*NameEntry.first, IsTemporary);
-    }
-    assert(IsTemporary && "Cannot rename non-temporary symbols");
-    AddSuffix = true;
+MCSymbol *MCContext::createRenamableSymbol(const Twine &Name,
+                                           bool AlwaysAddSuffix,
+                                           bool IsTemporary) {
+  SmallString<128> NewName;
+  Name.toVector(NewName);
+  size_t NameLen = NewName.size();
+
+  MCSymbolTableEntry &NameEntry = getSymbolTableEntry(NewName.str());
+  MCSymbolTableEntry *EntryPtr = &NameEntry;
+  while (AlwaysAddSuffix || EntryPtr->second.Used) {
+    AlwaysAddSuffix = false;
+
+    NewName.resize(NameLen);
+    raw_svector_ostream(NewName) << NameEntry.second.NextUniqueID++;
+    EntryPtr = &getSymbolTableEntry(NewName.str());
   }
-  llvm_unreachable("Infinite loop");
+
+  EntryPtr->second.Used = true;
+  return createSymbolImpl(EntryPtr, IsTemporary);
 }
 
 MCSymbol *MCContext::createTempSymbol(const Twine &Name, bool AlwaysAddSuffix) {
   if (!UseNamesOnTempLabels)
     return createSymbolImpl(nullptr, /*IsTemporary=*/true);
-
-  SmallString<128> NameSV;
-  raw_svector_ostream(NameSV) << MAI->getPrivateGlobalPrefix() << Name;
-  return createSymbol(NameSV, AlwaysAddSuffix, true);
+  return createRenamableSymbol(MAI->getPrivateGlobalPrefix() + Name,
+                               AlwaysAddSuffix, /*IsTemporary=*/true);
 }
 
 MCSymbol *MCContext::createNamedTempSymbol(const Twine &Name) {
-  SmallString<128> NameSV;
-  raw_svector_ostream(NameSV) << MAI->getPrivateGlobalPrefix() << Name;
-  return createSymbol(NameSV, true, false);
+  return createRenamableSymbol(MAI->getPrivateGlobalPrefix() + Name, true,
+                               /*IsTemporary=*/!SaveTempLabels);
 }
 
 MCSymbol *MCContext::createLinkerPrivateTempSymbol() {
@@ -323,9 +323,9 @@ MCSymbol *MCContext::createLinkerPrivateTempSymbol() {
 }
 
 MCSymbol *MCContext::createLinkerPrivateSymbol(const Twine &Name) {
-  SmallString<128> NameSV;
-  raw_svector_ostream(NameSV) << MAI->getLinkerPrivateGlobalPrefix() << Name;
-  return createSymbol(NameSV, true, false);
+  return createRenamableSymbol(MAI->getLinkerPrivateGlobalPrefix() + Name,
+                               /*AlwaysAddSuffix=*/true,
+                               /*IsTemporary=*/false);
 }
 
 MCSymbol *MCContext::createTempSymbol() { return createTempSymbol("tmp"); }
@@ -372,7 +372,7 @@ MCSymbol *MCContext::getDirectionalLocalSymbol(unsigned LocalLabelVal,
 MCSymbol *MCContext::lookupSymbol(const Twine &Name) const {
   SmallString<128> NameSV;
   StringRef NameRef = Name.toStringRef(NameSV);
-  return Symbols.lookup(NameRef);
+  return Symbols.lookup(NameRef).Symbol;
 }
 
 void MCContext::setSymbolValue(MCStreamer &Streamer, const Twine &Sym,
@@ -389,9 +389,8 @@ wasm::WasmSignature *MCContext::createWasmSignature() {
   return new (WasmSignatureAllocator.Allocate()) wasm::WasmSignature;
 }
 
-MCSymbolXCOFF *
-MCContext::createXCOFFSymbolImpl(const StringMapEntry<bool> *Name,
-                                 bool IsTemporary) {
+MCSymbolXCOFF *MCContext::createXCOFFSymbolImpl(const MCSymbolTableEntry *Name,
+                                                bool IsTemporary) {
   if (!Name)
     return new (nullptr, *this) MCSymbolXCOFF(nullptr, IsTemporary);
 
@@ -431,15 +430,13 @@ MCContext::createXCOFFSymbolImpl(const StringMapEntry<bool> *Name,
   else
     ValidName.append(InvalidName);
 
-  auto NameEntry = UsedNames.insert(std::make_pair(ValidName.str(), true));
-  assert((NameEntry.second || !NameEntry.first->second) &&
-         "This name is used somewhere else.");
-  // Mark the name as used for a non-section symbol.
-  NameEntry.first->second = true;
+  MCSymbolTableEntry &NameEntry = getSymbolTableEntry(ValidName.str());
+  assert(!NameEntry.second.Used && "This name is used somewhere else.");
+  NameEntry.second.Used = true;
   // Have the MCSymbol object itself refer to the copy of the string
-  // that is embedded in the UsedNames entry.
-  MCSymbolXCOFF *XSym = new (&*NameEntry.first, *this)
-      MCSymbolXCOFF(&*NameEntry.first, IsTemporary);
+  // that is embedded in the symbol table entry.
+  MCSymbolXCOFF *XSym =
+      new (&NameEntry, *this) MCSymbolXCOFF(&NameEntry, IsTemporary);
   XSym->setSymbolTableName(MCSymbolXCOFF::getUnqualifiedName(OriginalName));
   return XSym;
 }
@@ -485,7 +482,8 @@ MCSectionELF *MCContext::createELFSectionImpl(StringRef Section, unsigned Type,
                                               bool Comdat, unsigned UniqueID,
                                               const MCSymbolELF *LinkedToSym) {
   MCSymbolELF *R;
-  MCSymbol *&Sym = Symbols[Section];
+  MCSymbolTableEntry &SymEntry = getSymbolTableEntry(Section);
+  MCSymbol *Sym = SymEntry.second.Symbol;
   // A section symbol can not redefine regular symbols. There may be multiple
   // sections with the same name, in which case the first such section wins.
   if (Sym && Sym->isDefined() &&
@@ -494,10 +492,10 @@ MCSectionELF *MCContext::createELFSectionImpl(StringRef Section, unsigned Type,
   if (Sym && Sym->isUndefined()) {
     R = cast<MCSymbolELF>(Sym);
   } else {
-    auto NameIter = UsedNames.insert(std::make_pair(Section, false)).first;
-    R = new (&*NameIter, *this) MCSymbolELF(&*NameIter, /*isTemporary*/ false);
+    SymEntry.second.Used = true;
+    R = new (&SymEntry, *this) MCSymbolELF(&SymEntry, /*isTemporary*/ false);
     if (!Sym)
-      Sym = R;
+      SymEntry.second.Symbol = R;
   }
   R->setBinding(ELF::STB_LOCAL);
   R->setType(ELF::STT_SECTION);
@@ -750,8 +748,9 @@ MCSectionWasm *MCContext::getWasmSection(const Twine &Section, SectionKind Kind,
 
   StringRef CachedName = Entry.first.SectionName;
 
-  MCSymbol *Begin = createSymbol(CachedName, true, false);
-  Symbols[Begin->getName()] = Begin;
+  MCSymbol *Begin = createRenamableSymbol(CachedName, true, false);
+  // Begin always has a different name than CachedName... see #48596.
+  getSymbolTableEntry(Begin->getName()).second.Symbol = Begin;
   cast<MCSymbolWasm>(Begin)->setType(wasm::WASM_SYMBOL_TYPE_SECTION);
 
   MCSectionWasm *Result = new (WasmAllocator.Allocate())
diff --git a/llvm/lib/MC/MCParser/AsmParser.cpp b/llvm/lib/MC/MCParser/AsmParser.cpp
index ddbf99b..6e7ec68 100644
--- a/llvm/lib/MC/MCParser/AsmParser.cpp
+++ b/llvm/lib/MC/MCParser/AsmParser.cpp
@@ -1044,11 +1044,12 @@ bool AsmParser::Run(bool NoInitialTextSection, bool NoFinalize) {
   if (!NoFinalize) {
     if (MAI.hasSubsectionsViaSymbols()) {
       for (const auto &TableEntry : getContext().getSymbols()) {
-        MCSymbol *Sym = TableEntry.getValue();
+        MCSymbol *Sym = TableEntry.getValue().Symbol;
         // Variable symbols may not be marked as defined, so check those
         // explicitly. If we know it's a variable, we have a definition for
         // the purposes of this check.
-        if (Sym->isTemporary() && !Sym->isVariable() && !Sym->isDefined())
+        if (Sym && Sym->isTemporary() && !Sym->isVariable() &&
+            !Sym->isDefined())
           // FIXME: We would really like to refer back to where the symbol was
           // first referenced for a source location. We need to add something
           // to track that. Currently, we just point to the end of the file.
diff --git a/llvm/lib/MC/MCParser/MasmParser.cpp b/llvm/lib/MC/MCParser/MasmParser.cpp
index 5ed66fb..67f3540 100644
--- a/llvm/lib/MC/MCParser/MasmParser.cpp
+++ b/llvm/lib/MC/MCParser/MasmParser.cpp
@@ -1421,11 +1421,12 @@ bool MasmParser::Run(bool NoInitialTextSection, bool NoFinalize) {
   if (!NoFinalize) {
     if (MAI.hasSubsectionsViaSymbols()) {
       for (const auto &TableEntry : getContext().getSymbols()) {
-        MCSymbol *Sym = TableEntry.getValue();
+        MCSymbol *Sym = TableEntry.getValue().Symbol;
         // Variable symbols may not be marked as defined, so check those
         // explicitly. If we know it's a variable, we have a definition for
         // the purposes of this check.
-        if (Sym->isTemporary() && !Sym->isVariable() && !Sym->isDefined())
+        if (Sym && Sym->isTemporary() && !Sym->isVariable() &&
+            !Sym->isDefined())
           // FIXME: We would really like to refer back to where the symbol was
           // first referenced for a source location. We need to add something
           // to track that. Currently, we just point to the end of the file.
diff --git a/llvm/lib/MC/MCSymbol.cpp b/llvm/lib/MC/MCSymbol.cpp
index 4017225..867a472 100644
--- a/llvm/lib/MC/MCSymbol.cpp
+++ b/llvm/lib/MC/MCSymbol.cpp
@@ -27,7 +27,7 @@ static MCDummyFragment SentinelFragment(nullptr);
 // Sentinel value for the absolute pseudo fragment.
 MCFragment *MCSymbol::AbsolutePseudoFragment = &SentinelFragment;
 
-void *MCSymbol::operator new(size_t s, const StringMapEntry<bool> *Name,
+void *MCSymbol::operator new(size_t s, const MCSymbolTableEntry *Name,
                              MCContext &Ctx) {
   // We may need more space for a Name to account for alignment.  So allocate
   // space for the storage type and not the name pointer.