[JITLink][MachO][arm64] Add support for splitting compact-unwind sections.

CompactUnwindSplitter splits compact-unwind sections on record boundaries and
adds keep-alive edges from target functions back to their respective records.

In MachO_arm64.cpp, a CompactUnwindSplitter pass is added to the pre-prune pass
list when setting up the standard pipeline.

This patch does not provide runtime support for compact-unwind, but is a first
step towards enabling it.

llvm/lib/ExecutionEngine/JITLink/MachOLinkGraphBuilder.cpp

@@ -701,5 +701,119 @@ Error MachOLinkGraphBuilder::graphifyCStringSection(
   return Error::success();
 }
 
+Error CompactUnwindSplitter::operator()(LinkGraph &G) {
+  auto *CUSec = G.findSectionByName(CompactUnwindSectionName);
+  if (!CUSec)
+    return Error::success();
+
+  if (!G.getTargetTriple().isOSBinFormatMachO())
+    return make_error<JITLinkError>(
+        "Error linking " + G.getName() +
+        ": compact unwind splitting not supported on non-macho target " +
+        G.getTargetTriple().str());
+
+  unsigned CURecordSize = 0;
+  unsigned PersonalityEdgeOffset = 0;
+  unsigned LSDAEdgeOffset = 0;
+  switch (G.getTargetTriple().getArch()) {
+  case Triple::aarch64:
+  case Triple::x86_64:
+    // 64-bit compact-unwind record format:
+    // Range start: 8 bytes.
+    // Range size:  4 bytes.
+    // CU encoding: 4 bytes.
+    // Personality: 8 bytes.
+    // LSDA:        8 bytes.
+    CURecordSize = 32;
+    PersonalityEdgeOffset = 16;
+    LSDAEdgeOffset = 24;
+    break;
+  default:
+    return make_error<JITLinkError>(
+        "Error linking " + G.getName() +
+        ": compact unwind splitting not supported on " +
+        G.getTargetTriple().getArchName());
+  }
+
+  std::vector<Block *> OriginalBlocks(CUSec->blocks().begin(),
+                                      CUSec->blocks().end());
+  LLVM_DEBUG({
+    dbgs() << "In " << G.getName() << " splitting compact unwind section "
+           << CompactUnwindSectionName << " containing "
+           << OriginalBlocks.size() << " initial blocks...\n";
+  });
+
+  while (!OriginalBlocks.empty()) {
+    auto *B = OriginalBlocks.back();
+    OriginalBlocks.pop_back();
+
+    if (B->getSize() == 0) {
+      LLVM_DEBUG({
+        dbgs() << "  Skipping empty block at "
+               << formatv("{0:x16}", B->getAddress()) << "\n";
+      });
+      continue;
+    }
+
+    LLVM_DEBUG({
+      dbgs() << "  Splitting block at " << formatv("{0:x16}", B->getAddress())
+             << " into " << (B->getSize() / CURecordSize)
+             << " compact unwind record(s)\n";
+    });
+
+    if (B->getSize() % CURecordSize)
+      return make_error<JITLinkError>(
+          "Error splitting compact unwind record in " + G.getName() +
+          ": block at " + formatv("{0:x}", B->getAddress()) + " has size " +
+          formatv("{0:x}", B->getSize()) +
+          " (not a multiple of CU record size of " +
+          formatv("{0:x}", CURecordSize) + ")");
+
+    unsigned NumBlocks = B->getSize() / CURecordSize;
+    LinkGraph::SplitBlockCache C;
+
+    for (unsigned I = 0; I != NumBlocks; ++I) {
+      auto &CURec = G.splitBlock(*B, CURecordSize, &C);
+      bool AddedKeepAlive = false;
+
+      for (auto &E : CURec.edges()) {
+        if (E.getOffset() == 0) {
+          LLVM_DEBUG({
+            dbgs() << "    Updating compact unwind record at "
+                   << formatv("{0:x16}", CURec.getAddress()) << " to point to "
+                   << (E.getTarget().hasName() ? E.getTarget().getName()
+                                               : StringRef())
+                   << " (at " << formatv("{0:x16}", E.getTarget().getAddress())
+                   << ")\n";
+          });
+
+          if (E.getTarget().isExternal())
+            return make_error<JITLinkError>(
+                "Error adding keep-alive edge for compact unwind record at " +
+                formatv("{0:x}", CURec.getAddress()) + ": target " +
+                E.getTarget().getName() + " is an external symbol");
+          auto &TgtBlock = E.getTarget().getBlock();
+          auto &CURecSym =
+              G.addAnonymousSymbol(CURec, 0, CURecordSize, 0, false);
+          TgtBlock.addEdge(Edge::KeepAlive, 0, CURecSym, 0);
+          AddedKeepAlive = true;
+        } else if (E.getOffset() != PersonalityEdgeOffset &&
+                   E.getOffset() != LSDAEdgeOffset)
+          return make_error<JITLinkError>("Unexpected edge at offset " +
+                                          formatv("{0:x}", E.getOffset()) +
+                                          " in compact unwind record at " +
+                                          formatv("{0:x}", CURec.getAddress()));
+      }
+
+      if (!AddedKeepAlive)
+        return make_error<JITLinkError>(
+            "Error adding keep-alive edge for compact unwind record at " +
+            formatv("{0:x}", CURec.getAddress()) +
+            ": no outgoing target edge at offset 0");
+    }
+  }
+  return Error::success();
+}
+
 } // end namespace jitlink
 } // end namespace llvm

llvm/lib/ExecutionEngine/JITLink/MachOLinkGraphBuilder.h

@@ -231,6 +231,17 @@ private:
   StringMap<SectionParserFunction> CustomSectionParserFunctions;
 };
 
+/// A pass to split up __LD,__compact_unwind sections.
+class CompactUnwindSplitter {
+public:
+  CompactUnwindSplitter(StringRef CompactUnwindSectionName)
+      : CompactUnwindSectionName(CompactUnwindSectionName) {}
+  Error operator()(LinkGraph &G);
+
+private:
+  StringRef CompactUnwindSectionName;
+};
+
 } // end namespace jitlink
 } // end namespace llvm
 

llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp

@@ -683,6 +683,10 @@ void link_MachO_arm64(std::unique_ptr<LinkGraph> G,
     else
       Config.PrePrunePasses.push_back(markAllSymbolsLive);
 
+    // Add compact unwind splitter pass.
+    Config.PrePrunePasses.push_back(
+        CompactUnwindSplitter("__LD,__compact_unwind"));
+
     // Add an in-place GOT/Stubs pass.
     Config.PostPrunePasses.push_back(
         PerGraphGOTAndPLTStubsBuilder_MachO_arm64::asPass);

llvm/test/ExecutionEngine/JITLink/AArch64/MachO_arm64_compact_unwind.s

@@ -0,0 +1,20 @@
+# REQUIRES: asserts
+# RUN: llvm-mc -triple=arm64-apple-ios -filetype=obj -o %t %s
+# RUN: llvm-jitlink -noexec -debug-only=jitlink %t 2>&1 | FileCheck %s
+#
+# Check that splitting of compact-unwind sections works.
+#
+# CHECK: splitting {{.*}} __LD,__compact_unwind containing 1 initial blocks...
+# CHECK:   Splitting {{.*}} into 1 compact unwind record(s)
+# CHECK:     Updating {{.*}} to point to _main {{.*}}
+
+	.section	__TEXT,__text,regular,pure_instructions
+	.globl	_main
+	.p2align	2
+_main:
+	.cfi_startproc
+	ret
+	.cfi_endproc
+
+.subsections_via_symbols
+