[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
diff --git a/llvm/include/llvm/MC/MCSectionXCOFF.h b/llvm/include/llvm/MC/MCSectionXCOFF.h
index 611eb69..5cc6921 100644
--- a/llvm/include/llvm/MC/MCSectionXCOFF.h
+++ b/llvm/include/llvm/MC/MCSectionXCOFF.h
@@ -48,7 +48,7 @@
            "Invalid or unhandled type for csect.");
     assert(QualName != nullptr && "QualName is needed.");
     QualName->setStorageClass(SC);
-    QualName->setContainingCsect(this);
+    QualName->setRepresentedCsect(this);
   }
 
 public:
diff --git a/llvm/include/llvm/MC/MCSymbolXCOFF.h b/llvm/include/llvm/MC/MCSymbolXCOFF.h
index 07dfb5d..6815f42 100644
--- a/llvm/include/llvm/MC/MCSymbolXCOFF.h
+++ b/llvm/include/llvm/MC/MCSymbolXCOFF.h
@@ -36,21 +36,6 @@
     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 @@
     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
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 6d96205..8c51203 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1749,11 +1749,6 @@
     // 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 @@
       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 @@
       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)
diff --git a/llvm/lib/CodeGen/MachineModuleInfo.cpp b/llvm/lib/CodeGen/MachineModuleInfo.cpp
index ce5961a..a55943aff 100644
--- a/llvm/lib/CodeGen/MachineModuleInfo.cpp
+++ b/llvm/lib/CodeGen/MachineModuleInfo.cpp
@@ -105,15 +105,6 @@
   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;
 }
diff --git a/llvm/lib/MC/CMakeLists.txt b/llvm/lib/MC/CMakeLists.txt
index 9ef61d7..9dca8d7 100644
--- a/llvm/lib/MC/CMakeLists.txt
+++ b/llvm/lib/MC/CMakeLists.txt
@@ -44,6 +44,7 @@
   MCSubtargetInfo.cpp
   MCSymbol.cpp
   MCSymbolELF.cpp
+  MCSymbolXCOFF.cpp
   MCTargetOptions.cpp
   MCTargetOptionsCommandFlags.cpp
   MCValue.cpp
diff --git a/llvm/lib/MC/MCSymbolXCOFF.cpp b/llvm/lib/MC/MCSymbolXCOFF.cpp
new file mode 100644
index 0000000..53aec88
--- /dev/null
+++ b/llvm/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;
+}
diff --git a/llvm/lib/MC/XCOFFObjectWriter.cpp b/llvm/lib/MC/XCOFFObjectWriter.cpp
index 2c05442..fc0f067 100644
--- a/llvm/lib/MC/XCOFFObjectWriter.cpp
+++ b/llvm/lib/MC/XCOFFObjectWriter.cpp
@@ -313,6 +313,12 @@
   }
 }
 
+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 @@
       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 @@
       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.");
 
diff --git a/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp b/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp
index 4e39873..0a256cb 100644
--- a/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp
+++ b/llvm/lib/Target/PowerPC/PPCAsmPrinter.cpp
@@ -1606,7 +1606,6 @@
           : 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 @@
   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 @@
     // 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);
diff --git a/llvm/lib/Target/PowerPC/PPCISelLowering.cpp b/llvm/lib/Target/PowerPC/PPCISelLowering.cpp
index 3223efb..16b818e 100644
--- a/llvm/lib/Target/PowerPC/PPCISelLowering.cpp
+++ b/llvm/lib/Target/PowerPC/PPCISelLowering.cpp
@@ -5135,14 +5135,14 @@
         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 =