From fddc63383b0b0c848da3a7147327481c88f6c44e Mon Sep 17 00:00:00 2001
From: Teng Qin <qinteng@fb.com>
Date: Wed, 8 Nov 2017 20:04:18 -0800
Subject: [PATCH] Add a multi-type version of iterateSymbolsWithType

Summary:
When using `folly::symbolizer`, it's very often that we want to use `iterateSymbolsWithType` iterate through symbols of a few types using the same callback. Current approach would require iterating the section multiple times.

This Diff adds `iterateSymbolsWithTypes`, which is basically just `iterateSymbolsWithType` but accepts symbol types.

This Diff also updated implementation of `getDefinitionByAddress` and `getSymbolByName` which currently does two iterations for `STT_OBJECT` and `STT_FUNC`.

Reviewed By: yfeldblum

Differential Revision: D6279651

fbshipit-source-id: a661dd15f18e4f2f63dbcca615f5a86d92e528ea
---
 folly/experimental/symbolizer/Elf-inl.h | 13 +++++++++++++
 folly/experimental/symbolizer/Elf.cpp   |  8 ++++----
 folly/experimental/symbolizer/Elf.h     |  6 ++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/folly/experimental/symbolizer/Elf-inl.h b/folly/experimental/symbolizer/Elf-inl.h
index bd9b0a40..3e2f4dcf 100644
--- a/folly/experimental/symbolizer/Elf-inl.h
+++ b/folly/experimental/symbolizer/Elf-inl.h
@@ -97,5 +97,18 @@ const ElfSym* ElfFile::iterateSymbolsWithType(
   });
 }
 
+template <class Fn>
+const ElfSym* ElfFile::iterateSymbolsWithTypes(
+    const ElfShdr& section,
+    std::initializer_list<uint32_t> types,
+    Fn fn) const {
+  // N.B. st_info has the same representation on 32- and 64-bit platforms
+  return iterateSymbols(section, [&](const ElfSym& sym) -> bool {
+    auto const elfType = ELF32_ST_TYPE(sym.st_info);
+    auto const it = std::find(types.begin(), types.end(), elfType);
+    return it != types.end() && fn(sym);
+  });
+}
+
 } // namespace symbolizer
 } // namespace folly
diff --git a/folly/experimental/symbolizer/Elf.cpp b/folly/experimental/symbolizer/Elf.cpp
index a0dfe9ce..7be8f69d 100644
--- a/folly/experimental/symbolizer/Elf.cpp
+++ b/folly/experimental/symbolizer/Elf.cpp
@@ -360,8 +360,8 @@ ElfFile::Symbol ElfFile::getDefinitionByAddress(uintptr_t address) const {
       return false;
     };
 
-    return iterateSymbolsWithType(section, STT_OBJECT, findSymbols) ||
-        iterateSymbolsWithType(section, STT_FUNC, findSymbols);
+    return iterateSymbolsWithTypes(
+        section, {STT_OBJECT, STT_FUNC}, findSymbols);
   };
 
   // Try the .dynsym section first if it exists, it's smaller.
@@ -399,8 +399,8 @@ ElfFile::Symbol ElfFile::getSymbolByName(const char* name) const {
       return false;
     };
 
-    return iterateSymbolsWithType(section, STT_OBJECT, findSymbols) ||
-        iterateSymbolsWithType(section, STT_FUNC, findSymbols);
+    return iterateSymbolsWithTypes(
+        section, {STT_OBJECT, STT_FUNC}, findSymbols);
   };
 
   // Try the .dynsym section first if it exists, it's smaller.
diff --git a/folly/experimental/symbolizer/Elf.h b/folly/experimental/symbolizer/Elf.h
index c097be69..477d8798 100644
--- a/folly/experimental/symbolizer/Elf.h
+++ b/folly/experimental/symbolizer/Elf.h
@@ -23,6 +23,7 @@
 #include <link.h> // For ElfW()
 
 #include <cstdio>
+#include <initializer_list>
 #include <stdexcept>
 #include <system_error>
 
@@ -156,6 +157,11 @@ class ElfFile {
   template <class Fn>
   const ElfSym*
   iterateSymbolsWithType(const ElfShdr& section, uint32_t type, Fn fn) const;
+  template <class Fn>
+  const ElfSym* iterateSymbolsWithTypes(
+      const ElfShdr& section,
+      std::initializer_list<uint32_t> types,
+      Fn fn) const;
 
   /**
    * Find symbol definition by address.
-- 
2.34.1