[cxx-interop] Fix miscompilation for inferred shared references

When a base class is annotated as shared reference we can occasionally
infer that the derived types also need to be shared references.
Unfortunately, we did not generate the correct code for some of those
scenarios. When the reference counted base is not at the offset zero we
need to do offset adjustments before we pass the pointer to the
reference counting functions. We did not do those offset calculations.

I looked into implementing the codegen for the offset calculation
directly in Swift but it needed significantly more work than I
anticipated. We need to invoke the frontend to get the path to the base
class and we also need to deal with virtual inheritance, alignment and
some other considerations.

This PR ends up generating a Clang shim instead and the derived to base
conversion happens in this shim. As a result, we piggy-back on Clang
making all the correct offset calculations. This patch also had to
change how certain aspects of shared references are implemented to be
compatible with this approach:
* Instead of always looking at the base classes to querry the
  retain/release operations we now propagate the corresponding
  annotations once per types. This also has the beneficial effects that
  we traverse the inheritance hierarchy less often.
* To generate the correct diagnostics, I reuse the result of the
  refcount operation query.
* We do not want these generated functions to be inherited, so added a
  set to exempt them from cloning.
* Tweaked the lookup logic for retain/release a bit as these generated
  clang methods are not found by lookup. We rely on looking up the
  imported methods instead.

rdar://166227787
rdar://165635002

Author: Gabor Horvath

include/swift/AST/ClangModuleLoader.h

@@ -216,9 +216,13 @@ class ClangModuleLoader : public ModuleLoader {
                                           DeclContext *newContext,
                                           ClangInheritanceInfo inheritance) = 0;
 
-  /// Returnes the original method if \param decl is a clone from a base class
+  /// Returns the original method if \param decl is a clone from a base class
   virtual ValueDecl *getOriginalForClonedMember(const ValueDecl *decl) = 0;
 
+  /// Returns true if we synthesize this member for every type so no need to
+  /// clone it for the derived classes.
+  virtual bool isPerTypeSynthesizedMember(const ValueDecl *decl) = 0;
+
   /// Emits diagnostics for any declarations named name
   /// whose direct declaration context is a TU.
   virtual void diagnoseTopLevelValue(const DeclName &name) = 0;

include/swift/ClangImporter/ClangImporter.h

@@ -650,6 +650,7 @@ class ClangImporter final : public ClangModuleLoader {
                                   ClangInheritanceInfo inheritance) override;
 
   ValueDecl *getOriginalForClonedMember(const ValueDecl *decl) override;
+  bool isPerTypeSynthesizedMember(const ValueDecl *decl) override;
 
   /// Emits diagnostics for any declarations named name
   /// whose direct declaration context is a TU.

include/swift/ClangImporter/ClangImporterRequests.h

@@ -500,7 +500,7 @@ struct CustomRefCountingOperationResult {
 
   CustomRefCountingOperationResultKind kind;
   ValueDecl *operation;
-  std::string name;
+  StringRef name;
 };
 
 class CustomRefCountingOperation

lib/ClangImporter/ClangImporter.cpp

@@ -6565,6 +6565,7 @@ TinyPtrVector<ValueDecl *> ClangRecordMemberLookup::evaluate(
       auto namedMember = dyn_cast<ValueDecl>(member);
       if (!namedMember || !namedMember->hasName() ||
           namedMember->getName().getBaseName() != name ||
+          clangModuleLoader->isPerTypeSynthesizedMember(namedMember) ||
           clangModuleLoader->getOriginalForClonedMember(namedMember))
         continue;
 
@@ -8033,6 +8034,16 @@ ValueDecl *ClangImporter::Implementation::getOriginalForClonedMember(
   return nullptr;
 }
 
+bool ClangImporter::Implementation::isPerTypeSynthesizedMember(
+    const ValueDecl *decl) {
+  return perTypeSynthesizedMembers.contains(decl);
+}
+
+void ClangImporter::Implementation::markPerTypeSynthesizedMember(
+    const ValueDecl *decl) {
+  perTypeSynthesizedMembers.insert(decl);
+}
+
 size_t ClangImporter::Implementation::getImportedBaseMemberDeclArity(
     const ValueDecl *valueDecl) {
   if (auto *func = dyn_cast<FuncDecl>(valueDecl)) {
@@ -8053,6 +8064,10 @@ ValueDecl *ClangImporter::getOriginalForClonedMember(const ValueDecl *decl) {
   return Impl.getOriginalForClonedMember(decl);
 }
 
+bool ClangImporter::isPerTypeSynthesizedMember(const ValueDecl *decl) {
+  return Impl.isPerTypeSynthesizedMember(decl);
+}
+
 void ClangImporter::diagnoseTopLevelValue(const DeclName &name) {
   Impl.diagnoseTopLevelValue(name);
 }
@@ -8206,12 +8221,15 @@ importer::getValueDeclsForName(NominalTypeDecl *decl, StringRef name) {
   auto clangDecl = decl->getClangDecl();
   llvm::SmallVector<ValueDecl *, 1> results;
 
-  if (name.starts_with(".")) {
+  if (name.consume_front(".")) {
     // Look for a member of decl instead of a global.
-    StringRef memberName = name.drop_front(1);
-    if (memberName.empty())
+    if (name.empty())
       return {};
-    auto declName = DeclName(ctx.getIdentifier(memberName));
+    auto declName = DeclName(ctx.getIdentifier(name));
+    auto swiftLookupResults = decl->lookupDirect(declName);
+    if (!swiftLookupResults.empty())
+      return SmallVector<ValueDecl *, 1>(swiftLookupResults.begin(),
+                                         swiftLookupResults.end());
     auto allResults = evaluateOrDefault(
         ctx.evaluator, ClangRecordMemberLookup({decl, declName}), {});
     return SmallVector<ValueDecl *, 1>(allResults.begin(), allResults.end());
@@ -8242,9 +8260,9 @@ importer::getValueDeclsForName(NominalTypeDecl *decl, StringRef name) {
   return results;
 }
 
-static const clang::RecordDecl *
-getRefParentOrDiag(const clang::RecordDecl *decl, ASTContext &ctx,
-                   ClangImporter::Implementation *importerImpl) {
+const clang::RecordDecl *
+importer::getRefParentOrDiag(const clang::RecordDecl *decl, ASTContext &ctx,
+                             ClangImporter::Implementation *importerImpl) {
   auto refParentDecls = getRefParentDecls(decl, ctx, importerImpl);
   if (refParentDecls.empty())
     return nullptr;
@@ -8257,11 +8275,11 @@ getRefParentOrDiag(const clang::RecordDecl *decl, ASTContext &ctx,
     assert(refParentDecl && "refParentDecl is null inside getRefParentOrDiag");
     for (const auto *attr : refParentDecl->getAttrs()) {
       if (const auto swiftAttr = llvm::dyn_cast<clang::SwiftAttrAttr>(attr)) {
-        const auto &attribute = swiftAttr->getAttribute();
-        if (attribute.starts_with(retainPrefix))
-          uniqueRetainDecls.insert(attribute.drop_front(retainPrefix.size()));
-        else if (attribute.starts_with(releasePrefix))
-          uniqueReleaseDecls.insert(attribute.drop_front(releasePrefix.size()));
+        auto attribute = swiftAttr->getAttribute();
+        if (attribute.consume_front(retainPrefix))
+          uniqueRetainDecls.insert(attribute);
+        else if (attribute.consume_front(releasePrefix))
+          uniqueReleaseDecls.insert(attribute);
       }
     }
   }
@@ -8811,19 +8829,12 @@ CustomRefCountingOperationResult CustomRefCountingOperation::evaluate(
     Evaluator &evaluator, CustomRefCountingOperationDescriptor desc) const {
   auto swiftDecl = desc.decl;
   auto operation = desc.kind;
-  auto &ctx = swiftDecl->getASTContext();
 
-  std::string operationStr = operation == CustomRefCountingOperationKind::retain
-                                 ? "retain:"
-                                 : "release:";
+  StringRef operationStr = operation == CustomRefCountingOperationKind::retain
+                               ? "retain:"
+                               : "release:";
 
   auto decl = cast<clang::RecordDecl>(swiftDecl->getClangDecl());
-
-  if (!hasImportAsRefAttr(decl)) {
-    if (auto parentRefDecl = getRefParentOrDiag(decl, ctx, nullptr))
-      decl = parentRefDecl;
-  }
-
   if (!decl->hasAttrs())
     return {CustomRefCountingOperationResult::noAttribute, nullptr, ""};
 
@@ -8842,10 +8853,8 @@ CustomRefCountingOperationResult CustomRefCountingOperation::evaluate(
   if (retainReleaseAttrs.size() > 1)
     return {CustomRefCountingOperationResult::tooManyAttributes, nullptr, ""};
 
-  auto name = retainReleaseAttrs.front()
-                  ->getAttribute()
-                  .drop_front(StringRef(operationStr).size())
-                  .str();
+  auto name = retainReleaseAttrs.front()->getAttribute().drop_front(
+      operationStr.size());
 
   if (name == "immortal")
     return {CustomRefCountingOperationResult::immortal, nullptr, name};

lib/ClangImporter/ImportDecl.cpp

@@ -2190,6 +2190,39 @@ namespace {
       return true;
     }
 
+    std::pair<CustomRefCountingOperationResult,
+              CustomRefCountingOperationResult>
+    addRefCountOperationsIfRequired(ClassDecl *nominal,
+                                    clang::RecordDecl *clangType) {
+      auto &context = Impl.SwiftContext;
+      auto nonInheritedRefCountingOperations = [&] {
+        auto retainResult = evaluateOrDefault(
+            context.evaluator,
+            CustomRefCountingOperation(
+                {nominal, CustomRefCountingOperationKind::retain}),
+            {});
+        auto releaseResult = evaluateOrDefault(
+            context.evaluator,
+            CustomRefCountingOperation(
+                {nominal, CustomRefCountingOperationKind::release}),
+            {});
+        return std::make_pair(retainResult, releaseResult);
+      };
+      auto clangDecl = dyn_cast<clang::CXXRecordDecl>(clangType);
+      if (!clangDecl)
+        return nonInheritedRefCountingOperations();
+      auto baseClangDecl = dyn_cast_or_null<clang::CXXRecordDecl>(
+          getRefParentOrDiag(clangDecl, context, nullptr));
+      if (!baseClangDecl || baseClangDecl == clangDecl)
+        return nonInheritedRefCountingOperations();
+
+      auto baseSwiftDecl = cast<ClassDecl>(
+          Impl.importDecl(baseClangDecl, getActiveSwiftVersion()));
+
+      return synthesizer.addRefCountOperations(nominal, clangDecl,
+                                               baseSwiftDecl, baseClangDecl);
+    }
+
     Decl *VisitRecordDecl(const clang::RecordDecl *decl) {
       // Track whether this record contains fields we can't reference in Swift
       // as stored properties.
@@ -2786,7 +2819,11 @@ namespace {
       }
 
       if (auto classDecl = dyn_cast<ClassDecl>(result)) {
-        validateForeignReferenceType(decl, classDecl);
+        auto operations = addRefCountOperationsIfRequired(
+            classDecl, const_cast<clang::RecordDecl *>(decl));
+
+        validateForeignReferenceType(decl, classDecl, operations.first,
+                                     operations.second);
 
         auto availability = Impl.SwiftContext.getSwift58Availability();
         if (!availability.isAlwaysAvailable()) {
@@ -2853,8 +2890,10 @@ namespace {
       }
     }
 
-    void validateForeignReferenceType(const clang::RecordDecl *decl,
-                                      ClassDecl *classDecl) {
+    void validateForeignReferenceType(
+        const clang::RecordDecl *decl, ClassDecl *classDecl,
+        CustomRefCountingOperationResult retainOperation,
+        CustomRefCountingOperationResult releaseOperation) {
 
       enum class RetainReleaseOperationKind {
         notAfunction,
@@ -2937,11 +2976,6 @@ namespace {
         return RetainReleaseOperationKind::valid;
       };
 
-      auto retainOperation = evaluateOrDefault(
-          Impl.SwiftContext.evaluator,
-          CustomRefCountingOperation(
-              {classDecl, CustomRefCountingOperationKind::retain}),
-          {});
       if (retainOperation.kind ==
           CustomRefCountingOperationResult::noAttribute) {
         HeaderLoc loc(decl->getLocation());
@@ -3008,11 +3042,6 @@ namespace {
                CustomRefCountingOperationResult::immortal);
       }
 
-      auto releaseOperation = evaluateOrDefault(
-          Impl.SwiftContext.evaluator,
-          CustomRefCountingOperation(
-              {classDecl, CustomRefCountingOperationKind::release}),
-          {});
       if (releaseOperation.kind ==
           CustomRefCountingOperationResult::noAttribute) {
         HeaderLoc loc(decl->getLocation());

lib/ClangImporter/ImporterImpl.h

@@ -48,6 +48,7 @@
 #include "clang/Serialization/ModuleFileExtension.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/SmallBitVector.h"
@@ -692,6 +693,7 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
 
   // Map all cloned methods back to the original member
   llvm::DenseMap<ValueDecl *, ValueDecl *> clonedMembers;
+  llvm::DenseSet<const ValueDecl *> perTypeSynthesizedMembers;
 
   // Keep track of methods that are unavailale in each class.
   // We need this set because these methods will be imported lazily. We don't
@@ -714,6 +716,8 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
                                   ClangInheritanceInfo inheritance);
 
   ValueDecl *getOriginalForClonedMember(const ValueDecl *decl);
+  bool isPerTypeSynthesizedMember(const ValueDecl *decl);
+  void markPerTypeSynthesizedMember(const ValueDecl *decl);
 
   static size_t getImportedBaseMemberDeclArity(const ValueDecl *valueDecl);
 
@@ -2231,6 +2235,11 @@ getImplicitObjectParamAnnotation(const clang::FunctionDecl *FD) {
   return nullptr;
 }
 
+/// Find a unique base class that is annotated as SHARED_REFERENCE if any.
+const clang::RecordDecl *
+getRefParentOrDiag(const clang::RecordDecl *decl, ASTContext &ctx,
+                   ClangImporter::Implementation *importerImpl);
+
 } // end namespace importer
 } // end namespace swift