[JITLink] Allow multiple relocations at same offset in EHFrameEdgeFixer (#68252)

The pass only requires that it can determine a uniquely identified
target at some offsets. Multiple relocations at the same offset are fine
otherwise and will be required when adding exception handling support
for RISC-V.

Author: hahnjo

llvm/lib/ExecutionEngine/JITLink/EHFrameSupport.cpp

@@ -126,16 +126,23 @@ Error EHFrameEdgeFixer::processBlock(ParseContext &PC, Block &B) {
   }
 
   // Find the offsets of any existing edges from this block.
-  BlockEdgeMap BlockEdges;
+  BlockEdgesInfo BlockEdges;
   for (auto &E : B.edges())
     if (E.isRelocation()) {
-      if (BlockEdges.count(E.getOffset()))
-        return make_error<JITLinkError>(
-            "Multiple relocations at offset " +
-            formatv("{0:x16}", E.getOffset()) + " in " + EHFrameSectionName +
-            " block at address " + formatv("{0:x16}", B.getAddress()));
-
-      BlockEdges[E.getOffset()] = EdgeTarget(E);
+      // Check if we already saw more than one relocation at this offset.
+      if (BlockEdges.Multiple.contains(E.getOffset()))
+        continue;
+
+      // Otherwise check if we previously had exactly one relocation at this
+      // offset. If so, we now have a second one and move it from the TargetMap
+      // into the Multiple set.
+      auto It = BlockEdges.TargetMap.find(E.getOffset());
+      if (It != BlockEdges.TargetMap.end()) {
+        BlockEdges.TargetMap.erase(It);
+        BlockEdges.Multiple.insert(E.getOffset());
+      } else {
+        BlockEdges.TargetMap[E.getOffset()] = EdgeTarget(E);
+      }
     }
 
   BinaryStreamReader BlockReader(
@@ -172,7 +179,7 @@ Error EHFrameEdgeFixer::processBlock(ParseContext &PC, Block &B) {
 
 Error EHFrameEdgeFixer::processCIE(ParseContext &PC, Block &B,
                                    size_t CIEDeltaFieldOffset,
-                                   const BlockEdgeMap &BlockEdges) {
+                                   const BlockEdgesInfo &BlockEdges) {
 
   LLVM_DEBUG(dbgs() << "    Record is CIE\n");
 
@@ -285,7 +292,7 @@ Error EHFrameEdgeFixer::processCIE(ParseContext &PC, Block &B,
 Error EHFrameEdgeFixer::processFDE(ParseContext &PC, Block &B,
                                    size_t CIEDeltaFieldOffset,
                                    uint32_t CIEDelta,
-                                   const BlockEdgeMap &BlockEdges) {
+                                   const BlockEdgesInfo &BlockEdges) {
   LLVM_DEBUG(dbgs() << "    Record is FDE\n");
 
   orc::ExecutorAddr RecordAddress = B.getAddress();
@@ -303,12 +310,17 @@ Error EHFrameEdgeFixer::processFDE(ParseContext &PC, Block &B,
 
   {
     // Process the CIE pointer field.
-    auto CIEEdgeItr = BlockEdges.find(CIEDeltaFieldOffset);
+    if (BlockEdges.Multiple.contains(CIEDeltaFieldOffset))
+      return make_error<JITLinkError>(
+          "CIE pointer field already has multiple edges at " +
+          formatv("{0:x16}", RecordAddress + CIEDeltaFieldOffset));
+
+    auto CIEEdgeItr = BlockEdges.TargetMap.find(CIEDeltaFieldOffset);
 
     orc::ExecutorAddr CIEAddress =
         RecordAddress + orc::ExecutorAddrDiff(CIEDeltaFieldOffset) -
         orc::ExecutorAddrDiff(CIEDelta);
-    if (CIEEdgeItr == BlockEdges.end()) {
+    if (CIEEdgeItr == BlockEdges.TargetMap.end()) {
       LLVM_DEBUG({
         dbgs() << "        Adding edge at "
                << (RecordAddress + CIEDeltaFieldOffset)
@@ -497,7 +509,7 @@ Error EHFrameEdgeFixer::skipEncodedPointer(uint8_t PointerEncoding,
 }
 
 Expected<Symbol *> EHFrameEdgeFixer::getOrCreateEncodedPointerEdge(
-    ParseContext &PC, const BlockEdgeMap &BlockEdges, uint8_t PointerEncoding,
+    ParseContext &PC, const BlockEdgesInfo &BlockEdges, uint8_t PointerEncoding,
     BinaryStreamReader &RecordReader, Block &BlockToFix,
     size_t PointerFieldOffset, const char *FieldName) {
   using namespace dwarf;
@@ -508,8 +520,8 @@ Expected<Symbol *> EHFrameEdgeFixer::getOrCreateEncodedPointerEdge(
   // If there's already an edge here then just skip the encoded pointer and
   // return the edge's target.
   {
-    auto EdgeI = BlockEdges.find(PointerFieldOffset);
-    if (EdgeI != BlockEdges.end()) {
+    auto EdgeI = BlockEdges.TargetMap.find(PointerFieldOffset);
+    if (EdgeI != BlockEdges.TargetMap.end()) {
       LLVM_DEBUG({
         dbgs() << "      Existing edge at "
                << (BlockToFix.getAddress() + PointerFieldOffset) << " to "
@@ -522,6 +534,10 @@ Expected<Symbol *> EHFrameEdgeFixer::getOrCreateEncodedPointerEdge(
         return std::move(Err);
       return EdgeI->second.Target;
     }
+
+    if (BlockEdges.Multiple.contains(PointerFieldOffset))
+      return make_error<JITLinkError>("Multiple relocations at offset " +
+                                      formatv("{0:x16}", PointerFieldOffset));
   }
 
   // Switch absptr to corresponding udata encoding.

llvm/lib/ExecutionEngine/JITLink/EHFrameSupportImpl.h

@@ -60,7 +60,11 @@ class EHFrameEdgeFixer {
     Edge::AddendT Addend = 0;
   };
 
-  using BlockEdgeMap = DenseMap<Edge::OffsetT, EdgeTarget>;
+  struct BlockEdgesInfo {
+    DenseMap<Edge::OffsetT, EdgeTarget> TargetMap;
+    DenseSet<Edge::OffsetT> Multiple;
+  };
+
   using CIEInfosMap = DenseMap<orc::ExecutorAddr, CIEInformation>;
 
   struct ParseContext {
@@ -82,9 +86,9 @@ class EHFrameEdgeFixer {
 
   Error processBlock(ParseContext &PC, Block &B);
   Error processCIE(ParseContext &PC, Block &B, size_t CIEDeltaFieldOffset,
-                   const BlockEdgeMap &BlockEdges);
+                   const BlockEdgesInfo &BlockEdges);
   Error processFDE(ParseContext &PC, Block &B, size_t CIEDeltaFieldOffset,
-                   uint32_t CIEDelta, const BlockEdgeMap &BlockEdges);
+                   uint32_t CIEDelta, const BlockEdgesInfo &BlockEdges);
 
   Expected<AugmentationInfo>
   parseAugmentationString(BinaryStreamReader &RecordReader);
@@ -94,9 +98,9 @@ class EHFrameEdgeFixer {
   Error skipEncodedPointer(uint8_t PointerEncoding,
                            BinaryStreamReader &RecordReader);
   Expected<Symbol *> getOrCreateEncodedPointerEdge(
-      ParseContext &PC, const BlockEdgeMap &BlockEdges, uint8_t PointerEncoding,
-      BinaryStreamReader &RecordReader, Block &BlockToFix,
-      size_t PointerFieldOffset, const char *FieldName);
+      ParseContext &PC, const BlockEdgesInfo &BlockEdges,
+      uint8_t PointerEncoding, BinaryStreamReader &RecordReader,
+      Block &BlockToFix, size_t PointerFieldOffset, const char *FieldName);
 
   Expected<Symbol &> getOrCreateSymbol(ParseContext &PC,
                                        orc::ExecutorAddr Addr);