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

[Alastair Houghton via llvm-commits]

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

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.

Also add a couple of `override` specifiers that were missing last time (#102586).

rdar://133473673

>From e3cde339908bd682879362b21ea605a96a186fca Mon Sep 17 00:00:00 2001
From: Alastair Houghton <ahoughton at apple.com>
Date: Mon, 2 Sep 2024 10:36:41 +0100
Subject: [PATCH 1/2] [RuntimeDyld][Windows] Allocate space for dllimport
 things. (#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
---
 llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp   |  5 ++++-
 .../ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp    | 10 ++++++++++
 llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h |  7 +++++++
 llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h | 10 ++++++++++
 4 files changed, 31 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..73b37ee0ff3311 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.cpp
@@ -119,4 +119,14 @@ 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;
+
+  return TargetNameOrErr->starts_with(getImportSymbolPrefix());
+}
+
 } // namespace llvm
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h
index 25e3783cf160b2..51d177c7bb8bec 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"
 
 namespace llvm {
 
@@ -45,6 +46,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 8aa50920e7791e1da98ef50c15bb0499f0554e21 Mon Sep 17 00:00:00 2001
From: Alastair Houghton <ahoughton at apple.com>
Date: Mon, 2 Sep 2024 10:40:10 +0100
Subject: [PATCH 2/2] [RuntimeDyld][Windows] Add missing `override` specifiers.

We were missing a couple of `override` specifiers.

rdar://133473673
---
 llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h
index 51d177c7bb8bec..8ca1c2e4efc3f4 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCOFF.h
@@ -46,9 +46,9 @@ class RuntimeDyldCOFF : public RuntimeDyldImpl {
 
   static constexpr StringRef getImportSymbolPrefix() { return "__imp_"; }
 
-  bool relocationNeedsDLLImportStub(const RelocationRef &R) const;
+  bool relocationNeedsDLLImportStub(const RelocationRef &R) const override;
 
-  unsigned sizeAfterAddingDLLImportStub(unsigned Size) const {
+  unsigned sizeAfterAddingDLLImportStub(unsigned Size) const override {
     return alignTo(Size, PointerSize) + PointerSize;
   }