[MC][ARM] Reuse symbol value in constant pool

Fix https://github.com/llvm/llvm-project/issues/55816

Before this patch, MCConstantExpr were reused, but MCSymbolExpr were
not. To reuse symbol value, this patch added a DenseMap to record the
symbol value.

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

3 files changed, 22 insertions, 12 deletions

diff --git a/llvm/include/llvm/MC/ConstantPools.h b/llvm/include/llvm/MC/ConstantPools.h
index 9fe0cce..7eac753 100644
--- a/llvm/include/llvm/MC/ConstantPools.h
+++ b/llvm/include/llvm/MC/ConstantPools.h
@@ -43,7 +43,8 @@ struct ConstantPoolEntry {
 class ConstantPool {
   using EntryVecTy = SmallVector<ConstantPoolEntry, 4>;
   EntryVecTy Entries;
-  std::map<int64_t, const MCSymbolRefExpr *> CachedEntries;
+  std::map<int64_t, const MCSymbolRefExpr *> CachedConstantEntries;
+  DenseMap<const MCSymbol *, const MCSymbolRefExpr *> CachedSymbolEntries;
 
 public:
   // Initialize a new empty constant pool
diff --git a/llvm/lib/MC/ConstantPools.cpp b/llvm/lib/MC/ConstantPools.cpp
index d8a08a4..4766769 100644
--- a/llvm/lib/MC/ConstantPools.cpp
+++ b/llvm/lib/MC/ConstantPools.cpp
@@ -39,25 +39,38 @@ void ConstantPool::emitEntries(MCStreamer &Streamer) {
 const MCExpr *ConstantPool::addEntry(const MCExpr *Value, MCContext &Context,
                                      unsigned Size, SMLoc Loc) {
   const MCConstantExpr *C = dyn_cast<MCConstantExpr>(Value);
+  const MCSymbolRefExpr *S = dyn_cast<MCSymbolRefExpr>(Value);
 
   // Check if there is existing entry for the same constant. If so, reuse it.
-  auto Itr = C ? CachedEntries.find(C->getValue()) : CachedEntries.end();
-  if (Itr != CachedEntries.end())
-    return Itr->second;
+  if (C) {
+    auto CItr = CachedConstantEntries.find(C->getValue());
+    if (CItr != CachedConstantEntries.end())
+      return CItr->second;
+  }
+
+  // Check if there is existing entry for the same symbol. If so, reuse it.
+  if (S) {
+    auto SItr = CachedSymbolEntries.find(&(S->getSymbol()));
+    if (SItr != CachedSymbolEntries.end())
+      return SItr->second;
+  }
 
   MCSymbol *CPEntryLabel = Context.createTempSymbol();
 
   Entries.push_back(ConstantPoolEntry(CPEntryLabel, Value, Size, Loc));
   const auto SymRef = MCSymbolRefExpr::create(CPEntryLabel, Context);
   if (C)
-    CachedEntries[C->getValue()] = SymRef;
+    CachedConstantEntries[C->getValue()] = SymRef;
+  if (S)
+    CachedSymbolEntries[&(S->getSymbol())] = SymRef;
   return SymRef;
 }
 
 bool ConstantPool::empty() { return Entries.empty(); }
 
 void ConstantPool::clearCache() {
-  CachedEntries.clear();
+  CachedConstantEntries.clear();
+  CachedSymbolEntries.clear();
 }
 
 //
diff --git a/llvm/test/MC/ARM/ldr-pseudo-wide.s b/llvm/test/MC/ARM/ldr-pseudo-wide.s
index b4deca5..c059ee02 100644
--- a/llvm/test/MC/ARM/ldr-pseudo-wide.s
+++ b/llvm/test/MC/ARM/ldr-pseudo-wide.s
@@ -36,9 +36,9 @@ f2:
 
   ldr.w r0, =foo
 @ CHECK-ARM: ldr r0, .Ltmp[[TMP2:[0-9]+]]
-@ CHECK-DARWIN-ARM: ldr r0, Ltmp2
+@ CHECK-DARWIN-ARM: ldr r0, Ltmp1
 @ CHECK-THUMB2: ldr.w r0, .Ltmp[[TMP2:[0-9]+]]
-@ CHECK-DARWIN-THUMB2: ldr.w r0, Ltmp2
+@ CHECK-DARWIN-THUMB2: ldr.w r0, Ltmp1
 @ CHECK-THUMB: error: instruction requires: thumb2
 @ CHECK-THUMB-NEXT:  ldr.w r0, =foo
 
@@ -56,12 +56,8 @@ f3:
 @ CHECK-NEXT: .long   65538
 @ CHECK: .Ltmp1:
 @ CHECK-NEXT: .long   foo
-@ CHECK: .Ltmp2:
-@ CHECK-NEXT: .long   foo
 
 @ CHECK-DARWIN: Ltmp0:
 @ CHECK-DARWIN-NEXT: .long   65538
 @ CHECK-DARWIN: Ltmp1:
 @ CHECK-DARWIN-NEXT: .long   foo
-@ CHECK-DARWIN: Ltmp2:
-@ CHECK-DARWIN-NEXT: .long   foo