[JITLink][COFF] Don't dead strip seh frame of exported function.

Adds keep-alive edges to pdata section to prevent dead strip of block when its parent function is alive.

Reviewed By: lhames

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

llvm/lib/ExecutionEngine/JITLink/COFF_x86_64.cpp

@@ -12,8 +12,8 @@
 
 #include "llvm/ExecutionEngine/JITLink/COFF_x86_64.h"
 #include "COFFLinkGraphBuilder.h"
-#include "EHFrameSupportImpl.h"
 #include "JITLinkGeneric.h"
+#include "SEHFrameSupport.h"
 #include "llvm/BinaryFormat/COFF.h"
 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
 #include "llvm/Object/COFF.h"
@@ -239,9 +239,10 @@ void link_COFF_x86_64(std::unique_ptr<LinkGraph> G,
   const Triple &TT = G->getTargetTriple();
   if (Ctx->shouldAddDefaultTargetPasses(TT)) {
     // Add a mark-live pass.
-    if (auto MarkLive = Ctx->getMarkLivePass(TT))
+    if (auto MarkLive = Ctx->getMarkLivePass(TT)) {
       Config.PrePrunePasses.push_back(std::move(MarkLive));
-    else
+      Config.PrePrunePasses.push_back(SEHFrameKeepAlivePass(".pdata"));
+    } else
       Config.PrePrunePasses.push_back(markAllSymbolsLive);
 
     // Add COFF edge lowering passes.

llvm/lib/ExecutionEngine/JITLink/SEHFrameSupport.h

@@ -0,0 +1,61 @@
+//===------- SEHFrameSupport.h - JITLink seh-frame utils --------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// SEHFrame utils for JITLink.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_JITLINK_SEHFRAMESUPPORT_H
+#define LLVM_EXECUTIONENGINE_JITLINK_SEHFRAMESUPPORT_H
+
+#include "llvm/ADT/Triple.h"
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+#include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/Support/Error.h"
+
+namespace llvm {
+namespace jitlink {
+/// This pass adds keep-alive edge from SEH frame sections
+/// to the parent function content block.
+class SEHFrameKeepAlivePass {
+public:
+  SEHFrameKeepAlivePass(StringRef SEHFrameSectionName)
+      : SEHFrameSectionName(SEHFrameSectionName) {}
+
+  Error operator()(LinkGraph &G) {
+    auto *S = G.findSectionByName(SEHFrameSectionName);
+    if (!S)
+      return Error::success();
+
+    // Simply consider every block pointed by seh frame block as parants.
+    // This adds some unnecessary keep-alive edges to unwind info blocks,
+    // (xdata) but these blocks are usually dead by default, so they wouldn't
+    // count for the fate of seh frame block.
+    for (auto *B : S->blocks()) {
+      auto &DummySymbol = G.addAnonymousSymbol(*B, 0, 0, false, false);
+      DenseSet<Block *> Children;
+      for (auto &E : B->edges()) {
+        auto &Sym = E.getTarget();
+        if (!Sym.isDefined())
+          continue;
+        Children.insert(&Sym.getBlock());
+      }
+      for (auto *Child : Children)
+        Child->addEdge(Edge(Edge::KeepAlive, 0, DummySymbol, 0));
+    }
+    return Error::success();
+  }
+
+private:
+  StringRef SEHFrameSectionName;
+};
+
+} // end namespace jitlink
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_JITLINK_SEHFRAMESUPPORT_H

llvm/test/ExecutionEngine/JITLink/X86/COFF_pdata_no_strip.s

@@ -0,0 +1,31 @@
+# RUN: llvm-mc -filetype=obj -triple=x86_64-windows-msvc %s -o %t
+# RUN: 
+# RUN: llvm-jitlink -abs __ImageBase=0xdeadbeaf -noexec %t \
+# RUN: -slab-allocate 100Kb -slab-address 0xfff00000 -slab-page-size 4096 \
+# RUN: -show-graph -noexec 2>&1 | FileCheck %s
+#
+# Check that basic seh frame inside pdata of alive function is not dead-stripped out.
+# CHECK: section .xdata:
+# CHECK-EMPTY:
+# CHECK-NEXT: block 0xfff00000 size = 0x00000008, align = 4, alignment-offset = 0
+# CHECK-NEXT: symbols:
+# CHECK-NEXT:   0xfff00000 (block + 0x00000000): size: 0x00000008, linkage: strong, scope: local, live  -   .xdata
+
+	.text
+
+	.def	main;
+	.scl	2;
+	.type	32;
+	.endef
+	.globl	main
+	.p2align	4, 0x90
+main:
+.seh_proc main
+	subq	$40, %rsp
+	.seh_stackalloc 40
+	.seh_endprologue
+	movl	$0, 36(%rsp)
+	nop
+	addq	$40, %rsp
+	retq
+	.seh_endproc

llvm/test/ExecutionEngine/JITLink/X86/COFF_pdata_strip.s

@@ -0,0 +1,43 @@
+# RUN: llvm-mc -filetype=obj -triple=x86_64-windows-msvc %s -o %t
+# RUN: 
+# RUN: llvm-jitlink -abs __ImageBase=0xdeadbeaf -noexec %t \
+# RUN: -slab-allocate 100Kb -slab-address 0xfff00000 -slab-page-size 4096 \
+# RUN: -show-graph -noexec 2>&1 | FileCheck %s
+#
+# Check that basic seh frame of dead block is dead-stripped out
+#
+# CHECK: section .func:
+# CHECK-EMPTY:
+# CHECK-NEXT: section .xdata:
+# CHECK-EMPTY:
+# CHECK-NEXT: section .pdata:
+# CHECK-EMPTY:
+
+	.text
+	
+	.def	main;
+	.scl	2;
+	.type	32;
+	.endef
+	.globl	main
+	.p2align	4, 0x90
+main:
+	retq
+
+	.section .func
+
+    .def	func;
+	.scl	3;
+	.type	32;
+	.endef
+	.p2align	4, 0x90
+func: 
+	.seh_proc func
+	subq	$40, %rsp
+	.seh_stackalloc 40
+	.seh_endprologue
+	movl	$0, 36(%rsp)
+	nop
+	addq	$40, %rsp
+	retq
+	.seh_endproc