[JITLink] Enable creation and management of mutable block content.

This patch introduces new operations on jitlink::Blocks: setMutableContent,
getMutableContent and getAlreadyMutableContent. The setMutableContent method
will set the block content data and size members and flag the content as
mutable. The getMutableContent method will return a mutable copy of the existing
content value, auto-allocating and populating a new mutable copy if the existing
content is marked immutable. The getAlreadyMutableMethod asserts that the
existing content is already mutable and returns it.

setMutableContent should be used when updating the block with totally new
content backed by mutable memory. It can be used to change the size of the
block. The argument value should *not* be shared with any other block.

getMutableContent should be used when clients want to modify the existing
content and are unsure whether it is mutable yet.

getAlreadyMutableContent should be used when clients want to modify the existing
content and know from context that it must already be immutable.

These operations reduce copy-modify-update boilerplate and unnecessary copies
introduced when clients couldn't me sure whether the existing content was
mutable or not.

include/llvm/ExecutionEngine/JITLink/JITLink.h

@@ -35,6 +35,7 @@
 namespace llvm {
 namespace jitlink {
 
+class LinkGraph;
 class Symbol;
 class Section;
 
@@ -138,8 +139,9 @@ private:
 
 protected:
   // bitfields for Block, allocated here to improve packing.
+  uint64_t ContentMutable : 1;
   uint64_t P2Align : 5;
-  uint64_t AlignmentOffset : 57;
+  uint64_t AlignmentOffset : 56;
 };
 
 using SectionOrdinal = unsigned;
@@ -158,11 +160,14 @@ private:
            "Alignment offset cannot exceed alignment");
     assert(AlignmentOffset <= MaxAlignmentOffset &&
            "Alignment offset exceeds maximum");
+    ContentMutable = false;
     P2Align = Alignment ? countTrailingZeros(Alignment) : 0;
     this->AlignmentOffset = AlignmentOffset;
   }
 
   /// Create a defined addressable for the given content.
+  /// The Content is assumed to be non-writable, and will be copied when
+  /// mutations are required.
   Block(Section &Parent, ArrayRef<char> Content, JITTargetAddress Address,
         uint64_t Alignment, uint64_t AlignmentOffset)
       : Addressable(Address, true), Parent(Parent), Data(Content.data()),
@@ -172,6 +177,26 @@ private:
            "Alignment offset cannot exceed alignment");
     assert(AlignmentOffset <= MaxAlignmentOffset &&
            "Alignment offset exceeds maximum");
+    ContentMutable = false;
+    P2Align = Alignment ? countTrailingZeros(Alignment) : 0;
+    this->AlignmentOffset = AlignmentOffset;
+  }
+
+  /// Create a defined addressable for the given content.
+  /// The content is assumed to be writable, and the caller is responsible
+  /// for ensuring that it lives for the duration of the Block's lifetime.
+  /// The standard way to achieve this is to allocate it on the Graph's
+  /// allocator.
+  Block(Section &Parent, MutableArrayRef<char> Content,
+        JITTargetAddress Address, uint64_t Alignment, uint64_t AlignmentOffset)
+      : Addressable(Address, true), Parent(Parent), Data(Content.data()),
+        Size(Content.size()) {
+    assert(isPowerOf2_64(Alignment) && "Alignment must be power of 2");
+    assert(AlignmentOffset < Alignment &&
+           "Alignment offset cannot exceed alignment");
+    assert(AlignmentOffset <= MaxAlignmentOffset &&
+           "Alignment offset exceeds maximum");
+    ContentMutable = true;
     P2Align = Alignment ? countTrailingZeros(Alignment) : 0;
     this->AlignmentOffset = AlignmentOffset;
   }
@@ -210,8 +235,44 @@ public:
   void setContent(ArrayRef<char> Content) {
     Data = Content.data();
     Size = Content.size();
+    ContentMutable = false;
   }
 
+  /// Get mutable content for this block.
+  ///
+  /// If this Block's content is not already mutable this will trigger a copy
+  /// of the existing immutable content to a new, mutable buffer allocated using
+  /// LinkGraph::allocateContent.
+  MutableArrayRef<char> getMutableContent(LinkGraph &G);
+
+  /// Get mutable content for this block.
+  ///
+  /// This block's content must already be mutable. It is a programmatic error
+  /// to call this on a block with immutable content -- consider using
+  /// getMutableContent instead.
+  MutableArrayRef<char> getAlreadyMutableContent() {
+    assert(ContentMutable && "Content is not mutable");
+    return MutableArrayRef<char>(const_cast<char *>(Data), Size);
+  }
+
+  /// Set mutable content for this block.
+  ///
+  /// The caller is responsible for ensuring that the memory pointed to by
+  /// MutableContent is not deallocated while pointed to by this block.
+  void setMutableContent(MutableArrayRef<char> MutableContent) {
+    Data = MutableContent.data();
+    Size = MutableContent.size();
+    ContentMutable = true;
+  }
+
+  /// Returns true if this block's content is mutable.
+  ///
+  /// This is primarily useful for asserting that a block is already in a
+  /// mutable state prior to modifying the content. E.g. when applying
+  /// fixups we expect the block to already be mutable as it should have been
+  /// copied to working memory.
+  bool isContentMutable() const { return ContentMutable; }
+
   /// Get the alignment for this content.
   uint64_t getAlignment() const { return 1ull << P2Align; }
 
@@ -268,7 +329,7 @@ public:
   }
 
 private:
-  static constexpr uint64_t MaxAlignmentOffset = (1ULL << 57) - 1;
+  static constexpr uint64_t MaxAlignmentOffset = (1ULL << 56) - 1;
 
   Section &Parent;
   const char *Data = nullptr;
@@ -858,10 +919,10 @@ public:
   /// Allocate a copy of the given string using the LinkGraph's allocator.
   /// This can be useful when renaming symbols or adding new content to the
   /// graph.
-  ArrayRef<char> allocateString(ArrayRef<char> Source) {
+  MutableArrayRef<char> allocateContent(ArrayRef<char> Source) {
     auto *AllocatedBuffer = Allocator.Allocate<char>(Source.size());
     llvm::copy(Source, AllocatedBuffer);
-    return ArrayRef<char>(AllocatedBuffer, Source.size());
+    return MutableArrayRef<char>(AllocatedBuffer, Source.size());
   }
 
   /// Allocate a copy of the given string using the LinkGraph's allocator.
@@ -871,12 +932,12 @@ public:
   /// Note: This Twine-based overload requires an extra string copy and an
   /// extra heap allocation for large strings. The ArrayRef<char> overload
   /// should be preferred where possible.
-  ArrayRef<char> allocateString(Twine Source) {
+  MutableArrayRef<char> allocateString(Twine Source) {
     SmallString<256> TmpBuffer;
     auto SourceStr = Source.toStringRef(TmpBuffer);
     auto *AllocatedBuffer = Allocator.Allocate<char>(SourceStr.size());
     llvm::copy(SourceStr, AllocatedBuffer);
-    return ArrayRef<char>(AllocatedBuffer, SourceStr.size());
+    return MutableArrayRef<char>(AllocatedBuffer, SourceStr.size());
   }
 
   /// Create a section with the given name, protection flags, and alignment.
@@ -898,6 +959,15 @@ public:
     return createBlock(Parent, Content, Address, Alignment, AlignmentOffset);
   }
 
+  /// Create a content block with initially mutable data.
+  Block &createMutableContentBlock(Section &Parent,
+                                   MutableArrayRef<char> MutableContent,
+                                   uint64_t Address, uint64_t Alignment,
+                                   uint64_t AlignmentOffset) {
+    return createBlock(Parent, MutableContent, Address, Alignment,
+                       AlignmentOffset);
+  }
+
   /// Create a zero-fill block.
   Block &createZeroFillBlock(Section &Parent, uint64_t Size, uint64_t Address,
                              uint64_t Alignment, uint64_t AlignmentOffset) {
@@ -1213,6 +1283,12 @@ private:
   ExternalSymbolSet AbsoluteSymbols;
 };
 
+inline MutableArrayRef<char> Block::getMutableContent(LinkGraph &G) {
+  if (!ContentMutable)
+    setMutableContent(G.allocateContent({Data, Size}));
+  return MutableArrayRef<char>(const_cast<char *>(Data), Size);
+}
+
 /// Enables easy lookup of blocks by addresses.
 class BlockAddressMap {
 public:

include/llvm/ExecutionEngine/JITLink/x86_64.h

@@ -258,10 +258,10 @@ inline bool isInRangeForImmS32(int64_t Value) {
 }
 
 /// Apply fixup expression for edge to block content.
-inline Error applyFixup(LinkGraph &G, Block &B, const Edge &E,
-                        char *BlockWorkingMem) {
+inline Error applyFixup(LinkGraph &G, Block &B, const Edge &E) {
   using namespace support;
 
+  char *BlockWorkingMem = B.getAlreadyMutableContent().data();
   char *FixupPtr = BlockWorkingMem + E.getOffset();
   JITTargetAddress FixupAddress = B.getAddress() + E.getOffset();
 

lib/ExecutionEngine/JITLink/ELF_x86_64.cpp

@@ -780,10 +780,11 @@ private:
     return Error::success();
   }
 
-  Error applyFixup(LinkGraph &G, Block &B, const Edge &E,
-                   char *BlockWorkingMem) const {
+  Error applyFixup(LinkGraph &G, Block &B, const Edge &E) const {
     using namespace ELF_x86_64_Edges;
     using namespace llvm::support;
+
+    char *BlockWorkingMem = B.getAlreadyMutableContent().data();
     char *FixupPtr = BlockWorkingMem + E.getOffset();
     JITTargetAddress FixupAddress = B.getAddress() + E.getOffset();
     switch (E.getKind()) {

lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp

@@ -404,7 +404,7 @@ void JITLinkerBase::copyBlockContentToWorkingMemory(
       memcpy(BlockDataPtr, B->getContent().data(), B->getContent().size());
 
       // Point the block's content to the fixed up buffer.
-      B->setContent({BlockDataPtr, B->getContent().size()});
+      B->setMutableContent({BlockDataPtr, B->getContent().size()});
 
       // Update block end pointer.
       LastBlockEnd = BlockDataPtr + B->getContent().size();

lib/ExecutionEngine/JITLink/JITLinkGeneric.h

@@ -159,8 +159,7 @@ private:
           continue;
 
         // Dispatch to LinkerImpl for fixup.
-        auto *BlockData = const_cast<char *>(B->getContent().data());
-        if (auto Err = impl().applyFixup(G, *B, E, BlockData))
+        if (auto Err = impl().applyFixup(G, *B, E))
           return Err;
       }
     }

lib/ExecutionEngine/JITLink/MachO_arm64.cpp

@@ -526,10 +526,10 @@ private:
     return 0;
   }
 
-  Error applyFixup(LinkGraph &G, Block &B, const Edge &E,
-                   char *BlockWorkingMem) const {
+  Error applyFixup(LinkGraph &G, Block &B, const Edge &E) const {
     using namespace support;
 
+    char *BlockWorkingMem = B.getAlreadyMutableContent().data();
     char *FixupPtr = BlockWorkingMem + E.getOffset();
     JITTargetAddress FixupAddress = B.getAddress() + E.getOffset();
 

lib/ExecutionEngine/JITLink/MachO_x86_64.cpp

@@ -522,8 +522,7 @@ static Error optimizeMachO_x86_64_GOTAndStubs(LinkGraph &G) {
           E.setTarget(GOTTarget);
           E.setKind(x86_64::Delta32);
           E.setAddend(E.getAddend() - 4);
-          auto *BlockData = reinterpret_cast<uint8_t *>(
-              const_cast<char *>(B->getContent().data()));
+          auto *BlockData = B->getMutableContent(G).data();
           BlockData[E.getOffset() - 2] = 0x8d;
           LLVM_DEBUG({
             dbgs() << "  Replaced GOT load wih LEA:\n    ";
@@ -577,9 +576,8 @@ public:
       : JITLinker(std::move(Ctx), std::move(G), std::move(PassConfig)) {}
 
 private:
-  Error applyFixup(LinkGraph &G, Block &B, const Edge &E,
-                   char *BlockWorkingMem) const {
-    return x86_64::applyFixup(G, B, E, BlockWorkingMem);
+  Error applyFixup(LinkGraph &G, Block &B, const Edge &E) const {
+    return x86_64::applyFixup(G, B, E);
   }
 };
 

unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp

@@ -137,6 +137,71 @@ TEST(LinkGraphTest, BlockAndSymbolIteration) {
   EXPECT_TRUE(llvm::count(G.defined_symbols(), &S4));
 }
 
+TEST(LinkGraphTest, ContentAccessAndUpdate) {
+  // Check that we can make a defined symbol external.
+  LinkGraph G("foo", Triple("x86_64-apple-darwin"), 8, support::little,
+              getGenericEdgeKindName);
+  auto &Sec = G.createSection("__data", RWFlags);
+
+  // Create an initial block.
+  auto &B = G.createContentBlock(Sec, BlockContent, 0x1000, 8, 0);
+
+  EXPECT_FALSE(B.isContentMutable()) << "Content unexpectedly mutable";
+  EXPECT_EQ(B.getContent().data(), BlockContent.data())
+      << "Unexpected block content data pointer";
+  EXPECT_EQ(B.getContent().size(), BlockContent.size())
+      << "Unexpected block content size";
+
+  // Expect that attempting to get already-mutable content fails if the
+  // content is not yet mutable.
+  EXPECT_DEATH({ (void)B.getAlreadyMutableContent(); },
+               "Content is not mutable")
+      << "Unexpected mutable access allowed to immutable data";
+
+  // Check that mutable content is copied on request as expected.
+  auto MutableContent = B.getMutableContent(G);
+  EXPECT_TRUE(B.isContentMutable()) << "Content unexpectedly immutable";
+  EXPECT_NE(MutableContent.data(), BlockContent.data())
+      << "Unexpected mutable content data pointer";
+  EXPECT_EQ(MutableContent.size(), BlockContent.size())
+      << "Unexpected mutable content size";
+  EXPECT_TRUE(std::equal(MutableContent.begin(), MutableContent.end(),
+                         BlockContent.begin()))
+      << "Unexpected mutable content value";
+
+  // Check that already-mutable content behaves as expected, with no
+  // further copies.
+  auto MutableContent2 = B.getMutableContent(G);
+  EXPECT_TRUE(B.isContentMutable()) << "Content unexpectedly immutable";
+  EXPECT_EQ(MutableContent2.data(), MutableContent.data())
+      << "Unexpected mutable content 2 data pointer";
+  EXPECT_EQ(MutableContent2.size(), MutableContent.size())
+      << "Unexpected mutable content 2 size";
+
+  // Check that getAlreadyMutableContent behaves as expected, with no
+  // further copies.
+  auto MutableContent3 = B.getMutableContent(G);
+  EXPECT_TRUE(B.isContentMutable()) << "Content unexpectedly immutable";
+  EXPECT_EQ(MutableContent3.data(), MutableContent.data())
+      << "Unexpected mutable content 2 data pointer";
+  EXPECT_EQ(MutableContent3.size(), MutableContent.size())
+      << "Unexpected mutable content 2 size";
+
+  // Set content back to immutable and check that everything behaves as
+  // expected again.
+  B.setContent(BlockContent);
+  EXPECT_FALSE(B.isContentMutable()) << "Content unexpectedly mutable";
+  EXPECT_EQ(B.getContent().data(), BlockContent.data())
+      << "Unexpected block content data pointer";
+  EXPECT_EQ(B.getContent().size(), BlockContent.size())
+      << "Unexpected block content size";
+
+  // Create an initially mutable block.
+  auto &B2 = G.createMutableContentBlock(Sec, MutableContent, 0x10000, 8, 0);
+
+  EXPECT_TRUE(B2.isContentMutable()) << "Expected B2 content to be mutable";
+}
+
 TEST(LinkGraphTest, MakeExternal) {
   // Check that we can make a defined symbol external.
   LinkGraph G("foo", Triple("x86_64-apple-darwin"), 8, support::little,