simpleperf: Support LBR data in branch_list.proto

In branch_list.proto:
   Add messages for LBR data.
In the inject cmd:
1. Support converting perf.data with LBR data to branch list file.
2. Support converting branch list file with LBR data to AutoFDO text
   format.
3. Support merging branch list files with LBR data.
4. Support binary filter in branch list files with LBR data.

Bug: 293953824
Test: run simpleperf_unit_test
Change-Id: Ieacb11f289f1979114b96f588146e1d6f573ebc0

6 files changed, 372 insertions, 124 deletions

diff --git a/simpleperf/BranchListFile.cpp b/simpleperf/BranchListFile.cpp
index e8f6b22c..253dcdbc 100644
--- a/simpleperf/BranchListFile.cpp
+++ b/simpleperf/BranchListFile.cpp
@@ -132,32 +132,8 @@ static UnorderedETMBranchMap BuildUnorderedETMBranchMap(const proto::ETMBinary&
 }
 
 bool StringToETMBinaryMap(const std::string& s, ETMBinaryMap& binary_map) {
-  proto::BranchList branch_list_proto;
-  if (!branch_list_proto.ParseFromString(s)) {
-    PLOG(ERROR) << "failed to read ETMBranchList msg";
-    return false;
-  }
-  if (branch_list_proto.magic() != ETM_BRANCH_LIST_PROTO_MAGIC) {
-    PLOG(ERROR) << "not in etm branch list format in branch_list.proto";
-    return false;
-  }
-
-  for (size_t i = 0; i < branch_list_proto.etm_data_size(); i++) {
-    const auto& binary_proto = branch_list_proto.etm_data(i);
-    BinaryKey key(binary_proto.path(), BuildId(binary_proto.build_id()));
-    if (binary_proto.has_kernel_info()) {
-      key.kernel_start_addr = binary_proto.kernel_info().kernel_start_addr();
-    }
-    ETMBinary& binary = binary_map[key];
-    auto dso_type = ToDsoType(binary_proto.type());
-    if (!dso_type) {
-      LOG(ERROR) << "invalid binary type " << binary_proto.type();
-      return false;
-    }
-    binary.dso_type = dso_type.value();
-    binary.branch_map = BuildUnorderedETMBranchMap(binary_proto);
-  }
-  return true;
+  LBRData lbr_data;
+  return ParseBranchListData(s, binary_map, lbr_data);
 }
 
 class ETMThreadTreeWhenRecording : public ETMThreadTree {
@@ -382,4 +358,83 @@ std::unique_ptr<ETMBranchListGenerator> ETMBranchListGenerator::Create(bool dump
 
 ETMBranchListGenerator::~ETMBranchListGenerator() {}
 
+bool LBRDataToString(const LBRData& data, std::string& s) {
+  proto::BranchList branch_list_proto;
+  branch_list_proto.set_magic(ETM_BRANCH_LIST_PROTO_MAGIC);
+  auto lbr_proto = branch_list_proto.mutable_lbr_data();
+  for (const LBRSample& sample : data.samples) {
+    auto sample_proto = lbr_proto->add_samples();
+    sample_proto->set_binary_id(sample.binary_id);
+    sample_proto->set_vaddr_in_file(sample.vaddr_in_file);
+    for (const LBRBranch& branch : sample.branches) {
+      auto branch_proto = sample_proto->add_branches();
+      branch_proto->set_from_binary_id(branch.from_binary_id);
+      branch_proto->set_to_binary_id(branch.to_binary_id);
+      branch_proto->set_from_vaddr_in_file(branch.from_vaddr_in_file);
+      branch_proto->set_to_vaddr_in_file(branch.to_vaddr_in_file);
+    }
+  }
+  for (const BinaryKey& binary : data.binaries) {
+    auto binary_proto = lbr_proto->add_binaries();
+    binary_proto->set_path(binary.path);
+    binary_proto->set_build_id(binary.build_id.ToString().substr(2));
+  }
+  if (!branch_list_proto.SerializeToString(&s)) {
+    LOG(ERROR) << "failed to serialize lbr data";
+    return false;
+  }
+  return true;
+}
+
+bool ParseBranchListData(const std::string& s, ETMBinaryMap& etm_data, LBRData& lbr_data) {
+  proto::BranchList branch_list_proto;
+  if (!branch_list_proto.ParseFromString(s)) {
+    PLOG(ERROR) << "failed to read ETMBranchList msg";
+    return false;
+  }
+  if (branch_list_proto.magic() != ETM_BRANCH_LIST_PROTO_MAGIC) {
+    PLOG(ERROR) << "not in etm branch list format in branch_list.proto";
+    return false;
+  }
+  for (size_t i = 0; i < branch_list_proto.etm_data_size(); i++) {
+    const auto& binary_proto = branch_list_proto.etm_data(i);
+    BinaryKey key(binary_proto.path(), BuildId(binary_proto.build_id()));
+    if (binary_proto.has_kernel_info()) {
+      key.kernel_start_addr = binary_proto.kernel_info().kernel_start_addr();
+    }
+    ETMBinary& binary = etm_data[key];
+    auto dso_type = ToDsoType(binary_proto.type());
+    if (!dso_type) {
+      LOG(ERROR) << "invalid binary type " << binary_proto.type();
+      return false;
+    }
+    binary.dso_type = dso_type.value();
+    binary.branch_map = BuildUnorderedETMBranchMap(binary_proto);
+  }
+  if (branch_list_proto.has_lbr_data()) {
+    const auto& lbr_data_proto = branch_list_proto.lbr_data();
+    lbr_data.samples.resize(lbr_data_proto.samples_size());
+    for (size_t i = 0; i < lbr_data_proto.samples_size(); ++i) {
+      const auto& sample_proto = lbr_data_proto.samples(i);
+      LBRSample& sample = lbr_data.samples[i];
+      sample.binary_id = sample_proto.binary_id();
+      sample.vaddr_in_file = sample_proto.vaddr_in_file();
+      sample.branches.resize(sample_proto.branches_size());
+      for (size_t j = 0; j < sample_proto.branches_size(); ++j) {
+        const auto& branch_proto = sample_proto.branches(j);
+        LBRBranch& branch = sample.branches[j];
+        branch.from_binary_id = branch_proto.from_binary_id();
+        branch.to_binary_id = branch_proto.to_binary_id();
+        branch.from_vaddr_in_file = branch_proto.from_vaddr_in_file();
+        branch.to_vaddr_in_file = branch_proto.to_vaddr_in_file();
+      }
+    }
+    for (size_t i = 0; i < lbr_data_proto.binaries_size(); ++i) {
+      const auto& binary_proto = lbr_data_proto.binaries(i);
+      lbr_data.binaries.emplace_back(binary_proto.path(), BuildId(binary_proto.build_id()));
+    }
+  }
+  return true;
+}
+
 }  // namespace simpleperf
diff --git a/simpleperf/BranchListFile.h b/simpleperf/BranchListFile.h
index 031a026f..de97ef88 100644
--- a/simpleperf/BranchListFile.h
+++ b/simpleperf/BranchListFile.h
@@ -129,7 +129,6 @@ struct ETMBinary {
 };
 
 using ETMBinaryMap = std::unordered_map<BinaryKey, ETMBinary, BinaryKeyHash>;
-
 bool ETMBinaryMapToString(const ETMBinaryMap& binary_map, std::string& s);
 bool StringToETMBinaryMap(const std::string& s, ETMBinaryMap& binary_map);
 
@@ -167,6 +166,9 @@ struct LBRData {
   std::vector<BinaryKey> binaries;
 };
 
+bool LBRDataToString(const LBRData& data, std::string& s);
+bool ParseBranchListData(const std::string& s, ETMBinaryMap& etm_data, LBRData& lbr_data);
+
 // for testing
 std::string ETMBranchToProtoString(const std::vector<bool>& branch);
 std::vector<bool> ProtoStringToETMBranch(const std::string& s, size_t bit_size);
diff --git a/simpleperf/branch_list.proto b/simpleperf/branch_list.proto
index 29247933..6d6bfb17 100644
--- a/simpleperf/branch_list.proto
+++ b/simpleperf/branch_list.proto
@@ -26,6 +26,7 @@ message BranchList {
   // Should always be "simpleperf:EtmBranchList".
   string magic = 1;
   repeated ETMBinary etm_data = 2;
+  LBRData lbr_data = 3;
 }
 
 message ETMBinary {
@@ -65,3 +66,31 @@ message ETMBinary {
 
   KernelBinaryInfo kernel_info = 5;
 }
+
+message LBRData {
+  repeated Sample samples = 1;
+  repeated Binary binaries = 2;
+
+  message Sample {
+    // If binary_id >= 1, it refers to LBRData.binaries[binary_id - 1]. Otherwise, it's invalid.
+    uint32 binary_id = 1;
+    uint64 vaddr_in_file = 2;
+    repeated Branch branches = 3;
+
+    message Branch {
+      // If from_binary_id >= 1, it refers to LBRData.binaries[from_binary_id - 1]. Otherwise, it's
+      // invalid.
+      uint32 from_binary_id = 1;
+      // If to_binary_id >= 1, it refers to LBRData.binaries[to_binary_id - 1]. Otherwise, it's
+      // invalid.
+      uint32 to_binary_id = 2;
+      uint64 from_vaddr_in_file = 3;
+      uint64 to_vaddr_in_file = 4;
+    }
+  }
+
+  message Binary {
+    string path = 1;
+    string build_id = 2;
+  }
+}
diff --git a/simpleperf/cmd_inject.cpp b/simpleperf/cmd_inject.cpp
index be8f4173..182f4fd2 100644
--- a/simpleperf/cmd_inject.cpp
+++ b/simpleperf/cmd_inject.cpp
@@ -146,9 +146,9 @@ class PerfDataReader {
 
   std::string GetDataType() const { return GetDataType(*reader_); }
 
-  void SetCallback(const AutoFDOBinaryCallback& callback) { autofdo_callback_ = callback; }
-  virtual void SetCallback(const ETMBinaryCallback&) {}
-  virtual void SetCallback(const LBRDataCallback&) {}
+  void AddCallback(const AutoFDOBinaryCallback& callback) { autofdo_callback_ = callback; }
+  void AddCallback(const ETMBinaryCallback& callback) { etm_binary_callback_ = callback; }
+  void AddCallback(const LBRDataCallback& callback) { lbr_data_callback_ = callback; }
 
   virtual bool Read() {
     if (exclude_perf_) {
@@ -196,6 +196,8 @@ class PerfDataReader {
   std::optional<int> exclude_pid_;
   ThreadTree thread_tree_;
   AutoFDOBinaryCallback autofdo_callback_;
+  ETMBinaryCallback etm_binary_callback_;
+  LBRDataCallback lbr_data_callback_;
   // Store results for AutoFDO.
   std::unordered_map<const Dso*, AutoFDOBinaryInfo> autofdo_binary_map_;
 };
@@ -231,8 +233,6 @@ class ETMPerfDataReader : public PerfDataReader {
         etm_dump_option_(etm_dump_option),
         etm_thread_tree_(thread_tree_, exclude_pid_) {}
 
-  void SetCallback(const ETMBinaryCallback& callback) override { etm_binary_callback_ = callback; }
-
   bool Read() override {
     if (reader_->HasFeature(PerfFileFormat::FEAT_ETM_BRANCH_LIST)) {
       return ProcessETMBranchListFeature();
@@ -363,7 +363,6 @@ class ETMPerfDataReader : public PerfDataReader {
   }
 
   ETMDumpOption etm_dump_option_;
-  ETMBinaryCallback etm_binary_callback_;
   ETMThreadTreeWithFilter etm_thread_tree_;
   std::vector<uint8_t> aux_data_buffer_;
   std::unique_ptr<ETMDecoder> etm_decoder_;
@@ -372,12 +371,49 @@ class ETMPerfDataReader : public PerfDataReader {
   std::unordered_map<Dso*, ETMBinary> etm_binary_map_;
 };
 
+static std::optional<std::vector<AutoFDOBinaryInfo>> ConvertLBRDataToAutoFDO(
+    const LBRData& lbr_data) {
+  std::vector<AutoFDOBinaryInfo> binaries(lbr_data.binaries.size());
+  for (const LBRSample& sample : lbr_data.samples) {
+    if (sample.binary_id != 0) {
+      if (sample.binary_id > binaries.size()) {
+        LOG(ERROR) << "binary_id out of range";
+        return std::nullopt;
+      }
+      binaries[sample.binary_id - 1].AddAddress(sample.vaddr_in_file);
+    }
+    for (size_t i = 0; i < sample.branches.size(); ++i) {
+      const LBRBranch& branch = sample.branches[i];
+      if (branch.from_binary_id == 0) {
+        continue;
+      }
+      if (branch.from_binary_id > binaries.size()) {
+        LOG(ERROR) << "binary_id out of range";
+        return std::nullopt;
+      }
+      if (branch.from_binary_id == branch.to_binary_id) {
+        binaries[branch.from_binary_id - 1].AddBranch(branch.from_vaddr_in_file,
+                                                      branch.to_vaddr_in_file);
+      }
+      if (i > 0 && branch.from_binary_id == sample.branches[i - 1].to_binary_id) {
+        uint64_t begin = sample.branches[i - 1].to_vaddr_in_file;
+        uint64_t end = branch.from_vaddr_in_file;
+        // Use the same logic to skip bogus LBR data as AutoFDO.
+        if (end < begin || end - begin > (1 << 20)) {
+          continue;
+        }
+        binaries[branch.from_binary_id - 1].AddRange(begin, end);
+      }
+    }
+  }
+  return binaries;
+}
+
 class LBRPerfDataReader : public PerfDataReader {
  public:
   LBRPerfDataReader(std::unique_ptr<RecordFileReader> reader, bool exclude_perf,
                     const RegEx* binary_name_regex)
       : PerfDataReader(std::move(reader), exclude_perf, binary_name_regex) {}
-  void SetCallback(const LBRDataCallback& callback) override { lbr_data_callback_ = callback; }
 
  private:
   bool ProcessRecord(Record& r) override {
@@ -419,8 +455,16 @@ class LBRPerfDataReader : public PerfDataReader {
 
   bool PostProcess() override {
     if (autofdo_callback_) {
-      ConvertLBRDataToAutoFDO();
+      std::optional<std::vector<AutoFDOBinaryInfo>> binaries = ConvertLBRDataToAutoFDO(lbr_data_);
+      if (!binaries) {
+        return false;
+      }
+      for (const auto& [dso, binary_id] : dso_map_) {
+        autofdo_binary_map_[dso] = std::move(binaries.value()[binary_id - 1]);
+      }
       ProcessAutoFDOBinaryInfo();
+    } else if (lbr_data_callback_) {
+      lbr_data_callback_(lbr_data_);
     }
     return true;
   }
@@ -440,55 +484,25 @@ class LBRPerfDataReader : public PerfDataReader {
     if (auto it = dso_map_.find(dso); it != dso_map_.end()) {
       return it->second;
     }
-    uint32_t binary_id = static_cast<uint32_t>(lbr_data_.binaries.size() + 1);
+    lbr_data_.binaries.emplace_back(dso, 0);
+    uint32_t binary_id = static_cast<uint32_t>(lbr_data_.binaries.size());
     dso_map_[dso] = binary_id;
     return binary_id;
   }
 
-  void ConvertLBRDataToAutoFDO() {
-    std::vector<AutoFDOBinaryInfo> binaries(dso_map_.size());
-    for (const LBRSample& sample : lbr_data_.samples) {
-      if (sample.binary_id != 0) {
-        binaries[sample.binary_id - 1].AddAddress(sample.vaddr_in_file);
-      }
-      for (size_t i = 0; i < sample.branches.size(); ++i) {
-        const LBRBranch& branch = sample.branches[i];
-        if (branch.from_binary_id == 0) {
-          continue;
-        }
-        if (branch.from_binary_id == branch.to_binary_id) {
-          binaries[branch.from_binary_id - 1].AddBranch(branch.from_vaddr_in_file,
-                                                        branch.to_vaddr_in_file);
-        }
-        if (i > 0 && branch.from_binary_id == sample.branches[i - 1].to_binary_id) {
-          uint64_t begin = sample.branches[i - 1].to_vaddr_in_file;
-          uint64_t end = branch.from_vaddr_in_file;
-          // Use the same logic to skip bogus LBR data as AutoFDO.
-          if (end < begin || end - begin > (1 << 20)) {
-            continue;
-          }
-          binaries[branch.from_binary_id - 1].AddRange(begin, end);
-        }
-      }
-    }
-    for (const auto& [dso, binary_id] : dso_map_) {
-      autofdo_binary_map_[dso] = std::move(binaries[binary_id - 1]);
-    }
-  }
-
-  LBRDataCallback lbr_data_callback_;
   LBRData lbr_data_;
   // Map from dso to binary_id in lbr_data_.
   std::unordered_map<const Dso*, uint32_t> dso_map_;
 };
 
-// Read a protobuf file specified by etm_branch_list.proto, and generate ETMBinary.
-class ETMBranchListReader {
+// Read a protobuf file specified by branch_list.proto.
+class BranchListReader {
  public:
-  ETMBranchListReader(const std::string& filename, const RegEx* binary_name_regex)
+  BranchListReader(const std::string& filename, const RegEx* binary_name_regex)
       : filename_(filename), binary_filter_(binary_name_regex) {}
 
-  void SetCallback(const ETMBinaryCallback& callback) { callback_ = callback; }
+  void AddCallback(const ETMBinaryCallback& callback) { etm_binary_callback_ = callback; }
+  void AddCallback(const LBRDataCallback& callback) { lbr_data_callback_ = callback; }
 
   bool Read() {
     std::string s;
@@ -496,24 +510,84 @@ class ETMBranchListReader {
       PLOG(ERROR) << "failed to read " << filename_;
       return false;
     }
-    ETMBinaryMap binary_map;
-    if (!StringToETMBinaryMap(s, binary_map)) {
+    ETMBinaryMap etm_data;
+    LBRData lbr_data;
+    if (!ParseBranchListData(s, etm_data, lbr_data)) {
       PLOG(ERROR) << "file is in wrong format: " << filename_;
       return false;
     }
-    for (auto& [key, binary] : binary_map) {
+    if (etm_binary_callback_ && !etm_data.empty()) {
+      ProcessETMData(etm_data);
+    }
+    if (lbr_data_callback_ && !lbr_data.samples.empty()) {
+      ProcessLBRData(lbr_data);
+    }
+    return true;
+  }
+
+ private:
+  void ProcessETMData(ETMBinaryMap& etm_data) {
+    for (auto& [key, binary] : etm_data) {
       if (!binary_filter_.Filter(key.path)) {
         continue;
       }
-      callback_(key, binary);
+      etm_binary_callback_(key, binary);
     }
-    return true;
   }
 
- private:
+  void ProcessLBRData(LBRData& lbr_data) {
+    // 1. Check if we need to remove binaries.
+    std::vector<uint32_t> new_ids(lbr_data.binaries.size());
+    uint32_t next_id = 1;
+
+    for (size_t i = 0; i < lbr_data.binaries.size(); ++i) {
+      if (!binary_filter_.Filter(lbr_data.binaries[i].path)) {
+        new_ids[i] = 0;
+      } else {
+        new_ids[i] = next_id++;
+      }
+    }
+
+    if (next_id <= lbr_data.binaries.size()) {
+      // 2. Modify lbr_data.binaries.
+      for (size_t i = 0; i < lbr_data.binaries.size(); ++i) {
+        if (new_ids[i] != 0) {
+          size_t new_pos = new_ids[i] - 1;
+          lbr_data.binaries[new_pos] = lbr_data.binaries[i];
+        }
+      }
+      lbr_data.binaries.resize(next_id - 1);
+
+      // 3. Modify lbr_data.samples.
+      auto convert_id = [&](uint32_t& binary_id) {
+        if (binary_id != 0) {
+          binary_id = (binary_id <= new_ids.size()) ? new_ids[binary_id - 1] : 0;
+        }
+      };
+      std::vector<LBRSample> new_samples;
+      for (LBRSample& sample : lbr_data.samples) {
+        convert_id(sample.binary_id);
+        bool has_valid_binary_id = sample.binary_id != 0;
+        for (LBRBranch& branch : sample.branches) {
+          convert_id(branch.from_binary_id);
+          convert_id(branch.to_binary_id);
+          if (branch.from_binary_id != 0 || branch.to_binary_id != 0) {
+            has_valid_binary_id = true;
+          }
+        }
+        if (has_valid_binary_id) {
+          new_samples.emplace_back(std::move(sample));
+        }
+      }
+      lbr_data.samples = std::move(new_samples);
+    }
+    lbr_data_callback_(lbr_data);
+  }
+
   const std::string filename_;
   BinaryFilter binary_filter_;
-  ETMBinaryCallback callback_;
+  ETMBinaryCallback etm_binary_callback_;
+  LBRDataCallback lbr_data_callback_;
 };
 
 // Convert ETMBinary into AutoFDOBinaryInfo.
@@ -662,42 +736,82 @@ class AutoFDOWriter {
   std::unordered_map<BinaryKey, AutoFDOBinaryInfo, BinaryKeyHash> binary_map_;
 };
 
-// Merge ETMBinary.
-struct ETMBranchListMerger {
+// Merge branch list data.
+struct BranchListMerger {
   void AddETMBinary(const BinaryKey& key, ETMBinary& binary) {
-    auto it = binary_map.find(key);
-    if (it == binary_map.end()) {
-      binary_map[key] = std::move(binary);
-    } else {
+    if (auto it = etm_data_.find(key); it != etm_data_.end()) {
       it->second.Merge(binary);
+    } else {
+      etm_data_[key] = std::move(binary);
     }
   }
 
-  ETMBinaryMap binary_map;
-};
+  void AddLBRData(LBRData& lbr_data) {
+    // 1. Merge binaries.
+    std::vector<uint32_t> new_ids(lbr_data.binaries.size());
+    for (size_t i = 0; i < lbr_data.binaries.size(); i++) {
+      const BinaryKey& key = lbr_data.binaries[i];
+      if (auto it = lbr_binary_id_map_.find(key); it != lbr_binary_id_map_.end()) {
+        new_ids[i] = it->second;
+      } else {
+        uint32_t next_id = static_cast<uint32_t>(lbr_binary_id_map_.size()) + 1;
+        new_ids[i] = next_id;
+        lbr_binary_id_map_[key] = next_id;
+        lbr_data_.binaries.emplace_back(key);
+      }
+    }
 
-// Write branch lists to a protobuf file specified by etm_branch_list.proto.
-class ETMBranchListWriter {
- public:
-  bool Write(const std::string& output_filename, const ETMBinaryMap& binary_map) {
-    // Don't produce empty output file.
-    if (binary_map.empty()) {
-      LOG(INFO) << "Skip empty output file.";
-      unlink(output_filename.c_str());
-      return true;
+    // 2. Merge samples.
+    auto convert_id = [&](uint32_t& binary_id) {
+      if (binary_id != 0) {
+        binary_id = (binary_id <= new_ids.size()) ? new_ids[binary_id - 1] : 0;
+      }
+    };
+
+    for (LBRSample& sample : lbr_data.samples) {
+      convert_id(sample.binary_id);
+      for (LBRBranch& branch : sample.branches) {
+        convert_id(branch.from_binary_id);
+        convert_id(branch.to_binary_id);
+      }
+      lbr_data_.samples.emplace_back(std::move(sample));
     }
-    std::string s;
-    if (!ETMBinaryMapToString(binary_map, s)) {
-      LOG(ERROR) << "invalid ETMBinaryMap";
+  }
+
+  ETMBinaryMap& GetETMData() { return etm_data_; }
+
+  LBRData& GetLBRData() { return lbr_data_; }
+
+ private:
+  ETMBinaryMap etm_data_;
+  LBRData lbr_data_;
+  std::unordered_map<BinaryKey, uint32_t, BinaryKeyHash> lbr_binary_id_map_;
+};
+
+// Write branch lists to a protobuf file specified by branch_list.proto.
+static bool WriteBranchListFile(const std::string& output_filename, const ETMBinaryMap& etm_data,
+                                const LBRData& lbr_data) {
+  std::string s;
+  if (!etm_data.empty()) {
+    if (!ETMBinaryMapToString(etm_data, s)) {
       return false;
     }
-    if (!android::base::WriteStringToFile(s, output_filename)) {
-      PLOG(ERROR) << "failed to write to " << output_filename;
+  } else if (!lbr_data.samples.empty()) {
+    if (!LBRDataToString(lbr_data, s)) {
       return false;
     }
+  } else {
+    // Don't produce empty output file.
+    LOG(INFO) << "Skip empty output file.";
+    unlink(output_filename.c_str());
     return true;
   }
-};
+  if (!android::base::WriteStringToFile(s, output_filename)) {
+    PLOG(ERROR) << "failed to write to " << output_filename;
+    return false;
+  }
+  return true;
+}
 
 class InjectCommand : public Command {
  public:
@@ -877,7 +991,7 @@ class InjectCommand : public Command {
     auto afdo_callback = [&](const BinaryKey& key, AutoFDOBinaryInfo& binary) {
       autofdo_writer.AddAutoFDOBinary(key, binary);
     };
-    auto reader_callback = [&](PerfDataReader& reader) { reader.SetCallback(afdo_callback); };
+    auto reader_callback = [&](PerfDataReader& reader) { reader.AddCallback(afdo_callback); };
     if (!ReadPerfDataFiles(reader_callback)) {
       return false;
     }
@@ -885,36 +999,42 @@ class InjectCommand : public Command {
   }
 
   bool ConvertPerfDataToBranchList() {
-    ETMBranchListMerger branch_list_merger;
+    BranchListMerger merger;
     auto etm_callback = [&](const BinaryKey& key, ETMBinary& binary) {
-      branch_list_merger.AddETMBinary(key, binary);
+      merger.AddETMBinary(key, binary);
+    };
+    auto lbr_callback = [&](LBRData& lbr_data) { merger.AddLBRData(lbr_data); };
+
+    auto reader_callback = [&](PerfDataReader& reader) {
+      reader.AddCallback(etm_callback);
+      reader.AddCallback(lbr_callback);
     };
-    auto reader_callback = [&](PerfDataReader& reader) { reader.SetCallback(etm_callback); };
     if (!ReadPerfDataFiles(reader_callback)) {
       return false;
     }
-    ETMBranchListWriter branch_list_writer;
-    return branch_list_writer.Write(output_filename_, branch_list_merger.binary_map);
+    return WriteBranchListFile(output_filename_, merger.GetETMData(), merger.GetLBRData());
   }
 
   bool ConvertBranchListToAutoFDO() {
     // Step1 : Merge branch lists from all input files.
-    ETMBranchListMerger branch_list_merger;
-    auto callback = [&](const BinaryKey& key, ETMBinary& binary) {
-      branch_list_merger.AddETMBinary(key, binary);
+    BranchListMerger merger;
+    auto etm_callback = [&](const BinaryKey& key, ETMBinary& binary) {
+      merger.AddETMBinary(key, binary);
     };
+    auto lbr_callback = [&](LBRData& lbr_data) { merger.AddLBRData(lbr_data); };
     for (const auto& input_filename : input_filenames_) {
-      ETMBranchListReader reader(input_filename, binary_name_regex_.get());
-      reader.SetCallback(callback);
+      BranchListReader reader(input_filename, binary_name_regex_.get());
+      reader.AddCallback(etm_callback);
+      reader.AddCallback(lbr_callback);
       if (!reader.Read()) {
         return false;
       }
     }
 
-    // Step2: Convert ETMBinary to AutoFDOBinaryInfo.
+    // Step2: Convert ETMBinary and LBRData to AutoFDOBinaryInfo.
     AutoFDOWriter autofdo_writer;
     ETMBranchListToAutoFDOConverter converter;
-    for (auto& p : branch_list_merger.binary_map) {
+    for (auto& p : merger.GetETMData()) {
       const BinaryKey& key = p.first;
       ETMBinary& binary = p.second;
       std::unique_ptr<AutoFDOBinaryInfo> autofdo_binary = converter.Convert(key, binary);
@@ -924,6 +1044,16 @@ class InjectCommand : public Command {
         autofdo_writer.AddAutoFDOBinary(BinaryKey(key.path, key.build_id), *autofdo_binary);
       }
     }
+    if (!merger.GetLBRData().samples.empty()) {
+      LBRData& lbr_data = merger.GetLBRData();
+      std::optional<std::vector<AutoFDOBinaryInfo>> binaries = ConvertLBRDataToAutoFDO(lbr_data);
+      if (!binaries) {
+        return false;
+      }
+      for (size_t i = 0; i < binaries.value().size(); ++i) {
+        autofdo_writer.AddAutoFDOBinary(lbr_data.binaries[i], binaries.value()[i]);
+      }
+    }
 
     // Step3: Write AutoFDOBinaryInfo.
     return autofdo_writer.Write(output_filename_);
@@ -931,20 +1061,21 @@ class InjectCommand : public Command {
 
   bool ConvertBranchListToBranchList() {
     // Step1 : Merge branch lists from all input files.
-    ETMBranchListMerger branch_list_merger;
-    auto callback = [&](const BinaryKey& key, ETMBinary& binary) {
-      branch_list_merger.AddETMBinary(key, binary);
+    BranchListMerger merger;
+    auto etm_callback = [&](const BinaryKey& key, ETMBinary& binary) {
+      merger.AddETMBinary(key, binary);
     };
+    auto lbr_callback = [&](LBRData& lbr_data) { merger.AddLBRData(lbr_data); };
     for (const auto& input_filename : input_filenames_) {
-      ETMBranchListReader reader(input_filename, binary_name_regex_.get());
-      reader.SetCallback(callback);
+      BranchListReader reader(input_filename, binary_name_regex_.get());
+      reader.AddCallback(etm_callback);
+      reader.AddCallback(lbr_callback);
       if (!reader.Read()) {
         return false;
       }
     }
     // Step2: Write ETMBinary.
-    ETMBranchListWriter branch_list_writer;
-    return branch_list_writer.Write(output_filename_, branch_list_merger.binary_map);
+    return WriteBranchListFile(output_filename_, merger.GetETMData(), merger.GetLBRData());
   }
 
   std::unique_ptr<RegEx> binary_name_regex_;
diff --git a/simpleperf/cmd_inject_test.cpp b/simpleperf/cmd_inject_test.cpp
index 3b25553b..bdc2820f 100644
--- a/simpleperf/cmd_inject_test.cpp
+++ b/simpleperf/cmd_inject_test.cpp
@@ -226,12 +226,40 @@ TEST(cmd_inject, accept_missing_aux_data) {
 }
 
 TEST(cmd_inject, read_lbr_data) {
+  // Convert perf.data to AutoFDO text format.
+  std::string perf_data_path = GetTestData("lbr/perf_lbr.data");
   std::string data;
-  ASSERT_TRUE(RunInjectCmd({"-i", GetTestData("lbr/perf_lbr.data")}, &data));
+  ASSERT_TRUE(RunInjectCmd({"-i", perf_data_path}, &data));
   data.erase(std::remove(data.begin(), data.end(), '\r'), data.end());
 
   std::string expected_data;
   ASSERT_TRUE(android::base::ReadFileToString(
       GetTestData(std::string("lbr") + OS_PATH_SEPARATOR + "inject_lbr.data"), &expected_data));
-  ASSERT_NE(data.find(expected_data), data.npos);
+  ASSERT_EQ(data, expected_data);
+
+  // Convert perf.data to branch_list.proto format.
+  // Then convert branch_list.proto format to AutoFDO text format.
+  TemporaryFile branch_list_file;
+  close(branch_list_file.release());
+  ASSERT_TRUE(
+      RunInjectCmd({"-i", perf_data_path, "--output", "branch-list", "-o", branch_list_file.path}));
+  ASSERT_TRUE(RunInjectCmd({"-i", branch_list_file.path}, &data));
+  ASSERT_EQ(data, expected_data);
+
+  // Test binary filter on LBR data.
+  ASSERT_TRUE(RunInjectCmd({"-i", perf_data_path, "--binary", "no_lbr_test_loop"}, &data));
+  ASSERT_EQ(data.find("lbr_test_loop"), data.npos);
+
+  // Test binary filter on branch list file.
+  ASSERT_TRUE(RunInjectCmd({"-i", branch_list_file.path, "--binary", "no_lbr_test_loop"}, &data));
+  ASSERT_EQ(data.find("lbr_test_loop"), data.npos);
+
+  // Test multiple input files.
+  ASSERT_TRUE(RunInjectCmd(
+      {
+          "-i",
+          std::string(branch_list_file.path) + "," + branch_list_file.path,
+      },
+      &data));
+  ASSERT_NE(data.find("194d->1940:706"), data.npos);
 }
diff --git a/simpleperf/testdata/lbr/inject_lbr.data b/simpleperf/testdata/lbr/inject_lbr.data
index ffc9b5b5..6f141648 100644
--- a/simpleperf/testdata/lbr/inject_lbr.data
+++ b/simpleperf/testdata/lbr/inject_lbr.data
@@ -9,3 +9,6 @@
 2
 191b->1910:32
 194d->1940:353
+// build_id: 0x0000000000000000000000000000000000000000
+// /home/yabinc/lbr_test_loop
+