Add thunk vector member to BinaryAddressMapper. This will enable resolving

branches to thunks for more accurate CFG computation.

PiperOrigin-RevId: 703621698

Author: Propeller Team

Diff for: propeller/BUILD

@@ -269,6 +269,7 @@ cc_library(
     deps = [
         ":addr2cu",
         ":status_macros",
+        "@abseil-cpp//absl/container:btree",
         "@abseil-cpp//absl/container:flat_hash_map",
         "@abseil-cpp//absl/log",
         "@abseil-cpp//absl/log:check",
@@ -280,6 +281,7 @@ cc_library(
         "@llvm-project//llvm:DebugInfo",
         "@llvm-project//llvm:Object",
         "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:TargetParser",
     ],
 )
 
@@ -1080,6 +1082,7 @@ cc_test(
     name = "binary_content_test",
     srcs = ["binary_content_test.cc"],
     data = [
+        "//propeller/testdata:fake_thunks.bin",
         "//propeller/testdata:llvm_function_samples.binary",
         "//propeller/testdata:propeller_barebone_nopie_buildid",
         "//propeller/testdata:propeller_barebone_pie_nobuildid_bin",

Diff for: propeller/binary_address_mapper.cc

@@ -13,6 +13,7 @@
 #include "absl/algorithm/container.h"
 #include "absl/base/attributes.h"
 #include "absl/base/nullability.h"
+#include "absl/container/btree_map.h"
 #include "absl/container/btree_set.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
@@ -84,7 +85,9 @@ class BinaryAddressMapperBuilder {
           symtab,
       std::vector<llvm::object::BBAddrMap> bb_addr_map, PropellerStats &stats,
       absl::Nonnull<const PropellerOptions *> options
-          ABSL_ATTRIBUTE_LIFETIME_BOUND);
+          ABSL_ATTRIBUTE_LIFETIME_BOUND,
+      std::optional<absl::btree_map<uint64_t, llvm::object::ELFSymbolRef>>
+          thunk_map = std::nullopt);
 
   BinaryAddressMapperBuilder(const BinaryAddressMapperBuilder &) = delete;
   BinaryAddressMapperBuilder &operator=(const BinaryAddressMapper &) = delete;
@@ -131,6 +134,9 @@ class BinaryAddressMapperBuilder {
   int FilterDuplicateNameFunctions(
       absl::btree_set<int> &selected_functions) const;
 
+  // Create a sorted vector of thunks in the binary from `thunk_map_`.
+  std::optional<std::vector<ThunkInfo>> GetThunks();
+
   // BB address map of functions.
   std::vector<llvm::object::BBAddrMap> bb_addr_map_;
   // Non-zero sized function symbols from elf symbol table, indexed by
@@ -144,6 +150,10 @@ class BinaryAddressMapperBuilder {
 
   PropellerStats *stats_;
   const PropellerOptions *options_;
+
+  // Map of thunks by address.
+  std::optional<absl::btree_map<uint64_t, llvm::object::ELFSymbolRef>>
+      thunk_map_;
 };
 
 // Helper class for extracting intra-function paths from binary-address paths.
@@ -638,6 +648,17 @@ absl::btree_set<int> BinaryAddressMapperBuilder::SelectFunctions(
   return selected_functions;
 }
 
+std::optional<std::vector<ThunkInfo>> BinaryAddressMapperBuilder::GetThunks() {
+  if (!thunk_map_.has_value()) return std::nullopt;
+  std::vector<ThunkInfo> thunks;
+  for (const auto &thunk_entry : *thunk_map_) {
+    uint64_t thunk_address = thunk_entry.first;
+    llvm::object::ELFSymbolRef thunk_symbol = thunk_entry.second;
+    thunks.push_back({.address = thunk_address, .symbol = thunk_symbol});
+  }
+  return thunks;
+}
+
 std::vector<BbHandleBranchPath> BinaryAddressMapper::ExtractIntraFunctionPaths(
     const BinaryAddressBranchPath &address_path) const {
   return IntraFunctionPathsExtractor(this).Extract(address_path);
@@ -647,12 +668,15 @@ BinaryAddressMapperBuilder::BinaryAddressMapperBuilder(
     absl::flat_hash_map<uint64_t, llvm::SmallVector<llvm::object::ELFSymbolRef>>
         symtab,
     std::vector<llvm::object::BBAddrMap> bb_addr_map, PropellerStats &stats,
-    absl::Nonnull<const PropellerOptions *> options)
+    absl::Nonnull<const PropellerOptions *> options,
+    std::optional<absl::btree_map<uint64_t, llvm::object::ELFSymbolRef>>
+        thunk_map)
     : bb_addr_map_(std::move(bb_addr_map)),
       symtab_(std::move(symtab)),
       symbol_info_map_(GetSymbolInfoMap(symtab_, bb_addr_map_)),
       stats_(&stats),
-      options_(options) {
+      options_(options),
+      thunk_map_(std::move(thunk_map)) {
   stats_->bbaddrmap_stats.bbaddrmap_function_does_not_have_symtab_entry +=
       bb_addr_map_.size() - symbol_info_map_.size();
 }
@@ -661,11 +685,13 @@ BinaryAddressMapper::BinaryAddressMapper(
     absl::btree_set<int> selected_functions,
     std::vector<llvm::object::BBAddrMap> bb_addr_map,
     std::vector<BbHandle> bb_handles,
-    absl::flat_hash_map<int, FunctionSymbolInfo> symbol_info_map)
+    absl::flat_hash_map<int, FunctionSymbolInfo> symbol_info_map,
+    std::optional<std::vector<ThunkInfo>> thunks)
     : selected_functions_(std::move(selected_functions)),
       bb_handles_(std::move(bb_handles)),
       bb_addr_map_(std::move(bb_addr_map)),
-      symbol_info_map_(std::move(symbol_info_map)) {}
+      symbol_info_map_(std::move(symbol_info_map)),
+      thunks_(std::move(thunks)) {}
 
 absl::StatusOr<std::unique_ptr<BinaryAddressMapper>> BuildBinaryAddressMapper(
     const PropellerOptions &options, const BinaryContent &binary_content,
@@ -676,14 +702,16 @@ absl::StatusOr<std::unique_ptr<BinaryAddressMapper>> BuildBinaryAddressMapper(
   ASSIGN_OR_RETURN(bb_addr_map, ReadBbAddrMap(binary_content));
 
   return BinaryAddressMapperBuilder(ReadSymbolTable(binary_content),
-                                    std::move(bb_addr_map), stats, &options)
+                                    std::move(bb_addr_map), stats, &options,
+                                    ReadThunkSymbols(binary_content))
       .Build(hot_addresses);
 }
 
 std::unique_ptr<BinaryAddressMapper> BinaryAddressMapperBuilder::Build(
     const absl::flat_hash_set<uint64_t> *hot_addresses) && {
   std::optional<uint64_t> last_function_address;
   std::vector<BbHandle> bb_handles;
+  std::optional<std::vector<ThunkInfo>> thunks = GetThunks();
   absl::btree_set<int> selected_functions = SelectFunctions(hot_addresses);
   DropNonSelectedFunctions(selected_functions);
   for (int function_index : selected_functions) {
@@ -696,9 +724,10 @@ std::unique_ptr<BinaryAddressMapper> BinaryAddressMapperBuilder::Build(
       bb_handles.push_back({function_index, bb_index});
     last_function_address = function_bb_addr_map.getFunctionAddress();
   }
+
   return std::make_unique<BinaryAddressMapper>(
       std::move(selected_functions), std::move(bb_addr_map_),
-      std::move(bb_handles), std::move(symbol_info_map_));
+      std::move(bb_handles), std::move(symbol_info_map_), std::move(thunks));
 }
 
 }  // namespace propeller

Diff for: propeller/binary_address_mapper.h

@@ -18,6 +18,7 @@
 #include "absl/time/time.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/ELFTypes.h"
 #include "propeller/bb_handle.h"
 #include "propeller/binary_address_branch_path.h"
@@ -103,6 +104,12 @@ struct BbHandleBranchPath {
   }
 };
 
+struct ThunkInfo {
+  uint64_t address;
+  uint64_t target;
+  llvm::object::ELFSymbolRef symbol;
+};
+
 // Finds basic block entries from binary addresses.
 class BinaryAddressMapper {
  public:
@@ -120,7 +127,8 @@ class BinaryAddressMapper {
       absl::btree_set<int> selected_functions,
       std::vector<llvm::object::BBAddrMap> bb_addr_map,
       std::vector<BbHandle> bb_handles,
-      absl::flat_hash_map<int, FunctionSymbolInfo> symbol_info_map);
+      absl::flat_hash_map<int, FunctionSymbolInfo> symbol_info_map,
+      std::optional<std::vector<ThunkInfo>> thunks = std::nullopt);
 
   BinaryAddressMapper(const BinaryAddressMapper &) = delete;
   BinaryAddressMapper &operator=(const BinaryAddressMapper &) = delete;
@@ -141,6 +149,10 @@ class BinaryAddressMapper {
     return selected_functions_;
   }
 
+  const std::optional<std::vector<ThunkInfo>> &thunks() const {
+    return thunks_;
+  }
+
   // Returns the `bb_handles_` index associated with the binary address
   // `address` given a branch from/to this address based on `direction`.
   // It returns nullopt if the no `bb_handles_` index can be mapped.
@@ -268,6 +280,9 @@ class BinaryAddressMapper {
   // A map from function indices to their symbol info (function names and
   // section name).
   absl::flat_hash_map<int, FunctionSymbolInfo> symbol_info_map_;
+
+  // A vector of thunks in the binary, ordered in increasing order of address.
+  std::optional<std::vector<ThunkInfo>> thunks_;
 };
 
 // Builds a `BinaryAddressMapper` for binary represented by `binary_content` and

Diff for: propeller/binary_content.cc

@@ -8,6 +8,7 @@
 #include <utility>
 #include <vector>
 
+#include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
@@ -35,6 +36,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/TargetParser/Triple.h"
 #include "propeller/addr2cu.h"
 #include "propeller/status_macros.h"
 
@@ -241,6 +243,30 @@ absl::Status ELFFileUtil<ELFT>::InitializeKernelModule(
 }
 }  // namespace
 
+// Read AArch64 thunks from the symbol table and store them in sorted order.
+absl::btree_map<uint64_t, llvm::object::ELFSymbolRef> ReadAArch64ThunkSymbols(
+    const BinaryContent &binary_content) {
+  absl::btree_map<uint64_t, llvm::object::ELFSymbolRef> thunk_map;
+  for (llvm::object::SymbolRef sr : binary_content.object_file->symbols()) {
+    llvm::object::ELFSymbolRef symbol(sr);
+    uint8_t stt = symbol.getELFType();
+    if (stt != llvm::ELF::STT_FUNC) continue;
+    llvm::Expected<uint64_t> address = sr.getAddress();
+    if (!address || !*address) continue;
+
+    llvm::Expected<llvm::StringRef> func_name = symbol.getName();
+    if (!func_name || (!func_name->starts_with("__AArch64ADRPThunk_") &&
+                       !func_name->starts_with("__AArch64AbsLongThunk_")))
+      continue;
+
+    const uint64_t func_size = symbol.getSize();
+    if (func_size == 0) continue;
+
+    thunk_map.insert({*address, sr});
+  }
+  return thunk_map;
+}
+
 namespace propeller {
 absl::flat_hash_map<uint64_t, llvm::SmallVector<llvm::object::ELFSymbolRef>>
 ReadSymbolTable(const BinaryContent &binary_content) {
@@ -279,6 +305,14 @@ ReadSymbolTable(const BinaryContent &binary_content) {
   return symtab;
 }
 
+std::optional<absl::btree_map<uint64_t, llvm::object::ELFSymbolRef>>
+ReadThunkSymbols(const BinaryContent &binary_content) {
+  if (binary_content.object_file->getArch() == llvm::Triple::aarch64)
+    return ::ReadAArch64ThunkSymbols(binary_content);
+
+  return std::nullopt;
+}
+
 absl::StatusOr<std::vector<llvm::object::BBAddrMap>> ReadBbAddrMap(
     const BinaryContent &binary_content) {
   auto *elf_object = llvm::dyn_cast<llvm::object::ELFObjectFileBase>(

Diff for: propeller/binary_content.h

@@ -7,6 +7,7 @@
 #include <string>
 #include <vector>
 
+#include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -116,6 +117,14 @@ absl::StatusOr<int64_t> GetSymbolAddress(
 absl::flat_hash_map<uint64_t, llvm::SmallVector<llvm::object::ELFSymbolRef>>
 ReadSymbolTable(const BinaryContent &binary_content);
 
+// Returns an AArch64 binary's thunk symbols by reading from its symbol table.
+absl::btree_map<uint64_t, llvm::object::ELFSymbolRef> ReadAArch64ThunkSymbols(
+    const BinaryContent &binary_content);
+
+// Returns the binary's thunk symbols by reading from its symbol table.
+std::optional<absl::btree_map<uint64_t, llvm::object::ELFSymbolRef>>
+ReadThunkSymbols(const BinaryContent &binary_content);
+
 // Returns the binary's `BBAddrMap`s by calling LLVM-side decoding function
 // `ELFObjectFileBase::readBBAddrMap`. Returns error if the call fails or if the
 // result is empty.

Diff for: propeller/binary_content_test.cc

@@ -69,5 +69,23 @@ TEST(GetSymbolAddressTest, SymbolNotFound) {
       Not(IsOk()));
 }
 
+TEST(ThunkSymbolsTest, AArch64Thunks) {
+  const std::string binary = absl::StrCat(
+      ::testing::SrcDir(),
+      "_main/propeller/testdata/fake_thunks.bin");
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BinaryContent> binary_content,
+                       GetBinaryContent(binary));
+  EXPECT_THAT(ReadThunkSymbols(*binary_content), Optional(SizeIs(2)));
+}
+
+TEST(ThunkSymbolsTest, x86NoThunks) {
+  const std::string binary = absl::StrCat(
+      ::testing::SrcDir(),
+      "_main/propeller/testdata/propeller_sample_1.bin");
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BinaryContent> binary_content,
+                       GetBinaryContent(binary));
+  EXPECT_THAT(ReadThunkSymbols(*binary_content), Eq(std::nullopt));
+}
+
 }  // namespace
 }  // namespace propeller

Diff for: propeller/testdata/BUILD

@@ -36,6 +36,7 @@ exports_files([
     "bimodal_sample.x.bin",
     "call_from_simple_loop.protobuf",
     "clang_v0_labels.binary",
+    "fake_thunks.bin",
     "hot_and_cold_landing_pads.protobuf",
     "libro_sample.so",
     "llvm_function_samples.binary",
@@ -352,3 +353,25 @@ genrule(
     target_compatible_with = ["//third_party/bazel_platforms/cpu:aarch64"],
     tools = [":sample"],
 )
+
+cc_binary(
+    name = "fake_thunks",
+    srcs = ["fake_thunks.c"],
+)
+
+genrule(
+    name = "fake_thunks_bin",
+    srcs = ["fake_thunks.c"],
+    outs = [
+        "fake_thunks.bin.gen",
+    ],
+    cmd = "$(CC) $(CC_FLAGS) -g -O2 -Wl,-build-id -pie -fbasic-block-sections=labels $< -o " +
+          "$(RULEDIR)/fake_thunks.bin.gen",
+    exec_compatible_with = ["//third_party/bazel_platforms/cpu:aarch64"],
+    tags = [
+        "manual",
+        "requires-arch:arm",
+    ],
+    target_compatible_with = ["//third_party/bazel_platforms/cpu:aarch64"],
+    toolchains = _LLVM_PROPELLER_TESTDATA_TOOLCHAINS,
+)

Diff for: propeller/testdata/fake_thunks.bin

@@ -0,0 +1,19 @@
+/* fake_thunks.c */
+volatile int x = 1;
+
+__attribute__((noinline)) int __AArch64ADRPThunk_test1(int i) {
+  return x + i;
+}
+
+__attribute__((noinline)) int __AArch64ADRPThunk_test2(int i) {
+  return x + i + 1;
+}
+
+int sample1_func() { return 13; }
+
+int main(void) {
+  __AArch64ADRPThunk_test1(x);
+  __AArch64ADRPThunk_test2(x);
+
+  return 0;
+}