From 36baa33fc1b7da84428a64b502eb762888df2095 Mon Sep 17 00:00:00 2001
From: Nick Kledzik <kledzik@apple.com>
Date: Wed, 8 Jan 2014 02:52:58 +0000
Subject: [PATCH] [mach-o] properly extract atom content from subrange of
 section content

llvm-svn: 198728
---
 .../MachO/MachONormalizedFileToAtoms.cpp      | 31 ++++++++++--
 .../MachONormalizedFileToAtomsTests.cpp       | 50 +++++++++++--------
 2 files changed, 57 insertions(+), 24 deletions(-)

diff --git a/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp b/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
index 2032e0cea057..d3ff1e406483 100644
--- a/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
+++ b/lld/lib/ReaderWriter/MachO/MachONormalizedFileToAtoms.cpp
@@ -34,13 +34,38 @@ namespace lld {
 namespace mach_o {
 namespace normalized {
 
+static uint64_t nextSymbolAddress(const NormalizedFile &normalizedFile, 
+                                const Symbol &symbol) {
+  uint64_t symbolAddr = symbol.value;
+  uint8_t symbolSectionIndex = symbol.sect;
+  const Section &section = normalizedFile.sections[symbolSectionIndex - 1];
+  // If no symbol after this address, use end of section address.
+  uint64_t closestAddr = section.address + section.content.size();
+  for (const Symbol &s : normalizedFile.globalSymbols) {
+    if (s.sect != symbolSectionIndex)
+      continue;
+    uint64_t sValue = s.value;
+    if (sValue <= symbolAddr)
+      continue;
+    if (sValue < closestAddr)
+      closestAddr = s.value;
+  }
+  return closestAddr;
+}
+
 static ErrorOr<std::unique_ptr<lld::File>>
 normalizedObjectToAtoms(const NormalizedFile &normalizedFile, StringRef path) {
   std::unique_ptr<MachOFile> file(new MachOFile(path));
 
-  for (auto &sym : normalizedFile.globalSymbols) {
-    file->addDefinedAtom(sym.name,
-                         normalizedFile.sections[sym.sect - 1].content);
+ for (const Symbol &sym : normalizedFile.globalSymbols) {
+    // Mach-O symbol table does have size in it, so need to scan ahead
+    // to find symbol with next highest address.
+    const Section &section = normalizedFile.sections[sym.sect - 1];
+    uint64_t offset = sym.value - section.address;
+    uint64_t size = nextSymbolAddress(normalizedFile, sym) - sym.value;
+    ArrayRef<uint8_t> atomContent = llvm::makeArrayRef(&section.content[offset],
+                                                      size);
+    file->addDefinedAtom(sym.name, atomContent);
   }
 
   assert(normalizedFile.localSymbols.empty() &&
diff --git a/lld/unittests/MachOTests/MachONormalizedFileToAtomsTests.cpp b/lld/unittests/MachOTests/MachONormalizedFileToAtomsTests.cpp
index 9b2859a5961c..ddd5591ed08a 100644
--- a/lld/unittests/MachOTests/MachONormalizedFileToAtomsTests.cpp
+++ b/lld/unittests/MachOTests/MachONormalizedFileToAtomsTests.cpp
@@ -20,42 +20,50 @@ using llvm::ErrorOr;
 using namespace lld::mach_o::normalized;
 using namespace llvm::MachO;
 
-unsigned countDefinedAtoms(const lld::File &file) {
-  unsigned count = 0;
-  for (const auto &d : file.defined()) {
-    (void)d;
-    ++count;
-  }
-  return count;
-}
-
 TEST(ToAtomsTest, empty_obj_x86_64) {
   NormalizedFile f;
   f.arch = lld::MachOLinkingContext::arch_x86_64;
   ErrorOr<std::unique_ptr<const lld::File>> atom_f = normalizedToAtoms(f, "");
   EXPECT_FALSE(!atom_f);
-  EXPECT_EQ(0U, countDefinedAtoms(**atom_f));
+  EXPECT_EQ(0U, (*atom_f)->defined().size());
 }
 
 TEST(ToAtomsTest, basic_obj_x86_64) {
   NormalizedFile f;
   f.arch = lld::MachOLinkingContext::arch_x86_64;
   Section textSection;
-  static const uint8_t contentBytes[] = { 0x55, 0x48, 0x89, 0xE5,
-                                          0x31, 0xC0, 0x5D, 0xC3 };
+  static const uint8_t contentBytes[] = { 0x90, 0xC3, 0xC3 };
   const unsigned contentSize = sizeof(contentBytes) / sizeof(contentBytes[0]);
   textSection.content.insert(textSection.content.begin(), contentBytes,
                              &contentBytes[contentSize]);
   f.sections.push_back(textSection);
-  Symbol mainSymbol;
-  mainSymbol.name = "_main";
-  mainSymbol.type = N_SECT;
-  mainSymbol.sect = 1;
-  f.globalSymbols.push_back(mainSymbol);
+  Symbol fooSymbol;
+  fooSymbol.name = "_foo";
+  fooSymbol.type = N_SECT;
+  fooSymbol.sect = 1;
+  fooSymbol.value = 0;
+  f.globalSymbols.push_back(fooSymbol);
+  Symbol barSymbol;
+  barSymbol.name = "_bar";
+  barSymbol.type = N_SECT;
+  barSymbol.sect = 1;
+  barSymbol.value = 2;
+  f.globalSymbols.push_back(barSymbol);
+  
   ErrorOr<std::unique_ptr<const lld::File>> atom_f = normalizedToAtoms(f, "");
   EXPECT_FALSE(!atom_f);
-  EXPECT_EQ(1U, countDefinedAtoms(**atom_f));
-  const lld::DefinedAtom *singleAtom = *(*atom_f)->defined().begin();
-  llvm::ArrayRef<uint8_t> atomContent(singleAtom->rawContent());
-  EXPECT_EQ(0, memcmp(atomContent.data(), contentBytes, contentSize));
+  const lld::File &file = **atom_f;
+  EXPECT_EQ(2U, file.defined().size());
+  lld::File::defined_iterator it = file.defined().begin();
+  const lld::DefinedAtom *atom1 = *it;
+  ++it;
+  const lld::DefinedAtom *atom2 = *it;
+  EXPECT_TRUE(atom1->name().equals("_foo"));
+  EXPECT_EQ(2U, atom1->rawContent().size());
+  EXPECT_EQ(0x90, atom1->rawContent()[0]);
+  EXPECT_EQ(0xC3, atom1->rawContent()[1]);
+  
+  EXPECT_TRUE(atom2->name().equals("_bar"));
+  EXPECT_EQ(1U, atom2->rawContent().size());
+  EXPECT_EQ(0xC3, atom2->rawContent()[0]);
 }