Snap for 6424828 from 901a08a73484b0989530fb9fc7de0829d13d6757 to rvc-release

Change-Id: I3bbfe4f5d361406b4fd073014ca47a25eb2a931f

11 files changed, 370 insertions, 205 deletions

diff --git a/simpleperf/ETMDecoder.cpp b/simpleperf/ETMDecoder.cpp
index 6e65dd59..136bbcdf 100644
--- a/simpleperf/ETMDecoder.cpp
+++ b/simpleperf/ETMDecoder.cpp
@@ -119,10 +119,18 @@ class ETMV4IDecodeTree {
 
 // Similar to IPktDataIn<EtmV4ITrcPacket>, but add trace id.
 struct PacketCallback {
+  // packet callbacks are called in priority order.
+  enum Priority {
+    MAP_LOCATOR,
+    PACKET_TO_ELEMENT,
+  };
+
+  PacketCallback(Priority prio) : priority(prio) {}
   virtual ~PacketCallback() {}
   virtual ocsd_datapath_resp_t ProcessPacket(uint8_t trace_id, ocsd_datapath_op_t op,
                                              ocsd_trc_index_t index_sop,
                                              const EtmV4ITrcPacket* pkt) = 0;
+  const Priority priority;
 };
 
 // Receives packets from a packet decoder in OpenCSD library.
@@ -130,7 +138,13 @@ class PacketSink : public IPktDataIn<EtmV4ITrcPacket> {
  public:
   PacketSink(uint8_t trace_id) : trace_id_(trace_id) {}
 
-  void AddCallback(PacketCallback* callback) { callbacks_.push_back(callback); }
+  void AddCallback(PacketCallback* callback) {
+    auto it = std::lower_bound(callbacks_.begin(), callbacks_.end(), callback,
+                               [](const PacketCallback* c1, const PacketCallback* c2) {
+                                 return c1->priority < c2->priority;
+                               });
+    callbacks_.insert(it, callback);
+  }
 
   ocsd_datapath_resp_t PacketDataIn(ocsd_datapath_op_t op, ocsd_trc_index_t index_sop,
                                     const EtmV4ITrcPacket* pkt) override {
@@ -148,35 +162,92 @@ class PacketSink : public IPktDataIn<EtmV4ITrcPacket> {
   std::vector<PacketCallback*> callbacks_;
 };
 
-// Map (trace_id, ip address) to (binary_path, binary_offset), and read binary files.
-class MemAccess : public ITargetMemAccess {
+// For each trace_id, when given an addr, find the thread and map it belongs to.
+class MapLocator : public PacketCallback {
  public:
-  MemAccess(ThreadTree& thread_tree) : thread_tree_(thread_tree) {}
-
-  void ProcessPacket(uint8_t trace_id, const EtmV4ITrcPacket* packet) {
-    if (packet->getContext().updated_c) {
-      tid_map_[trace_id] = packet->getContext().ctxtID;
-      if (trace_id == trace_id_) {
-        // Invalidate the cached buffer when the last trace stream changes thread.
-        buffer_end_ = 0;
+  MapLocator(ThreadTree& thread_tree)
+      : PacketCallback(PacketCallback::MAP_LOCATOR), thread_tree_(thread_tree) {}
+
+  ThreadTree& GetThreadTree() { return thread_tree_; }
+
+  ocsd_datapath_resp_t ProcessPacket(uint8_t trace_id, ocsd_datapath_op_t op,
+                                     ocsd_trc_index_t index_sop,
+                                     const EtmV4ITrcPacket* pkt) override {
+    TraceData& data = trace_data_[trace_id];
+    if (op == OCSD_OP_DATA) {
+      if (pkt != nullptr && pkt->getContext().updated_c) {
+        int32_t new_tid = static_cast<int32_t>(pkt->getContext().ctxtID);
+        if (data.tid != new_tid) {
+          data.tid = new_tid;
+          data.thread = nullptr;
+          data.userspace_map = nullptr;
+        }
       }
+    } else if (op == OCSD_OP_RESET) {
+      data.tid = -1;
+      data.thread = nullptr;
+      data.userspace_map = nullptr;
+    }
+    return OCSD_RESP_CONT;
+  }
+
+  const MapEntry* FindMap(uint8_t trace_id, uint64_t addr) {
+    TraceData& data = trace_data_[trace_id];
+    if (data.userspace_map != nullptr && data.userspace_map->Contains(addr)) {
+      return data.userspace_map;
+    }
+    if (data.tid == -1) {
+      return nullptr;
+    }
+    if (data.thread == nullptr) {
+      data.thread = thread_tree_.FindThread(data.tid);
+      if (data.thread == nullptr) {
+        return nullptr;
+      }
+    }
+    data.userspace_map = data.thread->maps->FindMapByAddr(addr);
+    if (data.userspace_map != nullptr) {
+      return data.userspace_map;
     }
+    // We don't cache kernel map. Because kernel map can start from 0 and overlap all userspace
+    // maps.
+    return thread_tree_.GetKernelMaps().FindMapByAddr(addr);
   }
 
-  ocsd_err_t ReadTargetMemory(const ocsd_vaddr_t address, uint8_t cs_trace_id, ocsd_mem_space_acc_t,
+ private:
+  struct TraceData {
+    int32_t tid = -1;  // thread id, -1 if invalid
+    const ThreadEntry* thread = nullptr;
+    const MapEntry* userspace_map = nullptr;
+  };
+
+  ThreadTree& thread_tree_;
+  TraceData trace_data_[256];
+};
+
+// Map (trace_id, ip address) to (binary_path, binary_offset), and read binary files.
+class MemAccess : public ITargetMemAccess {
+ public:
+  MemAccess(MapLocator& map_locator) : map_locator_(map_locator) {}
+
+  ocsd_err_t ReadTargetMemory(const ocsd_vaddr_t address, uint8_t trace_id, ocsd_mem_space_acc_t,
                               uint32_t* num_bytes, uint8_t* p_buffer) override {
-    if (cs_trace_id == trace_id_ && address >= buffer_start_ &&
-        address + *num_bytes <= buffer_end_) {
-      if (buffer_ == nullptr) {
+    TraceData& data = trace_data_[trace_id];
+    const MapEntry* map = map_locator_.FindMap(trace_id, address);
+    // fast path
+    if (map != nullptr && map == data.buffer_map && address >= data.buffer_start &&
+        address + *num_bytes <= data.buffer_end) {
+      if (data.buffer == nullptr) {
         *num_bytes = 0;
       } else {
-        memcpy(p_buffer, buffer_ + (address - buffer_start_), *num_bytes);
+        memcpy(p_buffer, data.buffer + (address - data.buffer_start), *num_bytes);
       }
       return OCSD_OK;
     }
 
+    // slow path
     size_t copy_size = 0;
-    if (const MapEntry* map = FindMap(cs_trace_id, address); map != nullptr) {
+    if (map != nullptr) {
       llvm::MemoryBuffer* memory = GetMemoryBuffer(map->dso);
       if (memory != nullptr) {
         uint64_t offset = address - map->start_addr + map->pgoff;
@@ -187,47 +258,36 @@ class MemAccess : public ITargetMemAccess {
         }
       }
       // Update the last buffer cache.
-      trace_id_ = cs_trace_id;
-      buffer_ = memory == nullptr ? nullptr : (memory->getBufferStart() + map->pgoff);
-      buffer_start_ = map->start_addr;
-      buffer_end_ = map->get_end_addr();
+      data.buffer_map = map;
+      data.buffer = memory == nullptr ? nullptr : (memory->getBufferStart() + map->pgoff);
+      data.buffer_start = map->start_addr;
+      data.buffer_end = map->get_end_addr();
     }
     *num_bytes = copy_size;
     return OCSD_OK;
   }
 
  private:
-  const MapEntry* FindMap(uint8_t trace_id, uint64_t address) {
-    if (auto it = tid_map_.find(trace_id); it != tid_map_.end()) {
-      if (ThreadEntry* thread = thread_tree_.FindThread(it->second); thread != nullptr) {
-        if (const MapEntry* map = thread_tree_.FindMap(thread, address);
-            !thread_tree_.IsUnknownDso(map->dso)) {
-          return map;
-        }
-      }
-    }
-    return nullptr;
-  }
-
   llvm::MemoryBuffer* GetMemoryBuffer(Dso* dso) {
-    if (auto it = memory_buffers_.find(dso); it != memory_buffers_.end()) {
-      return it->second.get();
+    auto it = elf_map_.find(dso);
+    if (it == elf_map_.end()) {
+      ElfStatus status;
+      auto res = elf_map_.emplace(dso, ElfFile::Open(dso->GetDebugFilePath(), &status));
+      it = res.first;
     }
-    auto buffer_or_err = llvm::MemoryBuffer::getFile(dso->GetDebugFilePath());
-    llvm::MemoryBuffer* buffer = buffer_or_err ? buffer_or_err.get().release() : nullptr;
-    memory_buffers_.emplace(dso, buffer);
-    return buffer;
+    return it->second ? it->second->GetMemoryBuffer() : nullptr;
   }
 
-  // map from trace id to thread id
-  std::unordered_map<uint8_t, pid_t> tid_map_;
-  ThreadTree& thread_tree_;
-  std::unordered_map<Dso*, std::unique_ptr<llvm::MemoryBuffer>> memory_buffers_;
-  // cache of the last buffer
-  uint8_t trace_id_ = 0;
-  const char* buffer_ = nullptr;
-  uint64_t buffer_start_ = 0;
-  uint64_t buffer_end_ = 0;
+  struct TraceData {
+    const MapEntry* buffer_map = nullptr;
+    const char* buffer = nullptr;
+    uint64_t buffer_start = 0;
+    uint64_t buffer_end = 0;
+  };
+
+  MapLocator& map_locator_;
+  std::unordered_map<Dso*, std::unique_ptr<ElfFile>> elf_map_;
+  TraceData trace_data_[256];
 };
 
 class InstructionDecoder : public TrcIDecode {
@@ -253,9 +313,9 @@ struct ElementCallback {
 // Decode packets into elements.
 class PacketToElement : public PacketCallback, public ITrcGenElemIn {
  public:
-  PacketToElement(ThreadTree& thread_tree, const std::unordered_map<uint8_t, EtmV4Config>& configs,
+  PacketToElement(MapLocator& map_locator, const std::unordered_map<uint8_t, EtmV4Config>& configs,
                   DecodeErrorLogger& error_logger)
-      : mem_access_(thread_tree) {
+      : PacketCallback(PacketCallback::PACKET_TO_ELEMENT), mem_access_(map_locator) {
     for (auto& p : configs) {
       uint8_t trace_id = p.first;
       const EtmV4Config& config = p.second;
@@ -274,9 +334,6 @@ class PacketToElement : public PacketCallback, public ITrcGenElemIn {
   ocsd_datapath_resp_t ProcessPacket(uint8_t trace_id, ocsd_datapath_op_t op,
                                      ocsd_trc_index_t index_sop,
                                      const EtmV4ITrcPacket* pkt) override {
-    if (pkt != nullptr) {
-      mem_access_.ProcessPacket(trace_id, pkt);
-    }
     return element_decoders_[trace_id].PacketDataIn(op, index_sop, pkt);
   }
 
@@ -336,105 +393,82 @@ class DataDumper : public ElementCallback {
   ocsdMsgLogger stdout_logger_;
 };
 
-// Base class for parsing executed instruction ranges from etm data.
-class InstrRangeParser {
- public:
-  InstrRangeParser(ThreadTree& thread_tree, const ETMDecoder::CallbackFn& callback)
-      : thread_tree_(thread_tree), callback_(callback) {}
-
-  virtual ~InstrRangeParser() {}
-
- protected:
-  ThreadTree& thread_tree_;
-  ETMDecoder::CallbackFn callback_;
-};
-
 // It decodes each ETMV4IPacket into TraceElements, and generates ETMInstrRanges from TraceElements.
 // Decoding each packet is slow, but ensures correctness.
-class BasicInstrRangeParser : public InstrRangeParser, public ElementCallback {
+class InstrRangeParser : public ElementCallback {
+ private:
+  struct TraceData {
+    ETMInstrRange instr_range;
+    bool wait_for_branch_to_addr_fix = false;
+  };
+
  public:
-  BasicInstrRangeParser(ThreadTree& thread_tree, const ETMDecoder::CallbackFn& callback)
-      : InstrRangeParser(thread_tree, callback) {}
+  InstrRangeParser(MapLocator& map_locator, const ETMDecoder::CallbackFn& callback)
+      : map_locator_(map_locator), callback_(callback) {}
 
   ocsd_datapath_resp_t ProcessElement(const ocsd_trc_index_t, uint8_t trace_id,
                                       const OcsdTraceElement& elem,
                                       const ocsd_instr_info* next_instr) override {
-    if (elem.getType() == OCSD_GEN_TRC_ELEM_PE_CONTEXT) {
-      if (elem.getContext().ctxt_id_valid) {
-        // trace_id is associated with a new thread.
-        pid_t new_tid = elem.getContext().context_id;
-        auto& tid = tid_map_[trace_id];
-        if (tid != new_tid) {
-          tid = new_tid;
-          if (trace_id == current_map_.trace_id) {
-            current_map_.Invalidate();
-          }
-        }
-      }
-    } else if (elem.getType() == OCSD_GEN_TRC_ELEM_INSTR_RANGE) {
-      if (!FindMap(trace_id, elem.st_addr)) {
+    if (elem.getType() == OCSD_GEN_TRC_ELEM_INSTR_RANGE) {
+      TraceData& data = trace_data_[trace_id];
+      const MapEntry* map = map_locator_.FindMap(trace_id, elem.st_addr);
+      if (map == nullptr) {
+        FlushData(data);
         return OCSD_RESP_CONT;
       }
-      instr_range_.dso = current_map_.map->dso;
-      instr_range_.start_addr = current_map_.ToVaddrInFile(elem.st_addr);
-      instr_range_.end_addr = current_map_.ToVaddrInFile(elem.en_addr - elem.last_instr_sz);
+      uint64_t start_addr = map->GetVaddrInFile(elem.st_addr);
+      auto& instr_range = data.instr_range;
+
+      if (data.wait_for_branch_to_addr_fix) {
+        // OpenCSD may cache a list of InstrRange elements, making it inaccurate to get branch to
+        // address from next_instr->branch_addr. So fix it by using the start address of the next
+        // InstrRange element.
+        instr_range.branch_to_addr = start_addr;
+      }
+      FlushData(data);
+      instr_range.dso = map->dso;
+      instr_range.start_addr = start_addr;
+      instr_range.end_addr = map->GetVaddrInFile(elem.en_addr - elem.last_instr_sz);
       bool end_with_branch =
           elem.last_i_type == OCSD_INSTR_BR || elem.last_i_type == OCSD_INSTR_BR_INDIRECT;
       bool branch_taken = end_with_branch && elem.last_instr_exec;
       if (elem.last_i_type == OCSD_INSTR_BR && branch_taken) {
-        instr_range_.branch_to_addr = current_map_.ToVaddrInFile(next_instr->branch_addr);
+        // It is based on the assumption that we only do immediate branch inside a binary,
+        // which may not be true for all cases. TODO: http://b/151665001.
+        instr_range.branch_to_addr = map->GetVaddrInFile(next_instr->branch_addr);
+        data.wait_for_branch_to_addr_fix = true;
       } else {
-        instr_range_.branch_to_addr = 0;
+        instr_range.branch_to_addr = 0;
       }
-      instr_range_.branch_taken_count = branch_taken ? 1 : 0;
-      instr_range_.branch_not_taken_count = branch_taken ? 0 : 1;
-      callback_(instr_range_);
+      instr_range.branch_taken_count = branch_taken ? 1 : 0;
+      instr_range.branch_not_taken_count = branch_taken ? 0 : 1;
+
+    } else if (elem.getType() == OCSD_GEN_TRC_ELEM_TRACE_ON) {
+      // According to the ETM Specification, the Trace On element indicates a discontinuity in the
+      // instruction trace stream. So it cuts the connection between instr ranges.
+      FlushData(trace_data_[trace_id]);
     }
     return OCSD_RESP_CONT;
   }
 
- private:
-  struct CurrentMap {
-    int trace_id = -1;
-    const MapEntry* map = nullptr;
-    uint64_t addr_in_file = 0;
-
-    void Invalidate() { trace_id = -1; }
-
-    bool IsAddrInMap(uint8_t trace_id, uint64_t addr) {
-      return trace_id == this->trace_id && map != nullptr && addr >= map->start_addr &&
-             addr < map->get_end_addr();
+  void FinishData() {
+    for (auto& pair : trace_data_) {
+      FlushData(pair.second);
     }
+  }
 
-    uint64_t ToVaddrInFile(uint64_t addr) {
-      if (addr >= map->start_addr && addr < map->get_end_addr()) {
-        return addr - map->start_addr + addr_in_file;
-      }
-      return 0;
-    }
-  };
-
-  bool FindMap(uint8_t trace_id, uint64_t addr) {
-    if (current_map_.IsAddrInMap(trace_id, addr)) {
-      return true;
-    }
-    ThreadEntry* thread = thread_tree_.FindThread(tid_map_[trace_id]);
-    if (thread != nullptr) {
-      const MapEntry* map = thread_tree_.FindMap(thread, addr, false);
-      if (map != nullptr && !thread_tree_.IsUnknownDso(map->dso)) {
-        current_map_.trace_id = trace_id;
-        current_map_.map = map;
-        current_map_.addr_in_file =
-            map->dso->IpToVaddrInFile(map->start_addr, map->start_addr, map->pgoff);
-        return true;
-      }
+ private:
+  void FlushData(TraceData& data) {
+    if (data.instr_range.dso != nullptr) {
+      callback_(data.instr_range);
+      data.instr_range.dso = nullptr;
     }
-    return false;
+    data.wait_for_branch_to_addr_fix = false;
   }
 
-  std::unordered_map<uint8_t, pid_t> tid_map_;
-  CurrentMap current_map_;
-  ETMInstrRange instr_range_;
+  MapLocator& map_locator_;
+  std::unordered_map<uint8_t, TraceData> trace_data_;
+  ETMDecoder::CallbackFn callback_;
 };
 
 // Etm data decoding in OpenCSD library has two steps:
@@ -490,9 +524,9 @@ class ETMDecoderImpl : public ETMDecoder {
   }
 
   void RegisterCallback(const CallbackFn& callback) {
-    auto parser = std::make_unique<BasicInstrRangeParser>(thread_tree_, callback);
-    InstallElementCallback(parser.get());
-    instr_range_parser_.reset(parser.release());
+    InstallMapLocator();
+    instr_range_parser_.reset(new InstrRangeParser(*map_locator_, callback));
+    InstallElementCallback(instr_range_parser_.get());
   }
 
   bool ProcessData(const uint8_t* data, size_t size) override {
@@ -525,14 +559,33 @@ class ETMDecoderImpl : public ETMDecoder {
     return true;
   }
 
+  bool FinishData() override {
+    if (instr_range_parser_) {
+      instr_range_parser_->FinishData();
+    }
+    return true;
+  }
+
  private:
+  void InstallMapLocator() {
+    if (!map_locator_) {
+      map_locator_.reset(new MapLocator(thread_tree_));
+      InstallPacketCallback(map_locator_.get());
+    }
+  }
+
+  void InstallPacketCallback(PacketCallback* callback) {
+    for (auto& p : packet_sinks_) {
+      p.second.AddCallback(callback);
+    }
+  }
+
   void InstallElementCallback(ElementCallback* callback) {
     if (!packet_to_element_) {
+      InstallMapLocator();
       packet_to_element_.reset(
-          new PacketToElement(thread_tree_, configs_, decode_tree_.ErrorLogger()));
-      for (auto& p : packet_sinks_) {
-        p.second.AddCallback(packet_to_element_.get());
-      }
+          new PacketToElement(*map_locator_, configs_, decode_tree_.ErrorLogger()));
+      InstallPacketCallback(packet_to_element_.get());
     }
     packet_to_element_->AddCallback(callback);
   }
@@ -550,6 +603,7 @@ class ETMDecoderImpl : public ETMDecoder {
   // an index keeping processed etm data size
   size_t data_index_ = 0;
   std::unique_ptr<InstrRangeParser> instr_range_parser_;
+  std::unique_ptr<MapLocator> map_locator_;
 };
 
 }  // namespace
diff --git a/simpleperf/ETMDecoder.h b/simpleperf/ETMDecoder.h
index 3ebbd7dd..b9493acd 100644
--- a/simpleperf/ETMDecoder.h
+++ b/simpleperf/ETMDecoder.h
@@ -61,6 +61,7 @@ class ETMDecoder {
   virtual void RegisterCallback(const CallbackFn& callback) = 0;
 
   virtual bool ProcessData(const uint8_t* data, size_t size) = 0;
+  virtual bool FinishData() = 0;
 };
 
 }  // namespace simpleperf
\ No newline at end of file
diff --git a/simpleperf/cmd_inject.cpp b/simpleperf/cmd_inject.cpp
index 0a6719d8..ee27b0b8 100644
--- a/simpleperf/cmd_inject.cpp
+++ b/simpleperf/cmd_inject.cpp
@@ -80,6 +80,9 @@ class InjectCommand : public Command {
     if (!record_file_reader_->ReadDataSection([this](auto r) { return ProcessRecord(r.get()); })) {
       return false;
     }
+    if (etm_decoder_ && !etm_decoder_->FinishData()) {
+      return false;
+    }
     PostProcess();
     output_fp_.reset(nullptr);
     return true;
@@ -164,7 +167,7 @@ class InjectCommand : public Command {
       return;
     }
 
-    auto& binary = binary_map_[instr_range.dso->GetDebugFilePath()];
+    auto& binary = binary_map_[instr_range.dso];
     binary.range_count_map[AddrPair(instr_range.start_addr, instr_range.end_addr)] +=
         instr_range.branch_taken_count + instr_range.branch_not_taken_count;
     if (instr_range.branch_taken_count > 0) {
@@ -174,13 +177,23 @@ class InjectCommand : public Command {
   }
 
   void PostProcess() {
-    for (const auto& pair : binary_map_) {
-      const std::string& binary_path = pair.first;
-      const BinaryInfo& binary = pair.second;
+    // binary_map is used to store instruction ranges, which can have a large amount. And it has
+    // a larger access time (instruction ranges * executed time). So it's better to use
+    // unorder_maps to speed up access time. But we also want a stable output here, to compare
+    // output changes result from code changes. So generate a sorted output here.
+    std::vector<Dso*> dso_v;
+    for (auto& p : binary_map_) {
+      dso_v.emplace_back(p.first);
+    }
+    std::sort(dso_v.begin(), dso_v.end(), [](Dso* d1, Dso* d2) { return d1->Path() < d2->Path(); });
+    for (auto dso : dso_v) {
+      const BinaryInfo& binary = binary_map_[dso];
 
       // Write range_count_map.
-      fprintf(output_fp_.get(), "%zu\n", binary.range_count_map.size());
-      for (const auto& pair2 : binary.range_count_map) {
+      std::map<AddrPair, uint64_t> range_count_map(binary.range_count_map.begin(),
+                                                   binary.range_count_map.end());
+      fprintf(output_fp_.get(), "%zu\n", range_count_map.size());
+      for (const auto& pair2 : range_count_map) {
         const AddrPair& addr_range = pair2.first;
         uint64_t count = pair2.second;
 
@@ -192,8 +205,10 @@ class InjectCommand : public Command {
       fprintf(output_fp_.get(), "0\n");
 
       // Write branch_count_map.
-      fprintf(output_fp_.get(), "%zu\n", binary.branch_count_map.size());
-      for (const auto& pair2 : binary.branch_count_map) {
+      std::map<AddrPair, uint64_t> branch_count_map(binary.branch_count_map.begin(),
+                                                    binary.branch_count_map.end());
+      fprintf(output_fp_.get(), "%zu\n", branch_count_map.size());
+      for (const auto& pair2 : branch_count_map) {
         const AddrPair& branch = pair2.first;
         uint64_t count = pair2.second;
 
@@ -202,7 +217,7 @@ class InjectCommand : public Command {
       }
 
       // Write the binary path in comment.
-      fprintf(output_fp_.get(), "// %s\n\n", binary_path.c_str());
+      fprintf(output_fp_.get(), "// %s\n\n", dso->Path().c_str());
     }
   }
 
@@ -217,7 +232,7 @@ class InjectCommand : public Command {
   std::unique_ptr<FILE, decltype(&fclose)> output_fp_;
 
   // Store results for AutoFDO.
-  std::unordered_map<std::string, BinaryInfo> binary_map_;
+  std::unordered_map<Dso*, BinaryInfo> binary_map_;
 };
 
 }  // namespace
diff --git a/simpleperf/cmd_inject_test.cpp b/simpleperf/cmd_inject_test.cpp
index 0d66d6cd..1f057cb0 100644
--- a/simpleperf/cmd_inject_test.cpp
+++ b/simpleperf/cmd_inject_test.cpp
@@ -20,6 +20,7 @@
 #include "command.h"
 #include "get_test_data.h"
 #include "test_util.h"
+#include "utils.h"
 
 static std::unique_ptr<Command> InjectCmd() { return CreateCommandInstance("inject"); }
 
@@ -31,6 +32,10 @@ TEST(cmd_inject, smoke) {
   ASSERT_TRUE(android::base::ReadFileToString(tmpfile.path, &data));
   // Test that we can find instr range in etm_test_loop binary.
   ASSERT_NE(data.find("etm_test_loop"), std::string::npos);
+  std::string expected_data;
+  ASSERT_TRUE(android::base::ReadFileToString(
+      GetTestData(std::string("etm") + OS_PATH_SEPARATOR + "perf_inject.data"), &expected_data));
+  ASSERT_EQ(data, expected_data);
 }
 
 TEST(cmd_inject, binary_option) {
diff --git a/simpleperf/read_apk.cpp b/simpleperf/read_apk.cpp
index a6aae6fe..651a2412 100644
--- a/simpleperf/read_apk.cpp
+++ b/simpleperf/read_apk.cpp
@@ -96,11 +96,7 @@ std::unique_ptr<EmbeddedElf> ApkInspector::FindElfInApkByOffsetWithoutCache(
   }
 
   // We found something in the zip file at the right spot. Is it an ELF?
-  if (lseek(ahelper->GetFd(), found_entry.offset, SEEK_SET) != found_entry.offset) {
-    PLOG(ERROR) << "lseek() failed in " << apk_path << " offset " << found_entry.offset;
-    return nullptr;
-  }
-  if (IsValidElfFile(ahelper->GetFd()) != ElfStatus::NO_ERROR) {
+  if (IsValidElfFile(ahelper->GetFd(), found_entry.offset) != ElfStatus::NO_ERROR) {
     // Omit files that are not ELF files.
     return nullptr;
   }
diff --git a/simpleperf/read_elf.cpp b/simpleperf/read_elf.cpp
index 98cf1937..822d1e3e 100644
--- a/simpleperf/read_elf.cpp
+++ b/simpleperf/read_elf.cpp
@@ -43,6 +43,8 @@
 #define ELF_NOTE_GNU "GNU"
 #define NT_GNU_BUILD_ID 3
 
+using namespace simpleperf;
+
 std::ostream& operator<<(std::ostream& os, const ElfStatus& status) {
   switch (status) {
     case ElfStatus::NO_ERROR:
@@ -78,28 +80,14 @@ bool IsValidElfFileMagic(const char* buf, size_t buf_size) {
   return (buf_size >= 4u && memcmp(buf, elf_magic, 4) == 0);
 }
 
-ElfStatus IsValidElfFile(int fd) {
+ElfStatus IsValidElfFile(int fd, uint64_t file_offset) {
   char buf[4];
-  if (!android::base::ReadFully(fd, buf, 4)) {
+  if (!android::base::ReadFullyAtOffset(fd, buf, 4, file_offset)) {
     return ElfStatus::READ_FAILED;
   }
   return IsValidElfFileMagic(buf, 4) ? ElfStatus::NO_ERROR : ElfStatus::FILE_MALFORMED;
 }
 
-ElfStatus IsValidElfPath(const std::string& filename) {
-  if (!IsRegularFile(filename)) {
-    return ElfStatus::FILE_NOT_FOUND;
-  }
-  std::string mode = std::string("rb") + CLOSE_ON_EXEC_MODE;
-  FILE* fp = fopen(filename.c_str(), mode.c_str());
-  if (fp == nullptr) {
-    return ElfStatus::READ_FAILED;
-  }
-  ElfStatus result = IsValidElfFile(fileno(fp));
-  fclose(fp);
-  return result;
-}
-
 bool GetBuildIdFromNoteSection(const char* section, size_t section_size, BuildId* build_id) {
   const char* p = section;
   const char* end = p + section_size;
@@ -173,15 +161,16 @@ static ElfStatus GetBuildIdFromObjectFile(llvm::object::ObjectFile* obj, BuildId
 }
 
 struct BinaryWrapper {
-  llvm::object::OwningBinary<llvm::object::Binary> binary;
-  llvm::object::ObjectFile* obj;
-
-  BinaryWrapper() : obj(nullptr) {
-  }
+  std::unique_ptr<llvm::MemoryBuffer> buffer;
+  std::unique_ptr<llvm::object::Binary> binary;
+  llvm::object::ObjectFile* obj = nullptr;
 };
 
 static ElfStatus OpenObjectFile(const std::string& filename, uint64_t file_offset,
                                 uint64_t file_size, BinaryWrapper* wrapper) {
+  if (!IsRegularFile(filename)) {
+    return ElfStatus::FILE_NOT_FOUND;
+  }
   android::base::unique_fd fd = FileHelper::OpenReadOnly(filename);
   if (fd == -1) {
     return ElfStatus::READ_FAILED;
@@ -192,6 +181,10 @@ static ElfStatus OpenObjectFile(const std::string& filename, uint64_t file_offse
       return ElfStatus::READ_FAILED;
     }
   }
+  ElfStatus status = IsValidElfFile(fd, file_offset);
+  if (status != ElfStatus::NO_ERROR) {
+    return status;
+  }
   auto buffer_or_err = llvm::MemoryBuffer::getOpenFileSlice(fd, filename, file_size, file_offset);
   if (!buffer_or_err) {
     return ElfStatus::READ_FAILED;
@@ -200,9 +193,9 @@ static ElfStatus OpenObjectFile(const std::string& filename, uint64_t file_offse
   if (!binary_or_err) {
     return ElfStatus::READ_FAILED;
   }
-  wrapper->binary = llvm::object::OwningBinary<llvm::object::Binary>(std::move(binary_or_err.get()),
-                                                                        std::move(buffer_or_err.get()));
-  wrapper->obj = llvm::dyn_cast<llvm::object::ObjectFile>(wrapper->binary.getBinary());
+  wrapper->buffer = std::move(buffer_or_err.get());
+  wrapper->binary = std::move(binary_or_err.get());
+  wrapper->obj = llvm::dyn_cast<llvm::object::ObjectFile>(wrapper->binary.get());
   if (wrapper->obj == nullptr) {
     return ElfStatus::FILE_MALFORMED;
   }
@@ -215,9 +208,9 @@ static ElfStatus OpenObjectFileInMemory(const char* data, size_t size, BinaryWra
   if (!binary_or_err) {
     return ElfStatus::FILE_MALFORMED;
   }
-  wrapper->binary = llvm::object::OwningBinary<llvm::object::Binary>(std::move(binary_or_err.get()),
-                                                                std::move(buffer));
-  wrapper->obj = llvm::dyn_cast<llvm::object::ObjectFile>(wrapper->binary.getBinary());
+  wrapper->buffer = std::move(buffer);
+  wrapper->binary = std::move(binary_or_err.get());
+  wrapper->obj = llvm::dyn_cast<llvm::object::ObjectFile>(wrapper->binary.get());
   if (wrapper->obj == nullptr) {
     return ElfStatus::FILE_MALFORMED;
   }
@@ -225,10 +218,6 @@ static ElfStatus OpenObjectFileInMemory(const char* data, size_t size, BinaryWra
 }
 
 ElfStatus GetBuildIdFromElfFile(const std::string& filename, BuildId* build_id) {
-  ElfStatus result = IsValidElfPath(filename);
-  if (result != ElfStatus::NO_ERROR) {
-    return result;
-  }
   return GetBuildIdFromEmbeddedElfFile(filename, 0, 0, build_id);
 }
 
@@ -442,10 +431,6 @@ ElfStatus MatchBuildId(llvm::object::ObjectFile* obj, const BuildId& expected_bu
 ElfStatus ParseSymbolsFromElfFile(const std::string& filename,
                                   const BuildId& expected_build_id,
                                   const std::function<void(const ElfFileSymbol&)>& callback) {
-  ElfStatus result = IsValidElfPath(filename);
-  if (result != ElfStatus::NO_ERROR) {
-    return result;
-  }
   return ParseSymbolsFromEmbeddedElfFile(filename, 0, 0, expected_build_id, callback);
 }
 
@@ -537,10 +522,6 @@ ElfStatus ReadMinExecutableVirtualAddressFromElfFile(const std::string& filename
                                                      const BuildId& expected_build_id,
                                                      uint64_t* min_vaddr,
                                                      uint64_t* file_offset_of_min_vaddr) {
-  ElfStatus result = IsValidElfPath(filename);
-  if (result != ElfStatus::NO_ERROR) {
-    return result;
-  }
   return ReadMinExecutableVirtualAddressFromEmbeddedElfFile(filename, 0, 0, expected_build_id,
                                                             min_vaddr, file_offset_of_min_vaddr);
 }
@@ -570,12 +551,8 @@ ElfStatus ReadMinExecutableVirtualAddressFromEmbeddedElfFile(const std::string&
 
 ElfStatus ReadSectionFromElfFile(const std::string& filename, const std::string& section_name,
                                  std::string* content) {
-  ElfStatus result = IsValidElfPath(filename);
-  if (result != ElfStatus::NO_ERROR) {
-    return result;
-  }
   BinaryWrapper wrapper;
-  result = OpenObjectFile(filename, 0, 0, &wrapper);
+  ElfStatus result = OpenObjectFile(filename, 0, 0, &wrapper);
   if (result != ElfStatus::NO_ERROR) {
     return result;
   }
@@ -587,3 +564,58 @@ ElfStatus ReadSectionFromElfFile(const std::string& filename, const std::string&
     return ElfStatus::FILE_MALFORMED;
   }
 }
+
+namespace {
+
+template <typename T>
+class ElfFileImpl {};
+
+template <typename ELFT>
+class ElfFileImpl<llvm::object::ELFFile<ELFT>> : public ElfFile {
+ public:
+  ElfFileImpl(BinaryWrapper&& wrapper, const llvm::object::ELFFile<ELFT>* elf)
+      : wrapper_(std::move(wrapper)), elf_(elf) {}
+
+  llvm::MemoryBuffer* GetMemoryBuffer() override {
+    return wrapper_.buffer.get();
+  }
+
+ private:
+  BinaryWrapper wrapper_;
+  const llvm::object::ELFFile<ELFT>* elf_;
+};
+
+}  // namespace
+
+namespace simpleperf {
+
+std::unique_ptr<ElfFile> ElfFile::Open(const std::string& filename, ElfStatus* status) {
+  BinaryWrapper wrapper;
+  auto tuple = SplitUrlInApk(filename);
+  if (std::get<0>(tuple)) {
+    EmbeddedElf* elf = ApkInspector::FindElfInApkByName(std::get<1>(tuple), std::get<2>(tuple));
+    if (elf == nullptr) {
+      *status = ElfStatus::FILE_NOT_FOUND;
+    } else {
+      *status = OpenObjectFile(elf->filepath(), elf->entry_offset(), elf->entry_size(), &wrapper);
+    }
+  } else {
+    *status = OpenObjectFile(filename, 0, 0, &wrapper);
+  }
+  if (*status == ElfStatus::NO_ERROR) {
+    if (auto obj = llvm::dyn_cast<llvm::object::ELF32LEObjectFile>(wrapper.obj)) {
+      using elf_t = std::decay_t<decltype(*obj->getELFFile())>;
+      return std::unique_ptr<ElfFile>(
+          new ElfFileImpl<elf_t>(std::move(wrapper), obj->getELFFile()));
+    }
+    if (auto obj = llvm::dyn_cast<llvm::object::ELF64LEObjectFile>(wrapper.obj)) {
+      using elf_t = std::decay_t<decltype(*obj->getELFFile())>;
+      return std::unique_ptr<ElfFile>(
+          new ElfFileImpl<elf_t>(std::move(wrapper), obj->getELFFile()));
+    }
+    *status = ElfStatus::FILE_MALFORMED;
+  }
+  return nullptr;
+}
+
+}  // namespace simpleperf
\ No newline at end of file
diff --git a/simpleperf/read_elf.h b/simpleperf/read_elf.h
index 1ab9a1f3..f289ee17 100644
--- a/simpleperf/read_elf.h
+++ b/simpleperf/read_elf.h
@@ -84,11 +84,28 @@ ElfStatus ReadMinExecutableVirtualAddressFromEmbeddedElfFile(const std::string&
 ElfStatus ReadSectionFromElfFile(const std::string& filename, const std::string& section_name,
                                  std::string* content);
 
-// Expose the following functions for unit tests.
+namespace llvm {
+class MemoryBuffer;
+}
+
+namespace simpleperf {
+
+class ElfFile {
+ public:
+  static std::unique_ptr<ElfFile> Open(const std::string& filename, ElfStatus* status);
+  virtual ~ElfFile() {}
+
+  virtual llvm::MemoryBuffer* GetMemoryBuffer() = 0;
+
+ protected:
+  ElfFile() {}
+};
+
+}  // namespace simpleperf
+
 bool IsArmMappingSymbol(const char* name);
-ElfStatus IsValidElfFile(int fd);
+ElfStatus IsValidElfFile(int fd, uint64_t file_offset = 0);
 bool IsValidElfFileMagic(const char* buf, size_t buf_size);
-ElfStatus IsValidElfPath(const std::string& filename);
 bool GetBuildIdFromNoteSection(const char* section, size_t section_size, BuildId* build_id);
 
 #endif  // SIMPLE_PERF_READ_ELF_H_
diff --git a/simpleperf/read_elf_test.cpp b/simpleperf/read_elf_test.cpp
index 9f74366c..cc5aa52b 100644
--- a/simpleperf/read_elf_test.cpp
+++ b/simpleperf/read_elf_test.cpp
@@ -29,6 +29,8 @@
 #define ELF_NOTE_GNU "GNU"
 #define NT_GNU_BUILD_ID 3
 
+using namespace simpleperf;
+
 TEST(read_elf, GetBuildIdFromNoteSection) {
   BuildId build_id;
   std::vector<char> data;
@@ -137,7 +139,12 @@ TEST(read_elf, arm_mapping_symbol) {
   ASSERT_FALSE(IsArmMappingSymbol("$a_no_dot"));
 }
 
-TEST(read_elf, IsValidElfPath) {
+TEST(read_elf, ElfFile_Open) {
+  auto IsValidElfPath = [](const std::string& path) {
+    ElfStatus status;
+    ElfFile::Open(path, &status);
+    return status;
+  };
   ASSERT_NE(ElfStatus::NO_ERROR, IsValidElfPath("/dev/zero"));
   TemporaryFile tmp_file;
   ASSERT_EQ(ElfStatus::READ_FAILED, IsValidElfPath(tmp_file.path));
diff --git a/simpleperf/testdata/etm/perf_inject.data b/simpleperf/testdata/etm/perf_inject.data
new file mode 100644
index 00000000..4c121a32
--- /dev/null
+++ b/simpleperf/testdata/etm/perf_inject.data
@@ -0,0 +1,24 @@
+10
+1000-1004:1
+100c-1050:1
+1054-106c:1
+105c-106c:100
+1070-1070:1
+1074-1080:100
+1084-1088:1
+108c-10a0:1
+10a4-10b0:1
+10e0-10ec:1
+0
+9
+1004->100c:1
+1050->10e0:1
+106c->1074:100
+1070->1084:1
+1080->105c:100
+1088->0:1
+10a0->1054:1
+10b0->0:1
+10ec->0:1
+// /data/local/tmp/etm_test_loop
+
diff --git a/simpleperf/thread_tree.cpp b/simpleperf/thread_tree.cpp
index 17061ce7..6babf8cc 100644
--- a/simpleperf/thread_tree.cpp
+++ b/simpleperf/thread_tree.cpp
@@ -195,9 +195,9 @@ void ThreadTree::InsertMap(MapSet& maps, const MapEntry& entry) {
   maps.version++;
 }
 
-static const MapEntry* FindMapByAddr(const MapSet& maps, uint64_t addr) {
-  auto it = maps.maps.upper_bound(addr);
-  if (it != maps.maps.begin()) {
+const MapEntry* MapSet::FindMapByAddr(uint64_t addr) const {
+  auto it = maps.upper_bound(addr);
+  if (it != maps.begin()) {
     --it;
     if (it->second->get_end_addr() > addr) {
       return it->second;
@@ -209,19 +209,19 @@ static const MapEntry* FindMapByAddr(const MapSet& maps, uint64_t addr) {
 const MapEntry* ThreadTree::FindMap(const ThreadEntry* thread, uint64_t ip, bool in_kernel) {
   const MapEntry* result = nullptr;
   if (!in_kernel) {
-    result = FindMapByAddr(*thread->maps, ip);
+    result = thread->maps->FindMapByAddr(ip);
   } else {
-    result = FindMapByAddr(kernel_maps_, ip);
+    result = kernel_maps_.FindMapByAddr(ip);
   }
   return result != nullptr ? result : &unknown_map_;
 }
 
 const MapEntry* ThreadTree::FindMap(const ThreadEntry* thread, uint64_t ip) {
-  const MapEntry* result = FindMapByAddr(*thread->maps, ip);
+  const MapEntry* result = thread->maps->FindMapByAddr(ip);
   if (result != nullptr) {
     return result;
   }
-  result = FindMapByAddr(kernel_maps_, ip);
+  result = kernel_maps_.FindMapByAddr(ip);
   return result != nullptr ? result : &unknown_map_;
 }
 
diff --git a/simpleperf/thread_tree.h b/simpleperf/thread_tree.h
index e9412736..76d3403d 100644
--- a/simpleperf/thread_tree.h
+++ b/simpleperf/thread_tree.h
@@ -58,11 +58,24 @@ struct MapEntry {
   MapEntry() {}
 
   uint64_t get_end_addr() const { return start_addr + len; }
+
+  uint64_t Contains(uint64_t addr) const {
+    return addr >= start_addr && addr < get_end_addr();
+  }
+
+  uint64_t GetVaddrInFile(uint64_t addr) const {
+    if (Contains(addr)) {
+      return dso->IpToVaddrInFile(addr, start_addr, pgoff);
+    }
+    return 0;
+  }
 };
 
 struct MapSet {
   std::map<uint64_t, const MapEntry*> maps;  // Map from start_addr to a MapEntry.
   uint64_t version = 0u;  // incremented each time changing maps
+
+  const MapEntry* FindMapByAddr(uint64_t addr) const;
 };
 
 struct ThreadEntry {
@@ -97,6 +110,7 @@ class ThreadTree {
   void ExitThread(int pid, int tid);
   void AddKernelMap(uint64_t start_addr, uint64_t len, uint64_t pgoff,
                     const std::string& filename);
+  const MapSet& GetKernelMaps() { return kernel_maps_; }
   void AddThreadMap(int pid, int tid, uint64_t start_addr, uint64_t len,
                     uint64_t pgoff, const std::string& filename, uint32_t flags = 0);
   const MapEntry* FindMap(const ThreadEntry* thread, uint64_t ip,