From d766b8d736ba42b1f44795f859f6a2c25df955f6 Mon Sep 17 00:00:00 2001
From: Peter Rong <PeterRong@meta.com>
Date: Mon, 18 May 2026 15:58:59 -0700
Subject: [PATCH] [IRGenSIL] resolve objc_direct methods to tthunk if
 PreconditionThunk is enabled

---
 lib/IRGen/IRGenSIL.cpp                        | 22 +++++++
 .../objc_direct_precondition_thunk.swift      | 57 +++++++++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 test/IRGen/objc_direct_precondition_thunk.swift

diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
index 17a678d984ad..875efff61983 100644
--- a/lib/IRGen/IRGenSIL.cpp
+++ b/lib/IRGen/IRGenSIL.cpp
@@ -50,6 +50,8 @@
 #include "swift/SIL/TerminatorUtils.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclCXX.h"
+#include "clang/AST/DeclObjC.h"
+#include "clang/Basic/CodeGenOptions.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/CodeGen/CodeGenABITypes.h"
 #include "llvm/ADT/MapVector.h"
@@ -3211,6 +3213,26 @@ void IRGenSILFunction::visitFunctionRefBaseInst(FunctionRefBaseInst *i) {
     value = fnPtr;
     secondaryValue = nullptr;
   }
+  // For ObjC direct methods with precondition thunks enabled, resolve to
+  // the thunk instead of the true implementation. The thunk performs nil
+  // receiver checks (and class realization for class methods) before
+  // tail-calling the implementation.
+  if (auto *clangDecl = fn->getClangDecl()) {
+    if (auto *objcMethodDecl = dyn_cast<clang::ObjCMethodDecl>(clangDecl)) {
+      if (objcMethodDecl->isDirectMethod() &&
+          IGM.getClangCGM().getCodeGenOpts().ObjCDirectPreconditionThunk) {
+        const auto *cd = objcMethodDecl->getClassInterface();
+        bool receiverCanBeNull = true;
+        bool classObjectCanBeUnrealized = objcMethodDecl->isClassMethod();
+        llvm::Function *thunkFn =
+            clang::CodeGen::getObjCDirectMethodCallee(
+                IGM.getClangCGM(), objcMethodDecl, cd, receiverCanBeNull,
+                classObjectCanBeUnrealized);
+        value = thunkFn;
+      }
+    }
+  }
+
   FunctionPointer fp =
       FunctionPointer::forDirect(fpKind, value, secondaryValue, sig, useSignature);
   // Update the foreign no-throw information if needed.
diff --git a/test/IRGen/objc_direct_precondition_thunk.swift b/test/IRGen/objc_direct_precondition_thunk.swift
new file mode 100644
index 000000000000..4a9fdd90ccc4
--- /dev/null
+++ b/test/IRGen/objc_direct_precondition_thunk.swift
@@ -0,0 +1,57 @@
+// Test that Swift correctly dispatches to ObjC direct method thunks when
+// -fobjc-direct-precondition-thunk is enabled.
+
+// RUN: %target-swift-emit-ir -import-objc-header %S/../Inputs/objc_direct.h -Xcc -fobjc-direct-precondition-thunk -o - %s | %FileCheck %s
+
+// REQUIRES: objc_interop
+
+func markUsed<T>(_ t: T) {}
+
+let bar = Bar()
+
+// Instance method calls should go through the thunk.
+bar.directProperty = 123
+// CHECK: call void @"-[Bar setDirectProperty:]D_thunk"(ptr %{{[0-9]+}}, i32 {{.*}})
+
+markUsed(bar.directProperty)
+// CHECK: call i32 @"-[Bar directProperty]D_thunk"(ptr %{{[0-9]+}})
+
+bar.directProperty2 = 456
+// CHECK: call void @"-[Bar setDirectProperty2:]D_thunk"(ptr %{{[0-9]+}}, i32 {{.*}})
+
+markUsed(bar.directProperty2)
+// CHECK: call i32 @"-[Bar directProperty2]D_thunk"(ptr %{{[0-9]+}})
+
+bar[0] = 789
+// CHECK: call void @"-[Bar setObject:atIndexedSubscript:]D_thunk"(ptr %{{[0-9]+}}, i32 789, i32 0)
+
+markUsed(bar[0])
+// CHECK: call i32 @"-[Bar objectAtIndexedSubscript:]D_thunk"(ptr %{{[0-9]+}}, i32 0)
+
+markUsed(bar.directMethod())
+// CHECK: call {{.*}} @"-[Bar directMethod]D_thunk"(ptr %{{[0-9]+}})
+
+markUsed(bar.directMethod2())
+// CHECK: call {{.*}} @"-[Bar directMethod2]D_thunk"(ptr %{{[0-9]+}})
+
+// Class method calls should also go through the thunk.
+markUsed(Bar.directClassMethod())
+// CHECK: call {{.*}} @"+[Bar directClassMethod]D_thunk"(ptr
+
+markUsed(Bar.directClassMethod2())
+// CHECK: call {{.*}} @"+[Bar directClassMethod2]D_thunk"(ptr
+
+markUsed(bar.directProtocolMethod())
+// CHECK: call {{.*}} @"-[Bar directProtocolMethod]D_thunk"({{.*}})
+
+// Thunks should be declared with linkonce_odr hidden.
+// CHECK-DAG: define linkonce_odr hidden {{.*}} @"-[Bar setDirectProperty:]D_thunk"
+// CHECK-DAG: define linkonce_odr hidden {{.*}} @"-[Bar directProperty]D_thunk"
+// CHECK-DAG: define linkonce_odr hidden {{.*}} @"-[Bar directMethod]D_thunk"
+// CHECK-DAG: define linkonce_odr hidden {{.*}} @"+[Bar directClassMethod]D_thunk"
+
+// True implementations should be declared (external, with D suffix).
+// CHECK-DAG: declare {{.*}} @"-[Bar directProperty]D"
+// CHECK-DAG: declare {{.*}} @"-[Bar setDirectProperty:]D"
+// CHECK-DAG: declare {{.*}} @"-[Bar directMethod]D"
+// CHECK-DAG: declare {{.*}} @"+[Bar directClassMethod]D"