[interop][SwiftToCxx] dispatch Swift class methods correctly using the vtable

hyp · hyp · commit b7007cb74892 · 2023-02-09T20:20:53.000-08:00
diff --git a/include/swift/IRGen/IRABIDetailsProvider.h b/include/swift/IRGen/IRABIDetailsProvider.h
@@ -248,6 +248,40 @@ class IRABIDetailsProvider {
   llvm::MapVector<EnumElementDecl *, EnumElementInfo>
   getEnumTagMapping(const EnumDecl *ED);
 
+  /// Details how a specific method should be dispatched.
+  struct MethodDispatchInfo {
+    enum class Kind {
+      /// A direct call can be made to the underlying function.
+      Direct,
+      /// An indirect call that can be made via a static offset in a vtable.
+      IndirectVTableStaticOffset
+    };
+
+    static MethodDispatchInfo direct() {
+      return MethodDispatchInfo(Kind::Direct, 0);
+    }
+
+    static MethodDispatchInfo indirectVTableStaticOffset(size_t bitOffset) {
+      return MethodDispatchInfo(Kind::IndirectVTableStaticOffset, bitOffset);
+    }
+
+    Kind getKind() const { return kind; }
+    size_t getStaticBitOffset() const {
+      assert(kind == Kind::IndirectVTableStaticOffset);
+      return bitOffset;
+    }
+
+  private:
+    constexpr MethodDispatchInfo(Kind kind, size_t bitOffset)
+        : kind(kind), bitOffset(bitOffset) {}
+
+    Kind kind;
+    size_t bitOffset;
+  };
+
+  Optional<MethodDispatchInfo>
+  getMethodDispatchInfo(const AbstractFunctionDecl *funcDecl);
+
 private:
   std::unique_ptr<IRABIDetailsProviderImpl> impl;
 };
diff --git a/lib/IRGen/IRABIDetailsProvider.cpp b/lib/IRGen/IRABIDetailsProvider.cpp
@@ -18,6 +18,7 @@
 #include "GenericRequirement.h"
 #include "IRGen.h"
 #include "IRGenModule.h"
+#include "MetadataLayout.h"
 #include "NativeConventionSchema.h"
 
 // FIXME: This include should removed once getFunctionLoweredSignature() is
@@ -202,6 +203,40 @@ class IRABIDetailsProviderImpl {
     return result;
   }
 
+  using MethodDispatchInfo = IRABIDetailsProvider::MethodDispatchInfo;
+
+  Optional<MethodDispatchInfo>
+  getMethodDispatchInfo(const AbstractFunctionDecl *funcDecl) {
+    if (funcDecl->isSemanticallyFinal())
+      return MethodDispatchInfo::direct();
+    // If this is an override of an existing method, then lookup
+    // its base method in its base class.
+    if (auto *overridenDecl = funcDecl->getOverriddenDecl())
+      funcDecl = overridenDecl;
+    auto *parentClass = dyn_cast<ClassDecl>(funcDecl->getDeclContext());
+    if (!parentClass)
+      return MethodDispatchInfo::direct();
+    auto &layout = IGM.getMetadataLayout(parentClass);
+    if (!isa<ClassMetadataLayout>(layout))
+      return {};
+    auto &classLayout = cast<ClassMetadataLayout>(layout);
+    auto *mi = classLayout.getStoredMethodInfoIfPresent(
+        SILDeclRef(const_cast<AbstractFunctionDecl *>(funcDecl)));
+    if (!mi)
+      return {};
+    switch (mi->TheKind) {
+    case ClassMetadataLayout::MethodInfo::Kind::DirectImpl:
+      return MethodDispatchInfo::direct();
+    case ClassMetadataLayout::MethodInfo::Kind::Offset:
+      if (mi->TheOffset.isStatic()) {
+        return MethodDispatchInfo::indirectVTableStaticOffset(
+            /*bitOffset=*/mi->TheOffset.getStaticOffset().getValue());
+      }
+      return {};
+    }
+    llvm_unreachable("invalid kind");
+  }
+
   Lowering::TypeConverter typeConverter;
   // Default silOptions are sufficient, as we don't need to generated SIL.
   SILOptions silOpts;
@@ -405,3 +440,9 @@ llvm::MapVector<EnumElementDecl *, IRABIDetailsProvider::EnumElementInfo>
 IRABIDetailsProvider::getEnumTagMapping(const EnumDecl *ED) {
   return impl->getEnumTagMapping(ED);
 }
+
+Optional<IRABIDetailsProvider::MethodDispatchInfo>
+IRABIDetailsProvider::getMethodDispatchInfo(
+    const AbstractFunctionDecl *funcDecl) {
+  return impl->getMethodDispatchInfo(funcDecl);
+}
diff --git a/lib/IRGen/MetadataLayout.h b/lib/IRGen/MetadataLayout.h
@@ -273,6 +273,14 @@ class ClassMetadataLayout : public NominalMetadataLayout {
 
   MethodInfo getMethodInfo(IRGenFunction &IGF, SILDeclRef method) const;
 
+  const StoredMethodInfo *
+  getStoredMethodInfoIfPresent(SILDeclRef method) const {
+    auto it = MethodInfos.find(method);
+    if (it != MethodInfos.end())
+      return &it->second;
+    return nullptr;
+  }
+
   Offset getFieldOffset(IRGenFunction &IGF, VarDecl *field) const;
 
   /// Assuming that the given field offset is at a static offset in
diff --git a/lib/PrintAsClang/DeclAndTypePrinter.cpp b/lib/PrintAsClang/DeclAndTypePrinter.cpp
@@ -995,6 +995,13 @@ class DeclAndTypePrinter::Implementation
                              /*selfTypeDeclContext=*/typeDeclContext);
       if (!funcABI)
         return;
+      Optional<IRABIDetailsProvider::MethodDispatchInfo> dispatchInfo;
+      if (!isa<ConstructorDecl>(AFD)) {
+        dispatchInfo = owningPrinter.interopContext.getIrABIDetails()
+                           .getMethodDispatchInfo(AFD);
+        if (!dispatchInfo)
+          return;
+      }
       owningPrinter.prologueOS << cFuncPrologueOS.str();
 
       printDocumentationComment(AFD);
@@ -1018,7 +1025,7 @@ class DeclAndTypePrinter::Implementation
                                    funcABI->getSignature(),
                                    funcABI->getSymbolName(), resultTy,
                                    /*isStatic=*/isClassMethod,
-                                   /*isDefinition=*/false);
+                                   /*isDefinition=*/false, dispatchInfo);
       }
 
       DeclAndTypeClangFunctionPrinter defPrinter(
@@ -1041,7 +1048,7 @@ class DeclAndTypePrinter::Implementation
         defPrinter.printCxxMethod(typeDeclContext, AFD, funcABI->getSignature(),
                                   funcABI->getSymbolName(), resultTy,
                                   /*isStatic=*/isClassMethod,
-                                  /*isDefinition=*/true);
+                                  /*isDefinition=*/true, dispatchInfo);
       }
 
       // FIXME: SWIFT_WARN_UNUSED_RESULT
diff --git a/lib/PrintAsClang/PrintClangFunction.cpp b/lib/PrintAsClang/PrintClangFunction.cpp
@@ -1084,7 +1084,8 @@ void DeclAndTypeClangFunctionPrinter::printCxxThunkBody(
     const AbstractFunctionDecl *FD, const LoweredFunctionSignature &signature,
     StringRef swiftSymbolName, const NominalTypeDecl *typeDeclContext,
     const ModuleDecl *moduleContext, Type resultTy, const ParameterList *params,
-    bool hasThrows, const AnyFunctionType *funcType, bool isStaticMethod) {
+    bool hasThrows, const AnyFunctionType *funcType, bool isStaticMethod,
+    Optional<IRABIDetailsProvider::MethodDispatchInfo> dispatchInfo) {
   if (typeDeclContext)
     ClangSyntaxPrinter(os).printNominalTypeOutsideMemberDeclInnerStaticAssert(
         typeDeclContext);
@@ -1096,9 +1097,29 @@ void DeclAndTypeClangFunctionPrinter::printCxxThunkBody(
     os << "  void* opaqueError = nullptr;\n";
     os << "  void* _ctx = nullptr;\n";
   }
+  Optional<StringRef> indirectFunctionVar;
+  if (dispatchInfo &&
+      dispatchInfo->getKind() !=
+          IRABIDetailsProvider::MethodDispatchInfo::Kind::Direct) {
+    assert(dispatchInfo->getKind() == IRABIDetailsProvider::MethodDispatchInfo::
+                                          Kind::IndirectVTableStaticOffset);
+    auto vtableBitOffset = dispatchInfo->getStaticBitOffset();
+
+    os << "void ***selfPtr_ = reinterpret_cast<void ***>( "
+          "::swift::_impl::_impl_RefCountedClass::getOpaquePointer(*this));\n";
+    os << "void **vtable_ = *selfPtr_;\n";
+    os << "using FType = decltype(" << cxx_synthesis::getCxxImplNamespaceName()
+       << "::" << swiftSymbolName << ");\n";
+    os << "FType *fptr_ = reinterpret_cast<FType *>(*(vtable_ + "
+       << (vtableBitOffset / 8) << "));\n"; // FIXME: not 8
+    indirectFunctionVar = StringRef("fptr_");
+  }
   auto printCallToCFunc = [&](Optional<StringRef> additionalParam) {
-    os << cxx_synthesis::getCxxImplNamespaceName() << "::" << swiftSymbolName
-       << '(';
+    if (indirectFunctionVar)
+      os << "(* " << *indirectFunctionVar << ')';
+    else
+      os << cxx_synthesis::getCxxImplNamespaceName() << "::" << swiftSymbolName;
+    os << '(';
 
     bool needsComma = false;
     size_t paramIndex = 1;
@@ -1329,7 +1350,8 @@ static StringRef getConstructorName(const AbstractFunctionDecl *FD) {
 void DeclAndTypeClangFunctionPrinter::printCxxMethod(
     const NominalTypeDecl *typeDeclContext, const AbstractFunctionDecl *FD,
     const LoweredFunctionSignature &signature, StringRef swiftSymbolName,
-    Type resultTy, bool isStatic, bool isDefinition) {
+    Type resultTy, bool isStatic, bool isDefinition,
+    Optional<IRABIDetailsProvider::MethodDispatchInfo> dispatchInfo) {
   bool isConstructor = isa<ConstructorDecl>(FD);
   os << "  ";
 
@@ -1357,10 +1379,11 @@ void DeclAndTypeClangFunctionPrinter::printCxxMethod(
 
   os << " {\n";
   // FIXME: should it be objTy for resultTy?
-  printCxxThunkBody(
-      FD, signature, swiftSymbolName, typeDeclContext, FD->getModuleContext(),
-      resultTy, FD->getParameters(), FD->hasThrows(),
-      FD->getInterfaceType()->castTo<AnyFunctionType>(), isStatic);
+  printCxxThunkBody(FD, signature, swiftSymbolName, typeDeclContext,
+                    FD->getModuleContext(), resultTy, FD->getParameters(),
+                    FD->hasThrows(),
+                    FD->getInterfaceType()->castTo<AnyFunctionType>(), isStatic,
+                    dispatchInfo);
   os << "  }\n";
 }
 
diff --git a/lib/PrintAsClang/PrintClangFunction.h b/lib/PrintAsClang/PrintClangFunction.h
@@ -18,6 +18,7 @@
 #include "swift/Basic/LLVM.h"
 #include "swift/ClangImporter/ClangImporter.h"
 #include "swift/IRGen/GenericRequirement.h"
+#include "swift/IRGen/IRABIDetailsProvider.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Optional.h"
@@ -107,22 +108,21 @@ class DeclAndTypeClangFunctionPrinter {
 
   /// Print the body of the inline C++ function thunk that calls the underlying
   /// Swift function.
-  void printCxxThunkBody(const AbstractFunctionDecl *FD,
-                         const LoweredFunctionSignature &signature,
-                         StringRef swiftSymbolName,
-                         const NominalTypeDecl *typeDeclContext,
-                         const ModuleDecl *moduleContext, Type resultTy,
-                         const ParameterList *params, bool hasThrows = false,
-                         const AnyFunctionType *funcType = nullptr,
-                         bool isStaticMethod = false);
+  void printCxxThunkBody(
+      const AbstractFunctionDecl *FD, const LoweredFunctionSignature &signature,
+      StringRef swiftSymbolName, const NominalTypeDecl *typeDeclContext,
+      const ModuleDecl *moduleContext, Type resultTy,
+      const ParameterList *params, bool hasThrows = false,
+      const AnyFunctionType *funcType = nullptr, bool isStaticMethod = false,
+      Optional<IRABIDetailsProvider::MethodDispatchInfo> dispatchInfo = None);
 
   /// Print the Swift method as C++ method declaration/definition, including
   /// constructors.
-  void printCxxMethod(const NominalTypeDecl *typeDeclContext,
-                      const AbstractFunctionDecl *FD,
-                      const LoweredFunctionSignature &signature,
-                      StringRef swiftSymbolName, Type resultTy, bool isStatic,
-                      bool isDefinition);
+  void printCxxMethod(
+      const NominalTypeDecl *typeDeclContext, const AbstractFunctionDecl *FD,
+      const LoweredFunctionSignature &signature, StringRef swiftSymbolName,
+      Type resultTy, bool isStatic, bool isDefinition,
+      Optional<IRABIDetailsProvider::MethodDispatchInfo> dispatchInfo);
 
   /// Print the C++ getter/setter method signature.
   void printCxxPropertyAccessorMethod(const NominalTypeDecl *typeDeclContext,
diff --git a/test/Interop/SwiftToCxx/class/swift-class-virtual-method-dispatch-execution.cpp b/test/Interop/SwiftToCxx/class/swift-class-virtual-method-dispatch-execution.cpp
@@ -0,0 +1,94 @@
+// RUN: %empty-directory(%t)
+
+// RUN: %target-swift-frontend %S/swift-class-virtual-method-dispatch.swift -typecheck -module-name Class -clang-header-expose-decls=all-public -emit-clang-header-path %t/class.h
+
+// RUN: %target-interop-build-clangxx -c %s -I %t -o %t/swift-class-execution.o
+// RUN: %target-interop-build-swift %S/swift-class-virtual-method-dispatch.swift -o %t/swift-class-execution -Xlinker %t/swift-class-execution.o -module-name Class -Xfrontend -entry-point-function-name -Xfrontend swiftMain
+
+// RUN: %target-codesign %t/swift-class-execution
+// RUN: %target-run %t/swift-class-execution | %FileCheck %s
+
+// FIXME: pointer signing support.
+// UNSUPPORTED: CPU=arm64e
+
+#include "class.h"
+#include <assert.h>
+#include <cstdio>
+
+using namespace Class;
+
+int main() {
+  auto base = returnBaseClass();
+  auto derived = returnDerivedClass();
+  BaseClass derivedAsBase = derived;
+  auto derivedDerived = returnDerivedDerivedClass();
+  BaseClass derivedDerivedAsBase = derivedDerived;
+  DerivedClass derivedDerivedAsDerived = derivedDerived;
+
+  {
+    base.virtualMethod();
+// CHECK: BaseClass.virtualMethod
+    puts("after base invoke");
+    derived.virtualMethod();
+    derivedAsBase.virtualMethod();
+// CHECK-NEXT: after base invoke
+// CHECK-NEXT: DerivedClass.virtualMethod
+// CHECK-NEXT: DerivedClass.virtualMethod
+    puts("after derived invoke");
+    derivedDerived.virtualMethod();
+    derivedDerivedAsBase.virtualMethod();
+// CHECK-NEXT: after derived invoke
+// CHECK-NEXT: DerivedDerivedClass.virtualMethod
+// CHECK-NEXT: DerivedDerivedClass.virtualMethod
+  }
+
+  {
+    swift::Int x;
+    x = base.virtualMethodIntInt(5);
+    assert(x == 5);
+
+    x = derived.virtualMethodIntInt(5);
+    assert(x == -5);
+    x = derivedAsBase.virtualMethodIntInt(-13);
+    assert(x == 13);
+
+    x = derivedDerived.virtualMethodIntInt(789);
+    assert(x == -789);
+    x = derivedDerivedAsBase.virtualMethodIntInt(76);
+    assert(x == -76);
+  }
+
+  {
+    swift::Int x;
+    x = base.finalMethodInBase(5);
+    assert(x == 10);
+
+    x = derived.finalMethodInBase(10);
+    assert(x == 20);
+    x = derivedAsBase.finalMethodInBase(30);
+    assert(x == 60);
+
+    x = derivedDerived.finalMethodInBase(11);
+    assert(x == 22);
+    x = derivedDerivedAsBase.finalMethodInBase(-22);
+    assert(x == -44);
+  }
+
+  {
+    auto obj = derived.virtualMethodInDerived(base);
+    obj.virtualMethod();
+// CHECK-NEXT: BaseClass.virtualMethod
+    obj = derivedDerived.virtualMethodInDerived(base);
+    obj.virtualMethod();
+// CHECK-NEXT: DerivedDerivedClass.virtualMethod
+    obj = derivedDerivedAsDerived.virtualMethodInDerived(base);
+    obj.virtualMethod();
+// CHECK-NEXT: DerivedDerivedClass.virtualMethod
+  }
+
+  {
+    derivedDerived.methodInDerivedDerived();
+// CHECK-NEXT: DerivedDerivedClass.methodInDerivedDerived
+  }
+  return 0;
+}
diff --git a/test/Interop/SwiftToCxx/class/swift-class-virtual-method-dispatch.swift b/test/Interop/SwiftToCxx/class/swift-class-virtual-method-dispatch.swift
