From 6d0b97357f50d3319b5b5e95eca00d62a1644e58 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Sat, 14 Sep 2024 11:33:52 -0700
Subject: [PATCH] [BOLT] Allow sections in --print-only flag

While printing functions, expand --print-only flag to accept section
names. E.g., "--print-only=\.init" will only print functions from
".init" section.
---
 bolt/lib/Core/BinaryFunction.cpp   |  6 ++++++
 bolt/test/X86/print-only-section.s | 29 +++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)
 create mode 100644 bolt/test/X86/print-only-section.s

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 46bdf208be6ad..36c42fced93d0 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -165,6 +165,12 @@ bool shouldPrint(const BinaryFunction &Function) {
     }
   }
 
+  std::optional<StringRef> Origin = Function.getOriginSectionName();
+  if (Origin && llvm::any_of(opts::PrintOnly, [&](const std::string &Name) {
+        return Name == *Origin;
+      }))
+    return true;
+
   return false;
 }
 
diff --git a/bolt/test/X86/print-only-section.s b/bolt/test/X86/print-only-section.s
new file mode 100644
index 0000000000000..d580818ca4fc6
--- /dev/null
+++ b/bolt/test/X86/print-only-section.s
@@ -0,0 +1,29 @@
+## Check that --print-only flag works with sections.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-linux %s -o %t.o
+# RUN: ld.lld %t.o -o %t.exe
+# RUN: llvm-bolt %t.exe -o %t.out --print-cfg --print-only=unused_code 2>&1 \
+# RUN:   | FileCheck %s
+
+# CHECK: Binary Function "foo"
+# CHECK-NOT: Binary Function "_start"
+
+  .text
+  .globl _start
+  .type _start, %function
+_start:
+  .cfi_startproc
+  ret
+  .cfi_endproc
+  .size _start, .-_start
+
+  .section unused_code,"ax",@progbits
+  .globl foo
+  .type foo, %function
+foo:
+  .cfi_startproc
+  ret
+  .cfi_endproc
+  .size foo, .-foo