Add segment-aware VM addressing for universal binary support

Introduces per-arch VM address spaces to allow multiple Mach-O
architecture slices to coexist without address conflicts. RangeSink
gains an arch_index field which gets passed to all ForEachLoadCommand
calls. DualMap holds a map<int, RangeMap> for VM space per arch_index
while the file map remains shared.

Author: rjmansfield

src/bloaty.cc

@@ -1149,7 +1149,7 @@ bool RangeSink::IsVerboseForVMRange(uint64_t vmaddr, uint64_t vmsize) {
     RangeMap vm_map;
     RangeMap file_map;
     bool contains = false;
-    vm_map.AddRangeWithTranslation(vmaddr, vmsize, "", translator_->vm_map,
+    vm_map.AddRangeWithTranslation(vmaddr, vmsize, "", translator_->GetVMMap(arch_index_),
                                    false, &file_map);
     file_map.ForEachRange(
         [this, &contains](uint64_t fileoff, uint64_t filesize) {
@@ -1182,7 +1182,7 @@ bool RangeSink::IsVerboseForFileRange(uint64_t fileoff, uint64_t filesize) {
     RangeMap file_map;
     bool contains = false;
     file_map.AddRangeWithTranslation(fileoff, filesize, "",
-                                     translator_->file_map, false, &vm_map);
+                                     translator_->GetFileMap(), false, &vm_map);
     vm_map.ForEachRange([this, &contains](uint64_t vmaddr, uint64_t vmsize) {
       if (ContainsVerboseVMAddr(vmaddr, vmsize)) {
         contains = true;
@@ -1209,15 +1209,15 @@ void RangeSink::AddFileRange(const char* analyzer, string_view name,
   for (auto& pair : outputs_) {
     const std::string label = pair.second->Munge(name);
     if (translator_) {
-      bool ok = pair.first->file_map.AddRangeWithTranslation(
-          fileoff, filesize, label, translator_->file_map, verbose,
-          &pair.first->vm_map);
+      bool ok = pair.first->GetFileMap().AddRangeWithTranslation(
+          fileoff, filesize, label, translator_->GetFileMap(), verbose,
+          &pair.first->GetVMMap(arch_index_));
       if (!ok) {
         WARN("File range ($0, $1) for label $2 extends beyond base map",
              fileoff, filesize, name);
       }
     } else {
-      pair.first->file_map.AddRange(fileoff, filesize, label);
+      pair.first->GetFileMap().AddRange(fileoff, filesize, label);
     }
   }
 }
@@ -1235,10 +1235,10 @@ void RangeSink::AddFileRangeForVMAddr(const char* analyzer,
   assert(translator_);
   for (auto& pair : outputs_) {
     std::string label;
-    if (pair.first->vm_map.TryGetLabel(label_from_vmaddr, &label)) {
-      bool ok = pair.first->file_map.AddRangeWithTranslation(
-          file_offset, file_range.size(), label, translator_->file_map, verbose,
-          &pair.first->vm_map);
+    if (pair.first->GetVMMap(arch_index_).TryGetLabel(label_from_vmaddr, &label)) {
+      bool ok = pair.first->GetFileMap().AddRangeWithTranslation(
+          file_offset, file_range.size(), label, translator_->GetFileMap(), verbose,
+          &pair.first->GetVMMap(arch_index_));
       if (!ok) {
         WARN("File range ($0, $1) for label $2 extends beyond base map",
              file_offset, file_range.size(), label);
@@ -1264,11 +1264,11 @@ void RangeSink::AddFileRangeForFileRange(const char* analyzer,
   assert(translator_);
   for (auto& pair : outputs_) {
     std::string label;
-    if (pair.first->file_map.TryGetLabelForRange(
+    if (pair.first->GetFileMap().TryGetLabelForRange(
             from_file_offset, from_file_range.size(), &label)) {
-      bool ok = pair.first->file_map.AddRangeWithTranslation(
-          file_offset, file_range.size(), label, translator_->file_map, verbose,
-          &pair.first->vm_map);
+      bool ok = pair.first->GetFileMap().AddRangeWithTranslation(
+          file_offset, file_range.size(), label, translator_->GetFileMap(), verbose,
+          &pair.first->GetVMMap(arch_index_));
       if (!ok) {
         WARN("File range ($0, $1) for label $2 extends beyond base map",
              file_offset, file_range.size(), label);
@@ -1293,10 +1293,10 @@ void RangeSink::AddVMRangeForVMAddr(const char* analyzer,
   assert(translator_);
   for (auto& pair : outputs_) {
     std::string label;
-    if (pair.first->vm_map.TryGetLabel(label_from_vmaddr, &label)) {
-      bool ok = pair.first->vm_map.AddRangeWithTranslation(
-          addr, size, label, translator_->vm_map, verbose,
-          &pair.first->file_map);
+    if (pair.first->GetVMMap(arch_index_).TryGetLabel(label_from_vmaddr, &label)) {
+      bool ok = pair.first->GetVMMap(arch_index_).AddRangeWithTranslation(
+          addr, size, label, translator_->GetVMMap(arch_index_), verbose,
+          &pair.first->GetFileMap());
       if (!ok && verbose_level > 1) {
         WARN("VM range ($0, $1) for label $2 extends beyond base map", addr,
              size, label);
@@ -1318,9 +1318,9 @@ void RangeSink::AddVMRange(const char* analyzer, uint64_t vmaddr,
   assert(translator_);
   for (auto& pair : outputs_) {
     const std::string label = pair.second->Munge(name);
-    bool ok = pair.first->vm_map.AddRangeWithTranslation(
-        vmaddr, vmsize, label, translator_->vm_map, verbose,
-        &pair.first->file_map);
+    bool ok = pair.first->GetVMMap(arch_index_).AddRangeWithTranslation(
+        vmaddr, vmsize, label, translator_->GetVMMap(arch_index_), verbose,
+        &pair.first->GetFileMap());
     if (!ok) {
       WARN("VM range ($0, $1) for label $2 extends beyond base map", vmaddr,
            vmsize, name);
@@ -1361,8 +1361,8 @@ void RangeSink::AddRange(const char* analyzer, string_view name,
   }
 
   if (translator_) {
-    if (!translator_->vm_map.CoversRange(vmaddr, vmsize) ||
-        !translator_->file_map.CoversRange(fileoff, filesize)) {
+    if (!translator_->GetVMMap(arch_index_).CoversRange(vmaddr, vmsize) ||
+        !translator_->GetFileMap().CoversRange(fileoff, filesize)) {
       WARN("AddRange($0, $1, $2, $3, $4) will be ignored, because it is not "
            "covered by base map.",
            name.data(), vmaddr, vmsize, fileoff, filesize);
@@ -1374,11 +1374,11 @@ void RangeSink::AddRange(const char* analyzer, string_view name,
     const std::string label = pair.second->Munge(name);
     uint64_t common = std::min(vmsize, filesize);
 
-    pair.first->vm_map.AddDualRange(vmaddr, common, fileoff, label);
-    pair.first->file_map.AddDualRange(fileoff, common, vmaddr, label);
+    pair.first->GetVMMap(arch_index_).AddDualRange(vmaddr, common, fileoff, label);
+    pair.first->GetFileMap().AddDualRange(fileoff, common, vmaddr, label);
 
-    pair.first->vm_map.AddRange(vmaddr + common, vmsize - common, label);
-    pair.first->file_map.AddRange(fileoff + common, filesize - common, label);
+    pair.first->GetVMMap(arch_index_).AddRange(vmaddr + common, vmsize - common, label);
+    pair.first->GetFileMap().AddRange(fileoff + common, filesize - common, label);
   }
 }
 
@@ -1387,18 +1387,18 @@ uint64_t RangeSink::TranslateFileToVM(const char* ptr) {
   uint64_t offset = ptr - file_->data().data();
   uint64_t translated;
   if (!FileContainsPointer(ptr) ||
-      !translator_->file_map.Translate(offset, &translated)) {
+      !translator_->GetFileMap().Translate(offset, &translated)) {
     THROWF("Can't translate file offset ($0) to VM, contains: $1, map:\n$2",
            offset, FileContainsPointer(ptr),
-           translator_->file_map.DebugString().c_str());
+           translator_->GetFileMap().DebugString().c_str());
   }
   return translated;
 }
 
 std::string_view RangeSink::TranslateVMToFile(uint64_t address) {
   assert(translator_);
   uint64_t translated;
-  if (!translator_->vm_map.Translate(address, &translated) ||
+  if (!translator_->GetVMMap(arch_index_).Translate(address, &translated) ||
       translated > file_->data().size()) {
     THROWF("Can't translate VM pointer ($0) to file", address);
 
@@ -1864,14 +1864,33 @@ struct DualMaps {
   }
 
   void ComputeRollup(Rollup* rollup) {
+    // Collect all arch indices across all maps.
+    std::set<int> arch_indices;
     for (auto& map : maps_) {
-      map->vm_map.Compress();
-      map->file_map.Compress();
+      for (int id : map->VMArchIndices()) {
+        arch_indices.insert(id);
+      }
+    }
+    // Compress and roll up each arch's VM maps separately.
+    for (int arch_id : arch_indices) {
+      std::vector<const RangeMap*> vm_maps;
+      for (auto& map : maps_) {
+        if (map->HasVMMap(arch_id)) {
+          map->GetVMMap(arch_id).Compress();
+          vm_maps.push_back(&map->GetVMMap(arch_id));
+        }
+      }
+      if (!vm_maps.empty()) {
+        RangeMap::ComputeRollup(vm_maps,
+            [=](const std::vector<std::string>& keys, uint64_t addr, uint64_t end) {
+              return rollup->AddSizes(keys, end - addr, true);
+            });
+      }
+    }
+    // File maps are shared across all arches.
+    for (auto& map : maps_) {
+      map->GetFileMap().Compress();
     }
-    RangeMap::ComputeRollup(VmMaps(), [=](const std::vector<std::string>& keys,
-                                          uint64_t addr, uint64_t end) {
-      return rollup->AddSizes(keys, end - addr, true);
-    });
     RangeMap::ComputeRollup(
         FileMaps(),
         [=](const std::vector<std::string>& keys, uint64_t addr, uint64_t end) {
@@ -1895,7 +1914,38 @@ struct DualMaps {
   }
 
   void PrintFileMaps() { PrintMaps(FileMaps()); }
-  void PrintVMMaps() { PrintMaps(VmMaps()); }
+
+  void PrintVMMaps() {
+    std::set<int> arch_indices;
+    for (auto& map : maps_) {
+      for (int id : map->VMArchIndices()) {
+        arch_indices.insert(id);
+      }
+    }
+    for (int arch_id : arch_indices) {
+      if (arch_indices.size() > 1) {
+        auto it = arch_names_.find(arch_id);
+        if (it != arch_names_.end()) {
+          printf("%s:\n", it->second.c_str());
+        } else {
+          printf("arch %d:\n", arch_id);
+        }
+      }
+      std::vector<const RangeMap*> vm_maps;
+      for (auto& map : maps_) {
+        if (map->HasVMMap(arch_id)) {
+          vm_maps.push_back(&map->GetVMMap(arch_id));
+        }
+      }
+      if (!vm_maps.empty()) {
+        PrintMaps(vm_maps);
+      }
+    }
+  }
+
+  void SetArchNames(std::map<int, std::string> names) {
+    arch_names_ = std::move(names);
+  }
 
   std::string KeysToString(const std::vector<std::string>& keys) {
     std::string ret;
@@ -1921,23 +1971,16 @@ struct DualMaps {
   DualMap* base_map() { return maps_[0].get(); }
 
  private:
-  std::vector<const RangeMap*> VmMaps() const {
-    std::vector<const RangeMap*> ret;
-    for (const auto& map : maps_) {
-      ret.push_back(&map->vm_map);
-    }
-    return ret;
-  }
-
   std::vector<const RangeMap*> FileMaps() const {
     std::vector<const RangeMap*> ret;
     for (const auto& map : maps_) {
-      ret.push_back(&map->file_map);
+      ret.push_back(&map->GetFileMap());
     }
     return ret;
   }
 
   std::vector<std::unique_ptr<DualMap>> maps_;
+  std::map<int, std::string> arch_names_;
 };
 
 void Bloaty::ScanAndRollupFile(const std::string& filename, Rollup* rollup,
@@ -2002,14 +2045,19 @@ void Bloaty::ScanAndRollupFile(const std::string& filename, Rollup* rollup,
       rollup->file_total() + rollup->filtered_file_total();
   file->ProcessFile(sink_ptrs);
 
+  maps.SetArchNames(file->GetArchNames());
+
   // kInputFile source: Copy the base map to the filename sink(s).
   for (auto sink : filename_sink_ptrs) {
-    maps.base_map()->vm_map.ForEachRange(
-        [sink](uint64_t start, uint64_t length) {
-          sink->AddVMRange("inputfile_vmcopier", start, length,
-                           sink->input_file().filename());
-        });
-    maps.base_map()->file_map.ForEachRange(
+    for (int arch_id : maps.base_map()->VMArchIndices()) {
+      sink->set_arch_index(arch_id);
+      maps.base_map()->GetVMMap(arch_id).ForEachRange(
+          [sink](uint64_t start, uint64_t length) {
+            sink->AddVMRange("inputfile_vmcopier", start, length,
+                             sink->input_file().filename());
+          });
+    }
+    maps.base_map()->GetFileMap().ForEachRange(
         [sink](uint64_t start, uint64_t length) {
           sink->AddFileRange("inputfile_filecopier",
                              sink->input_file().filename(), start, length);

src/bloaty.h

@@ -20,12 +20,14 @@
 #define BLOATY_H_
 
 #include <stdlib.h>
+#include <cassert>
 #define __STDC_LIMIT_MACROS
 #define __STDC_FORMAT_MACROS
 #include <stdint.h>
 #include <inttypes.h>
 
 #include <list>
+#include <map>
 #include <memory>
 #include <set>
 #include <string>
@@ -127,6 +129,8 @@ class RangeSink {
   DataSource data_source() const { return data_source_; }
   const InputFile &input_file() const { return *file_; }
   bool IsBaseMap() const { return translator_ == nullptr; }
+  int arch_index() const { return arch_index_; }
+  void set_arch_index(int arch_index) { arch_index_ = arch_index; }
 
   // If vmsize or filesize is zero, this mapping is presumed not to exist in
   // that domain.  For example, .bss mappings don't exist in the file, and
@@ -239,6 +243,7 @@ class RangeSink {
   const DualMap* translator_;
   std::vector<std::pair<DualMap*, const NameMunger*>> outputs_;
   google::protobuf::Arena *arena_;
+  int arch_index_ = 0;
 };
 
 // NameMunger //////////////////////////////////////////////////////////////////
@@ -329,6 +334,11 @@ class ObjectFile {
   // given here, otherwise it is |this|.
   virtual void ProcessFile(const std::vector<RangeSink*>& sinks) const = 0;
 
+  // Optionally return human-readable names keyed by arch index, e.g.
+  // {0: "x86_64", 1: "arm64"}.  Only meaningful for formats with multiple VM
+  // address spaces, such as Mach-O universal binaries.
+  virtual std::map<int, std::string> GetArchNames() const { return {}; }
+
   virtual bool GetDisassemblyInfo(std::string_view symbol,
                                   DataSource symbol_source,
                                   DisassemblyInfo* info) const = 0;
@@ -376,11 +386,44 @@ std::string ItaniumDemangle(std::string_view symbol, DataSource source);
 
 // DualMap /////////////////////////////////////////////////////////////////////
 
-// Contains a RangeMap for VM space and file space for a given file.
+// Contains RangeMaps for VM space and file space for a given file.
+// For universal binaries, each architecture slice has its own VM address space
+// (keyed by arch_index). File offsets are shared across all arches and use a
+// single RangeMap. For single-arch binaries, only arch_index 0 is used.
 
 struct DualMap {
-  RangeMap vm_map;
+  RangeMap& GetVMMap(int arch_index) { return vm_maps_[arch_index]; }
+  const RangeMap& GetVMMap(int arch_index) const {
+    auto it = vm_maps_.find(arch_index);
+    if (it == vm_maps_.end()) {
+      static const RangeMap empty;
+      return empty;
+    }
+    return it->second;
+  }
+
+  RangeMap& GetFileMap() { return file_map; }
+  const RangeMap& GetFileMap() const { return file_map; }
+
+  bool HasVMMap(int arch_index) const {
+    return vm_maps_.find(arch_index) != vm_maps_.end();
+  }
+
+  std::vector<int> VMArchIndices() const {
+    std::vector<int> ids;
+    for (const auto& kv : vm_maps_) {
+      ids.push_back(kv.first);
+    }
+    return ids;
+  }
+
+  RangeMap& vm_map() { return GetVMMap(0); }
+  const RangeMap& vm_map() const { return GetVMMap(0); }
+
   RangeMap file_map;
+
+ private:
+  std::map<int, RangeMap> vm_maps_;
 };
 
 struct DisassemblyInfo {

src/disassemble.cc

@@ -212,7 +212,7 @@ std::string DisassembleFunction(const DisassemblyInfo& info) {
         } else {
           op_str = "<" + std::to_string(iter->second);
         }
-      } else if (info.symbol_map.vm_map.TryGetLabel(target, &label)) {
+      } else if (info.symbol_map.vm_map().TryGetLabel(target, &label)) {
         op_str = label;
       }
     }

src/dwarf.cc

@@ -440,7 +440,7 @@ void AddDIE(const dwarf::CU& cu, const GeneralDIE& die,
       // Unfortunately the location doesn't include a size, so we look that part
       // up in the symbol map.
       uint64_t size;
-      if (symbol_map.vm_map.TryGetSize(addr, &size)) {
+      if (symbol_map.GetVMMap(sink->arch_index()).TryGetSize(addr, &size)) {
         sink->AddVMRangeIgnoreDuplicate("dwarf_location", addr, size,
                                         cu.unit_name());
       } else {