From 1c27fe0578e52c52cdf2e53b63fef632cdce26ef Mon Sep 17 00:00:00 2001
From: Alastair Houghton <ahoughton@apple.com>
Date: Thu, 8 Aug 2024 18:53:17 +0100
Subject: [PATCH 1/2] [RuntimeDyld][Windows] Allocate space for dllimport
 things.

We weren't taking account of the space we require in the stubs for
things that are `dllimport`ed, and as a result we could hit the
assertion failure for running out of stub space.  Fix that.

rdar://133473673
---
 llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp  |  5 ++++-
 .../ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp   | 11 +++++++++++
 .../lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h |  7 +++++++
 .../lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h | 10 ++++++++++
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
index a9aaff42433f6..fc6353900e891 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
@@ -690,9 +690,12 @@ unsigned RuntimeDyldImpl::computeSectionStubBufSize(const ObjectFile &Obj,
     if (!(RelSecI == Section))
       continue;
 
-    for (const RelocationRef &Reloc : SI->relocations())
+    for (const RelocationRef &Reloc : SI->relocations()) {
       if (relocationNeedsStub(Reloc))
         StubBufSize += StubSize;
+      if (relocationNeedsDLLImportStub(Reloc))
+        StubBufSize = sizeAfterAddingDLLImportStub(StubBufSize);
+    }
   }
 
   // Get section data size and alignment
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
index 9255311f992d0..3ec01a307d1c5 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
@@ -118,4 +118,15 @@ bool RuntimeDyldCOFF::isCompatibleFile(const object::ObjectFile &Obj) const {
   return Obj.isCOFF();
 }
 
+bool RuntimeDyldCOFF::relocationNeedsDLLImportStub(const RelocationRef &R) const {
+  object::symbol_iterator Symbol = R.getSymbol();
+  Expected<StringRef> TargetNameOrErr = Symbol->getName();
+  if (!TargetNameOrErr)
+    return false;
+
+  StringRef TargetName = *TargetNameOrErr;
+
+  return TargetName.startswith(getImportSymbolPrefix());
+}
+
 } // namespace llvm
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h
index 41ee06c15448f..f58faebf6a004 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h
@@ -14,6 +14,7 @@
 #define LLVM_RUNTIME_DYLD_COFF_H
 
 #include "RuntimeDyldImpl.h"
+#include "llvm/Support/MathExtras.h"
 
 #define DEBUG_TYPE "dyld"
 
@@ -49,6 +50,12 @@ class RuntimeDyldCOFF : public RuntimeDyldImpl {
 
   static constexpr StringRef getImportSymbolPrefix() { return "__imp_"; }
 
+  bool relocationNeedsDLLImportStub(const RelocationRef &R) const;
+
+  unsigned sizeAfterAddingDLLImportStub(unsigned Size) const {
+    return alignTo(Size, PointerSize) + PointerSize;
+  }
+
 private:
   unsigned PointerSize;
   uint32_t PointerReloc;
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
index 501417db421a6..a7fdbc5e5b038 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
@@ -449,6 +449,16 @@ class RuntimeDyldImpl {
     return true;    // Conservative answer
   }
 
+  // Return true if the relocation R may require allocating a DLL import stub.
+  virtual bool relocationNeedsDLLImportStub(const RelocationRef &R) const {
+    return false;
+  }
+
+  // Add the size of a DLL import stub to the buffer size
+  virtual unsigned sizeAfterAddingDLLImportStub(unsigned Size) const {
+    return Size;
+  }
+
 public:
   RuntimeDyldImpl(RuntimeDyld::MemoryManager &MemMgr,
                   JITSymbolResolver &Resolver)

From 43d75207192e7017c41f20d7fb94d937516039f0 Mon Sep 17 00:00:00 2001
From: Alastair Houghton <alastair@alastairs-place.net>
Date: Thu, 8 Aug 2024 19:02:10 +0100
Subject: [PATCH 2/2] [RuntimeDyld][Windows] Eliminate a local variable.

No need for `TargetName` to exist separately, really.

Co-authored-by: Saleem Abdulrasool <compnerd@compnerd.org>
---
 llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
index 3ec01a307d1c5..e6455ce4a70e0 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
@@ -118,15 +118,14 @@ bool RuntimeDyldCOFF::isCompatibleFile(const object::ObjectFile &Obj) const {
   return Obj.isCOFF();
 }
 
-bool RuntimeDyldCOFF::relocationNeedsDLLImportStub(const RelocationRef &R) const {
+bool RuntimeDyldCOFF::relocationNeedsDLLImportStub(
+    const RelocationRef &R) const {
   object::symbol_iterator Symbol = R.getSymbol();
   Expected<StringRef> TargetNameOrErr = Symbol->getName();
   if (!TargetNameOrErr)
     return false;
 
-  StringRef TargetName = *TargetNameOrErr;
-
-  return TargetName.startswith(getImportSymbolPrefix());
+  return TargetNameOrErr->startswith(getImportSymbolPrefix());
 }
 
 } // namespace llvm