[llvm] [RuntimeDyld][Windows] Allocate space for dllimport things. (PR #102586)

[Alastair Houghton via llvm-commits]

https://github.com/al45tair created https://github.com/llvm/llvm-project/pull/102586

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

rdar://133473673

>From a823ade41b66908c23403b7e73ac99d564f8a488 Mon Sep 17 00:00:00 2001
From: Alastair Houghton <ahoughton at 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 7eb7da0138c972..5ac5532705dc49 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 25a2d8780fb56c..dc5e22e496cd6f 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
@@ -119,4 +119,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 41ee06c15448fd..f58faebf6a0049 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 e09c632842d6e9..de7630b9747ea4 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
@@ -455,6 +455,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 adaa01bf469a3ccf3c6fb9970ed4a509b0249a65 Mon Sep 17 00:00:00 2001
From: Alastair Houghton <alastair at 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 at compnerd.org>
---
 llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
index dc5e22e496cd6f..e75aceb0aa78c0 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
@@ -125,9 +125,7 @@ bool RuntimeDyldCOFF::relocationNeedsDLLImportStub(const RelocationRef &R) const
   if (!TargetNameOrErr)
     return false;
 
-  StringRef TargetName = *TargetNameOrErr;
-
-  return TargetName.startswith(getImportSymbolPrefix());
+  return TargetNameOrErr->startswith(getImportSymbolPrefix());
 }
 
 } // namespace llvm