[NFC][XCOFF][AIX] Refactor get/setContainingCsect

Summary:
For current architect, we always require setContainingCsect to be
called on every MCSymbol got used in XCOFF context.
This is very hard to achieve because symbols gets created everywhere
 and other MCSymbol types(ELF, COFF) do not have similar rules.
It's very easy to miss setting the containing csect, and we would
need to add a lot of XCOFF specialized code around some common code area.

This patch intendeds to do
1. Rely on getFragment().getParent() to get csect from labels.
2. Only use get/setRepresentedCsect (was get/setContainingCsect)
if symbol itself represents a csect.

Reviewers: DiggerLin, hubert.reinterpretcast, daltenty

Differential Revision: https://reviews.llvm.org/D77080

include/llvm/MC/MCSectionXCOFF.h

@@ -48,7 +48,7 @@ class MCSectionXCOFF final : public MCSection {
            "Invalid or unhandled type for csect.");
     assert(QualName != nullptr && "QualName is needed.");
     QualName->setStorageClass(SC);
-    QualName->setContainingCsect(this);
+    QualName->setRepresentedCsect(this);
   }
 
 public:

include/llvm/MC/MCSymbolXCOFF.h

@@ -36,21 +36,6 @@ public:
     return StorageClass.getValue();
   }
 
-  void setContainingCsect(MCSectionXCOFF *C) {
-    assert((!ContainingCsect || ContainingCsect == C) &&
-           "Trying to set a containing csect that doesn't match the one that"
-           "this symbol is already mapped to.");
-    ContainingCsect = C;
-  }
-
-  MCSectionXCOFF *getContainingCsect() const {
-    assert(ContainingCsect &&
-           "Trying to get containing csect but none was set.");
-    return ContainingCsect;
-  }
-
-  bool hasContainingCsect() const { return ContainingCsect != nullptr; }
-
   StringRef getUnqualifiedName() const {
     const StringRef name = getName();
     if (name.back() == ']') {
@@ -62,9 +47,15 @@ public:
     return name;
   }
 
+  bool hasRepresentedCsectSet() const { return RepresentedCsect != nullptr; }
+
+  MCSectionXCOFF *getRepresentedCsect() const;
+
+  void setRepresentedCsect(MCSectionXCOFF *C);
+
 private:
   Optional<XCOFF::StorageClass> StorageClass;
-  MCSectionXCOFF *ContainingCsect = nullptr;
+  MCSectionXCOFF *RepresentedCsect = nullptr;
 };
 
 } // end namespace llvm

lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1749,11 +1749,6 @@ void AsmPrinter::SetupMachineFunction(MachineFunction &MF) {
     // Get the function entry point symbol.
     CurrentFnSym = OutContext.getOrCreateSymbol(
         "." + cast<MCSymbolXCOFF>(CurrentFnDescSym)->getUnqualifiedName());
-
-    // Set the containing csect.
-    MCSectionXCOFF *FnEntryPointSec =
-        cast<MCSectionXCOFF>(getObjFileLowering().SectionForGlobal(&F, TM));
-    cast<MCSymbolXCOFF>(CurrentFnSym)->setContainingCsect(FnEntryPointSec);
   }
 
   CurrentFnSymForSize = CurrentFnSym;
@@ -1846,11 +1841,6 @@ void AsmPrinter::emitConstantPool() {
       if (!Sym->isUndefined())
         continue;
 
-      if (TM.getTargetTriple().isOSBinFormatXCOFF()) {
-        cast<MCSymbolXCOFF>(Sym)->setContainingCsect(
-            cast<MCSectionXCOFF>(CPSections[i].S));
-      }
-
       if (CurSection != CPSections[i].S) {
         OutStreamer->SwitchSection(CPSections[i].S);
         emitAlignment(Align(CPSections[i].Alignment));
@@ -1945,10 +1935,6 @@ void AsmPrinter::emitJumpTableInfo() {
       OutStreamer->emitLabel(GetJTISymbol(JTI, true));
 
     MCSymbol* JTISymbol = GetJTISymbol(JTI);
-    if (TM.getTargetTriple().isOSBinFormatXCOFF()) {
-      cast<MCSymbolXCOFF>(JTISymbol)->setContainingCsect(
-          cast<MCSectionXCOFF>(TLOF.getSectionForJumpTable(F, TM)));
-    }
     OutStreamer->emitLabel(JTISymbol);
 
     for (unsigned ii = 0, ee = JTBBs.size(); ii != ee; ++ii)

lib/CodeGen/MachineModuleInfo.cpp

@@ -105,15 +105,6 @@ ArrayRef<MCSymbol *> MMIAddrLabelMap::getAddrLabelSymbolToEmit(BasicBlock *BB) {
   Entry.Index = BBCallbacks.size() - 1;
   Entry.Fn = BB->getParent();
   MCSymbol *Sym = Context.createTempSymbol(!BB->hasAddressTaken());
-  if (Context.getObjectFileInfo()->getTargetTriple().isOSBinFormatXCOFF()) {
-    MCSymbol *FnEntryPointSym =
-        Context.lookupSymbol("." + Entry.Fn->getName());
-    assert(FnEntryPointSym && "The function entry pointer symbol should have"
-		              " already been initialized.");
-    MCSectionXCOFF *Csect =
-        cast<MCSymbolXCOFF>(FnEntryPointSym)->getContainingCsect();
-    cast<MCSymbolXCOFF>(Sym)->setContainingCsect(Csect);
-  }
   Entry.Symbols.push_back(Sym);
   return Entry.Symbols;
 }

lib/MC/CMakeLists.txt

@@ -44,6 +44,7 @@ add_llvm_component_library(LLVMMC
   MCSubtargetInfo.cpp
   MCSymbol.cpp
   MCSymbolELF.cpp
+  MCSymbolXCOFF.cpp
   MCTargetOptions.cpp
   MCTargetOptionsCommandFlags.cpp
   MCValue.cpp

lib/MC/MCSymbolXCOFF.cpp

@@ -0,0 +1,33 @@
+//===- lib/MC/MCSymbolXCOFF.cpp - XCOFF Code Symbol Representation --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/MCSectionXCOFF.h"
+
+using namespace llvm;
+
+MCSectionXCOFF *MCSymbolXCOFF::getRepresentedCsect() const {
+  assert(RepresentedCsect &&
+         "Trying to get csect representation of this symbol but none was set.");
+  assert((!getName().equals(getUnqualifiedName()) ||
+          RepresentedCsect->getCSectType() == XCOFF::XTY_ER) &&
+         "Symbol does not represent a csect; MCSectionXCOFF that represents "
+         "the symbol should not be (but is) set.");
+  return RepresentedCsect;
+}
+
+void MCSymbolXCOFF::setRepresentedCsect(MCSectionXCOFF *C) {
+  assert(C && "Assigned csect should not be null.");
+  assert((!RepresentedCsect || RepresentedCsect == C) &&
+         "Trying to set a csect that doesn't match the one that"
+         "this symbol is already mapped to.");
+  assert((!getName().equals(getUnqualifiedName()) ||
+          C->getCSectType() == XCOFF::XTY_ER) &&
+         "Symbol does not represent a csect; can only set a MCSectionXCOFF "
+         "representation for a csect.");
+  RepresentedCsect = C;
+}

lib/MC/XCOFFObjectWriter.cpp

@@ -313,6 +313,12 @@ CsectGroup &XCOFFObjectWriter::getCsectGroup(const MCSectionXCOFF *MCSec) {
   }
 }
 
+static MCSectionXCOFF *getContainingCsect(const MCSymbolXCOFF *XSym) {
+  if (XSym->isDefined())
+    return cast<MCSectionXCOFF>(XSym->getFragment()->getParent());
+  return XSym->getRepresentedCsect();
+}
+
 void XCOFFObjectWriter::executePostLayoutBinding(MCAssembler &Asm,
                                                  const MCAsmLayout &Layout) {
   if (TargetObjectWriter->is64Bit())
@@ -341,7 +347,7 @@ void XCOFFObjectWriter::executePostLayoutBinding(MCAssembler &Asm,
       continue;
 
     const MCSymbolXCOFF *XSym = cast<MCSymbolXCOFF>(&S);
-    const MCSectionXCOFF *ContainingCsect = XSym->getContainingCsect();
+    const MCSectionXCOFF *ContainingCsect = getContainingCsect(XSym);
 
     if (ContainingCsect->getCSectType() == XCOFF::XTY_ER) {
       // Handle undefined symbol.
@@ -394,7 +400,7 @@ void XCOFFObjectWriter::recordRelocation(MCAssembler &Asm,
       TargetObjectWriter->getRelocTypeAndSignSize(Target, Fixup, IsPCRel);
 
   const MCSectionXCOFF *SymASec =
-      cast<MCSymbolXCOFF>(SymA).getContainingCsect();
+      getContainingCsect(cast<MCSymbolXCOFF>(&SymA));
   assert(SectionMap.find(SymASec) != SectionMap.end() &&
          "Expected containing csect to exist in map.");
 

lib/Target/PowerPC/PPCAsmPrinter.cpp

@@ -1606,7 +1606,6 @@ void PPCAIXAsmPrinter::emitGlobalVariable(const GlobalVariable *GV) {
           : getObjFileLowering().SectionForGlobal(
                 GV, GVKind = getObjFileLowering().getKindForGlobal(GV, TM),
                 TM));
-  GVSym->setContainingCsect(Csect);
 
   // External global variables are already handled.
   if (GV->isDeclaration())
@@ -1649,7 +1648,7 @@ void PPCAIXAsmPrinter::emitFunctionDescriptor() {
   MCSectionSubPair Current = OutStreamer->getCurrentSection();
   // Emit function descriptor.
   OutStreamer->SwitchSection(
-      cast<MCSymbolXCOFF>(CurrentFnDescSym)->getContainingCsect());
+      cast<MCSymbolXCOFF>(CurrentFnDescSym)->getRepresentedCsect());
   // Emit function entry point address.
   OutStreamer->emitValue(MCSymbolRefExpr::create(CurrentFnSym, OutContext),
                          PointerSize);
@@ -1691,7 +1690,6 @@ void PPCAIXAsmPrinter::emitEndOfAsmFile(Module &M) {
     // Setup the csect for the current TC entry.
     MCSectionXCOFF *TCEntry = cast<MCSectionXCOFF>(
         getObjFileLowering().getSectionForTOCEntry(I.first));
-    cast<MCSymbolXCOFF>(I.second)->setContainingCsect(TCEntry);
     OutStreamer->SwitchSection(TCEntry);
 
     OutStreamer->emitLabel(I.second);

lib/Target/PowerPC/PPCISelLowering.cpp

@@ -5135,14 +5135,14 @@ static SDValue transformCallee(const SDValue &Callee, SelectionDAG &DAG,
         MCSymbolXCOFF *S = cast<MCSymbolXCOFF>(
             Context.getOrCreateSymbol(Twine(".") + Twine(FuncName)));
 
-        if (IsDeclaration && !S->hasContainingCsect()) {
+        if (IsDeclaration && !S->hasRepresentedCsectSet()) {
           // On AIX, an undefined symbol needs to be associated with a
           // MCSectionXCOFF to get the correct storage mapping class.
           // In this case, XCOFF::XMC_PR.
           MCSectionXCOFF *Sec = Context.getXCOFFSection(
               S->getName(), XCOFF::XMC_PR, XCOFF::XTY_ER, SC,
               SectionKind::getMetadata());
-          S->setContainingCsect(Sec);
+          S->setRepresentedCsect(Sec);
         }
 
         MVT PtrVT =