2 files changed, 66 insertions, 36 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
index 215742c0e52..0c19f8fee0d 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
@@ -100,27 +100,8 @@ Error MachOAtomGraphBuilder::parseSections() {
     if (auto EC = SecRef.getName(Name))
       return errorCodeToError(EC);
 
-    StringRef Content;
-
-    // If this is a virtual section, leave its content empty.
-    if (!SecRef.isVirtual()) {
-      if (auto EC = SecRef.getContents(Content))
-        return errorCodeToError(EC);
-      if (Content.size() != SecRef.getSize())
-        return make_error<JITLinkError>("Section content size does not match "
-                                        "declared size for " +
-                                        Name);
-    }
-
     unsigned SectionIndex = SecRef.getIndex() + 1;
 
-    LLVM_DEBUG({
-      dbgs() << "Adding section " << Name << ": "
-             << format("0x%016" PRIx64, SecRef.getAddress())
-             << ", size: " << Content.size()
-             << ", align: " << SecRef.getAlignment() << "\n";
-    });
-
     // FIXME: Get real section permissions
     // How, exactly, on MachO?
     sys::Memory::ProtectionFlags Prot;
@@ -132,12 +113,41 @@ Error MachOAtomGraphBuilder::parseSections() {
                                                        sys::Memory::MF_WRITE);
 
     auto &GenericSection = G->createSection(Name, Prot, SecRef.isBSS());
-    if (SecRef.isVirtual())
-      Sections[SectionIndex] =
-          MachOSection(GenericSection, SecRef.getAddress(),
-                       SecRef.getAlignment(), SecRef.getSize());
-    Sections[SectionIndex] = MachOSection(GenericSection, SecRef.getAddress(),
-                                          SecRef.getAlignment(), Content);
+
+    LLVM_DEBUG({
+      dbgs() << "Adding section " << Name << ": "
+             << format("0x%016" PRIx64, SecRef.getAddress())
+             << ", align: " << SecRef.getAlignment() << "\n";
+    });
+
+    assert(!Sections.count(SectionIndex) && "Section index already in use");
+
+    auto &MachOSec =
+        Sections
+            .try_emplace(SectionIndex, GenericSection, SecRef.getAddress(),
+                         SecRef.getAlignment())
+            .first->second;
+
+    if (!SecRef.isVirtual()) {
+      // If this section has content then record it.
+      StringRef Content;
+      if (auto EC = SecRef.getContents(Content))
+        return errorCodeToError(EC);
+      if (Content.size() != SecRef.getSize())
+        return make_error<JITLinkError>("Section content size does not match "
+                                        "declared size for " +
+                                        Name);
+      MachOSec.setContent(Content);
+    } else {
+      // If this is a zero-fill section then just record the size.
+      MachOSec.setZeroFill(SecRef.getSize());
+    }
+
+    uint32_t SectionFlags =
+        Obj.is64Bit() ? Obj.getSection64(SecRef.getRawDataRefImpl()).flags
+                      : Obj.getSection(SecRef.getRawDataRefImpl()).flags;
+
+    MachOSec.setNoDeadStrip(SectionFlags & MachO::S_ATTR_NO_DEAD_STRIP);
   }
 
   return Error::success();
@@ -290,7 +300,7 @@ Error MachOAtomGraphBuilder::addNonCustomAtoms() {
 
   LLVM_DEBUG(dbgs() << "MachOGraphBuilder setting atom content\n");
 
-  // Set atom contents.
+  // Set atom contents and any section-based flags.
   for (auto &KV : SecToAtoms) {
     auto &S = *KV.first;
     auto &SecAtoms = KV.second;
@@ -304,10 +314,16 @@ Error MachOAtomGraphBuilder::addNonCustomAtoms() {
         dbgs() << "  " << A << " to [ " << S.getAddress() + Offset << " .. "
                << S.getAddress() + LastAtomAddr << " ]\n";
       });
+
       if (S.isZeroFill())
         A.setZeroFill(LastAtomAddr - Offset);
       else
         A.setContent(S.getContent().substr(Offset, LastAtomAddr - Offset));
+
+      // If the section has no-dead-strip set then mark the atom as live.
+      if (S.isNoDeadStrip())
+        A.setLive(true);
+
       LastAtomAddr = Offset;
     }
   }
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h
index 540e2c366f0..f5cd0ea1853 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h
+++ b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h
@@ -34,17 +34,10 @@ protected:
   public:
     MachOSection() = default;
 
-    /// Create a MachO section with the given content.
+    /// Create a MachO section with the given address and alignment.
     MachOSection(Section &GenericSection, JITTargetAddress Address,
-                 unsigned Alignment, StringRef Content)
+                 unsigned Alignment)
         : Address(Address), GenericSection(&GenericSection),
-          ContentPtr(Content.data()), Size(Content.size()),
-          Alignment(Alignment) {}
-
-    /// Create a zero-fill MachO section with the given size.
-    MachOSection(Section &GenericSection, JITTargetAddress Address,
-                 unsigned Alignment, size_t ZeroFillSize)
-        : Address(Address), GenericSection(&GenericSection), Size(ZeroFillSize),
           Alignment(Alignment) {}
 
     /// Create a section without address, content or size (used for common
@@ -61,6 +54,19 @@ protected:
       return GenericSection->getName();
     }
 
+    MachOSection &setContent(StringRef Content) {
+      assert(!ContentPtr && !Size && "Content/zeroFill already set");
+      ContentPtr = Content.data();
+      Size = Content.size();
+      return *this;
+    }
+
+    MachOSection &setZeroFill(uint64_t Size) {
+      assert(!ContentPtr && !Size && "Content/zeroFill already set");
+      this->Size = Size;
+      return *this;
+    }
+
     bool isZeroFill() const { return !ContentPtr; }
 
     bool empty() const { return getSize() == 0; }
@@ -76,12 +82,20 @@ protected:
 
     unsigned getAlignment() const { return Alignment; }
 
+    MachOSection &setNoDeadStrip(bool NoDeadStrip) {
+      this->NoDeadStrip = NoDeadStrip;
+      return *this;
+    }
+
+    bool isNoDeadStrip() const { return NoDeadStrip; }
+
   private:
     JITTargetAddress Address = 0;
     Section *GenericSection = nullptr;
     const char *ContentPtr = nullptr;
-    size_t Size = 0;
+    uint64_t Size = 0;
     unsigned Alignment = 0;
+    bool NoDeadStrip = false;
   };
 
   using CustomAtomizeFunction = std::function<Error(MachOSection &S)>;